[cmd provider] add support of match size in logs

internal/util.go

@@ -86,13 +86,32 @@ func GetJSON(url string, obj interface{}, client *http.Client) (*http.Response,
 	return resp, json.Unmarshal(body, obj)
 }
 
-func ExtractSizeFromRsyncLog(content []byte) string {
-	// (?m) flag enables multi-line mode
-	re := regexp.MustCompile(`(?m)^Total file size: ([0-9\.]+[KMGTP]?) bytes`)
-	matches := re.FindAllSubmatch(content, -1)
-	// fmt.Printf("%q\n", matches)
-	if len(matches) == 0 {
+// FindAllSubmatchInFile calls re.FindAllSubmatch to find matches in given file
+func FindAllSubmatchInFile(fileName string, re *regexp.Regexp) (matches [][][]byte, err error) {
+	if fileName == "/dev/null" {
+		err = errors.New("Invalid log file")
+		return
+	}
+	if content, err := ioutil.ReadFile(fileName); err == nil {
+		matches = re.FindAllSubmatch(content, -1)
+		// fmt.Printf("FindAllSubmatchInFile: %q\n", matches)
+	}
+	return
+}
+
+// ExtractSizeFromLog uses a regexp to extract the size from log files
+func ExtractSizeFromLog(logFile string, re *regexp.Regexp) string {
+	matches, _ := FindAllSubmatchInFile(logFile, re)
+	if matches == nil || len(matches) == 0 {
 		return ""
 	}
+	// return the first capture group of the last occurrence
 	return string(matches[len(matches)-1][1])
 }
+
+// ExtractSizeFromRsyncLog extracts the size from rsync logs
+func ExtractSizeFromRsyncLog(logFile string) string {
+	// (?m) flag enables multi-line mode
+	re := regexp.MustCompile(`(?m)^Total file size: ([0-9\.]+[KMGTP]?) bytes`)
+	return ExtractSizeFromLog(logFile, re)
+}

worker/cmd_provider.go

@@ -3,11 +3,11 @@ package worker
 import (
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"regexp"
 	"time"
 
 	"github.com/anmitsu/go-shlex"
+	"github.com/tuna/tunasync/internal"
 )
 
 type cmdConfig struct {
@@ -18,13 +18,16 @@ type cmdConfig struct {
 	retry                       int
 	env                         map[string]string
 	failOnMatch                 string
+	sizePattern                 string
 }
 
 type cmdProvider struct {
 	baseProvider
 	cmdConfig
 	command     []string
+	dataSize    string
 	failOnMatch *regexp.Regexp
+	sizePattern *regexp.Regexp
 }
 
 func newCmdProvider(c cmdConfig) (*cmdProvider, error) {
@@ -59,6 +62,14 @@ func newCmdProvider(c cmdConfig) (*cmdProvider, error) {
 		}
 		provider.failOnMatch = failOnMatch
 	}
+	if len(c.sizePattern) > 0 {
+		var err error
+		sizePattern, err := regexp.Compile(c.sizePattern)
+		if err != nil {
+			return nil, errors.New("size-pattern regexp error: " + err.Error())
+		}
+		provider.sizePattern = sizePattern
+	}
 
 	return provider, nil
 }
@@ -71,7 +82,12 @@ func (p *cmdProvider) Upstream() string {
 	return p.upstreamURL
 }
 
+func (p *cmdProvider) DataSize() string {
+	return p.dataSize
+}
+
 func (p *cmdProvider) Run() error {
+	p.dataSize = ""
 	if err := p.Start(); err != nil {
 		return err
 	}
@@ -79,16 +95,18 @@ func (p *cmdProvider) Run() error {
 		return err
 	}
 	if p.failOnMatch != nil {
-		if logContent, err := ioutil.ReadFile(p.LogFile()); err == nil {
-			matches := p.failOnMatch.FindAllSubmatch(logContent, -1)
-			if len(matches) != 0 {
-				logger.Debug("Fail-on-match: %r", matches)
-				return errors.New(
-					fmt.Sprintf("Fail-on-match regexp found %d matches", len(matches)))
-			}
-		} else {
+		matches, err := internal.FindAllSubmatchInFile(p.LogFile(), p.failOnMatch)
+		fmt.Printf("FindAllSubmatchInFile: %q\n", matches)
+		if err != nil {
 			return err
 		}
+		if len(matches) != 0 {
+			logger.Debug("Fail-on-match: %r", matches)
+			return fmt.Errorf("Fail-on-match regexp found %d matches", len(matches))
+		}
+	}
+	if p.sizePattern != nil {
+		p.dataSize = internal.ExtractSizeFromLog(p.LogFile(), p.sizePattern)
 	}
 	return nil
 }

worker/config.go

@@ -131,6 +131,7 @@ type mirrorConfig struct {
 
 	Command       string   `toml:"command"`
 	FailOnMatch   string   `toml:"fail_on_match"`
+	SizePattern   string   `toml:"size_pattern"`
 	UseIPv6       bool     `toml:"use_ipv6"`
 	UseIPv4       bool     `toml:"use_ipv4"`
 	ExcludeFile   string   `toml:"exclude_file"`

worker/provider.go

@@ -113,6 +113,7 @@ func newMirrorProvider(mirror mirrorConfig, cfg *Config) mirrorProvider {
 			command:     mirror.Command,
 			workingDir:  mirrorDir,
 			failOnMatch: mirror.FailOnMatch,
+			sizePattern: mirror.SizePattern,
 			logDir:      logDir,
 			logFile:     filepath.Join(logDir, "latest.log"),
 			interval:    time.Duration(mirror.Interval) * time.Minute,

worker/provider_test.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strconv"
 	"testing"
 	"time"
 
@@ -365,7 +366,7 @@ sleep 5
 
 		})
 	})
-	Convey("Command Provider with fail-on-match regexp should work", t, func(ctx C) {
+	Convey("Command Provider with RegExprs should work", t, func(ctx C) {
 		tmpDir, err := ioutil.TempDir("", "tunasync")
 		defer os.RemoveAll(tmpDir)
 		So(err, ShouldBeNil)
@@ -376,22 +377,24 @@ sleep 5
 			upstreamURL: "http://mirrors.tuna.moe/",
 			command:     "uptime",
 			failOnMatch: "",
+			sizePattern: "",
 			workingDir:  tmpDir,
 			logDir:      tmpDir,
 			logFile:     tmpFile,
 			interval:    600 * time.Second,
 		}
 
-		Convey("when regexp matches", func() {
+		Convey("when fail-on-match regexp matches", func() {
 			c.failOnMatch = `[a-z]+`
 			provider, err := newCmdProvider(c)
 			So(err, ShouldBeNil)
 
 			err = provider.Run()
 			So(err, ShouldNotBeNil)
+			So(provider.DataSize(), ShouldBeEmpty)
 		})
 
-		Convey("when regexp does not match", func() {
+		Convey("when fail-on-match regexp does not match", func() {
 			c.failOnMatch = `load average_`
 			provider, err := newCmdProvider(c)
 			So(err, ShouldBeNil)
@@ -399,6 +402,49 @@ sleep 5
 			err = provider.Run()
 			So(err, ShouldBeNil)
 		})
+
+		Convey("when fail-on-match regexp meets /dev/null", func() {
+			c.failOnMatch = `load average_`
+			c.logFile = "/dev/null"
+			provider, err := newCmdProvider(c)
+			So(err, ShouldBeNil)
+
+			err = provider.Run()
+			So(err, ShouldNotBeNil)
+		})
+
+		Convey("when size-pattern regexp matches", func() {
+			c.sizePattern = `load average: ([\d\.]+)`
+			provider, err := newCmdProvider(c)
+			So(err, ShouldBeNil)
+
+			err = provider.Run()
+			So(err, ShouldBeNil)
+			So(provider.DataSize(), ShouldNotBeEmpty)
+			_, err = strconv.ParseFloat(provider.DataSize(), 32)
+			So(err, ShouldBeNil)
+		})
+
+		Convey("when size-pattern regexp does not match", func() {
+			c.sizePattern = `load ave: ([\d\.]+)`
+			provider, err := newCmdProvider(c)
+			So(err, ShouldBeNil)
+
+			err = provider.Run()
+			So(err, ShouldBeNil)
+			So(provider.DataSize(), ShouldBeEmpty)
+		})
+
+		Convey("when size-pattern regexp meets /dev/null", func() {
+			c.sizePattern = `load ave: ([\d\.]+)`
+			c.logFile = "/dev/null"
+			provider, err := newCmdProvider(c)
+			So(err, ShouldBeNil)
+
+			err = provider.Run()
+			So(err, ShouldNotBeNil)
+			So(provider.DataSize(), ShouldBeEmpty)
+		})
 	})
 }
 

worker/rsync_provider.go

@@ -2,7 +2,6 @@ package worker
 
 import (
 	"errors"
-	"io/ioutil"
 	"strings"
 	"time"
 
@@ -112,9 +111,7 @@ func (p *rsyncProvider) Run() error {
 	if err := p.Wait(); err != nil {
 		return err
 	}
-	if logContent, err := ioutil.ReadFile(p.LogFile()); err == nil {
-		p.dataSize = internal.ExtractSizeFromRsyncLog(logContent)
-	}
+	p.dataSize = internal.ExtractSizeFromRsyncLog(p.LogFile())
 	return nil
 }
 

worker/two_stage_rsync_provider.go

@@ -3,7 +3,6 @@ package worker
 import (
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"strings"
 	"time"
 
@@ -171,8 +170,6 @@ func (p *twoStageRsyncProvider) Run() error {
 			return err
 		}
 	}
-	if logContent, err := ioutil.ReadFile(p.LogFile()); err == nil {
-		p.dataSize = internal.ExtractSizeFromRsyncLog(logContent)
-	}
+	p.dataSize = internal.ExtractSizeFromRsyncLog(p.LogFile())
 	return nil
 }