feat: avoid unnecessary file override during the initial sync

Qingping Hou · Qingping Hou · commit 8ade48a6209c · 2020-03-23T00:31:22.000-07:00
diff --git a/main.go b/main.go
@@ -23,6 +23,7 @@ var (
 	FlagRunOnce        bool
 	FlagStatusAddr     = ":8087"
 	FlagExclude        []string
+	FlagScratch        bool
 
 	metricsSyncTime = prometheus.NewGauge(prometheus.GaugeOpts{
 		Namespace: "objinsync",
@@ -102,9 +103,10 @@ func main() {
 
 			puller := sync.NewPuller()
 			if FlagExclude != nil {
-				for _, exclude := range FlagExclude {
-					puller.AddExcludePattern(exclude)
-				}
+				puller.AddExcludePatterns(FlagExclude)
+			}
+			if !FlagScratch {
+				puller.PopulateChecksum(localDir)
 			}
 
 			pull := func() {
@@ -151,6 +153,13 @@ func main() {
 		&FlagStatusAddr, "status-addr", "s", ":8087", "binding address for status endpoint")
 	pullCmd.PersistentFlags().StringSliceVarP(
 		&FlagExclude, "exclude", "e", nil, "exclude files matching given pattern, see https://github.com/bmatcuk/doublestar#patterns for pattern spec")
+	pullCmd.PersistentFlags().BoolVarP(
+		&FlagScratch,
+		"scratch",
+		"",
+		false,
+		"skip checksums calculation and override all files during the initial sync",
+	)
 
 	rootCmd.AddCommand(pullCmd)
 	rootCmd.Execute()
diff --git a/pkg/sync/pull.go b/pkg/sync/pull.go
@@ -1,6 +1,8 @@
 package sync
 
 import (
+	"crypto/md5"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -56,8 +58,11 @@ type DownloadTask struct {
 	Uri       string
 	LocalPath string
 	Uid       string
+	// uid key is common suffix between local path and remote uri
+	UidKey string
 }
 
+// parse bucket and key out of remote object URI
 func parseObjectUri(uri string) (string, string, error) {
 	parts := strings.SplitN(uri, "//", 2)
 	if len(parts) != 2 {
@@ -73,6 +78,27 @@ func parseObjectUri(uri string) (string, string, error) {
 	return pathParts[0], pathParts[1], nil
 }
 
+func uidKeyFromLocalPath(localDir string, localPath string) (string, error) {
+	return filepath.Rel(localDir, localPath)
+}
+
+func uidFromLocalPath(localPath string) (string, error) {
+	f, err := os.Open(localPath)
+	if err != nil {
+		return "", fmt.Errorf("Invalid file path for checksum calculation: %s, err: %s", localPath, err)
+	}
+	defer f.Close()
+
+	h := md5.New()
+	if _, err := io.Copy(h, f); err != nil {
+		return "", fmt.Errorf("Failed to calculate checksum for file: %s, err: %s", localPath, err)
+	}
+
+	uid := hex.EncodeToString(h.Sum(nil))
+	// AWS S3 ETag is a quoted hex string
+	return fmt.Sprintf("\"%s\"", uid), nil
+}
+
 func (self *Puller) downloadHandler(task DownloadTask, downloader GenericDownloader) {
 	l := zap.S()
 
@@ -120,9 +146,19 @@ func (self *Puller) downloadHandler(task DownloadTask, downloader GenericDownloa
 
 	// update cache with new object ID
 	self.uidLock.Lock()
-	l.Debugw("Updaing uid cache", "key", task.Uri, "val", task.Uid)
-	self.uidCache[task.Uri] = task.Uid
-	defer self.uidLock.Unlock()
+	l.Debugw("Updaing uid cache", "key", task.UidKey, "val", task.Uid)
+	self.uidCache[task.UidKey] = task.Uid
+	self.uidLock.Unlock()
+}
+
+func (self *Puller) isPathExcluded(path string) bool {
+	for _, pattern := range self.exclude {
+		matched, _ := doublestar.Match(pattern, path)
+		if matched {
+			return true
+		}
+	}
+	return false
 }
 
 func (self *Puller) handlePageList(
@@ -153,16 +189,9 @@ func (self *Puller) handlePageList(
 			continue
 		}
 		// ignore file that matches exclude rules
-		shouldSkip := false
-		for _, pattern := range self.exclude {
-			matched, _ := doublestar.Match(pattern, relPath)
-			if matched {
-				l.Debugf("skipped %s due to exclude pattern: %s", uri, pattern)
-				shouldSkip = true
-				break
-			}
-		}
+		shouldSkip := self.isPathExcluded(relPath)
 		if shouldSkip {
+			l.Debugf("skipped %s due to exclude pattern", uri)
 			continue
 		}
 
@@ -178,10 +207,11 @@ func (self *Puller) handlePageList(
 
 		self.fileListedCnt += 1
 
+		uidKey := relPath
 		self.uidLock.Lock()
-		oldUid, ok := self.uidCache[uri]
+		oldUid, ok := self.uidCache[uidKey]
 		self.uidLock.Unlock()
-		l.Debugf("Comparing object UID: %s <> %s = %v", oldUid, newUid, oldUid == newUid)
+		l.Debugf("Comparing object UID: %s <> %s", oldUid, newUid)
 		if ok && oldUid == newUid {
 			// skip update if uid is the same
 			continue
@@ -192,6 +222,7 @@ func (self *Puller) handlePageList(
 			Uri:       uri,
 			LocalPath: localPath,
 			Uid:       newUid,
+			UidKey:    uidKey,
 		}
 	}
 	return true
@@ -219,8 +250,10 @@ type Puller struct {
 	filePulledCnt int
 }
 
-func (self *Puller) AddExcludePattern(pattern string) {
-	self.exclude = append(self.exclude, pattern)
+func (self *Puller) AddExcludePatterns(patterns []string) {
+	for _, pattern := range patterns {
+		self.exclude = append(self.exclude, pattern)
+	}
 }
 
 func (self *Puller) Pull(remoteUri string, localDir string) string {
@@ -321,6 +354,76 @@ func (self *Puller) Pull(remoteUri string, localDir string) string {
 	}
 }
 
+func (self *Puller) PopulateChecksum(localDir string) {
+	l := zap.S()
+
+	setFileChecksum := func(path string) {
+		f, err := os.Open(path)
+		if err != nil {
+			l.Errorf("Invalid file path for checksum calculation: %s, err: %s", path, err)
+		}
+		defer f.Close()
+
+		h := md5.New()
+		if _, err := io.Copy(h, f); err != nil {
+			l.Errorf("Failed to calculate checksum for file: %s, err: %s", path, err)
+		}
+
+		uidKey, err := uidKeyFromLocalPath(localDir, path)
+		if err != nil {
+			l.Errorf("Failed to calculate uidKey for file: %s under dir: %s, err: %s", path, localDir, err)
+			return
+		}
+
+		uid, err := uidFromLocalPath(path)
+		if err != nil {
+			l.Errorf("Failed to calculate UID: %s", err)
+			return
+		}
+
+		self.uidLock.Lock()
+		self.uidCache[uidKey] = uid
+		self.uidLock.Unlock()
+	}
+
+	err := filepath.Walk(localDir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		// ignore file that matches exclude rules
+		shouldSkip := false
+		relPath, err := filepath.Rel(localDir, path)
+		if err != nil {
+			l.Errorf("Got invalid path from filepath.Walk: %s, err: %s", path, err)
+			shouldSkip = true
+		} else {
+			if info.IsDir() {
+				// this is so that pattern `foo/**` also matches `foo`
+				relPath += "/"
+			}
+			shouldSkip = self.isPathExcluded(relPath)
+		}
+
+		if info.IsDir() {
+			if shouldSkip {
+				return filepath.SkipDir
+			}
+		} else {
+			if shouldSkip {
+				return nil
+			}
+
+			setFileChecksum(path)
+		}
+		return nil
+	})
+
+	if err != nil {
+		l.Errorf("Failed to walk directory for populating file checksum, err: %s", err)
+	}
+}
+
 func NewPuller() *Puller {
 	return &Puller{
 		workerCnt: 5,
diff --git a/pkg/sync/pull_internal_test.go b/pkg/sync/pull_internal_test.go
@@ -119,7 +119,7 @@ func TestDeleteStaleFile(t *testing.T) {
 func TestSkipObjectsWithoutChange(t *testing.T) {
 	p := NewPuller()
 	p.taskQueue = make(chan DownloadTask, 10)
-	p.uidCache["s3://foo/home/dags/b.file"] = "1"
+	p.uidCache["b.file"] = "\"1\""
 
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -135,11 +135,11 @@ func TestSkipObjectsWithoutChange(t *testing.T) {
 			Contents: []*s3.Object{
 				&s3.Object{
 					Key:  aws.String("home/dags/b.file"),
-					ETag: aws.String("1"),
+					ETag: aws.String("\"1\""),
 				},
 				&s3.Object{
 					Key:  aws.String("home/dags/bar/a.go"),
-					ETag: aws.String("1"),
+					ETag: aws.String("\"1\""),
 				},
 			},
 		},
@@ -168,31 +168,29 @@ func TestSkipExcludedObjects(t *testing.T) {
 		wg.Done()
 	}()
 
-	p.AddExcludePattern("airflow.cfg")
-	p.AddExcludePattern("webserver_config.py")
-	p.AddExcludePattern("config/**")
+	p.AddExcludePatterns([]string{"airflow.cfg", "webserver_config.py", "config/**"})
 	p.handlePageList(
 		&s3.ListObjectsV2Output{
 			Contents: []*s3.Object{
 				&s3.Object{
 					Key:  aws.String("home/dags/b.file"),
-					ETag: aws.String("1"),
+					ETag: aws.String("\"1\""),
 				},
 				&s3.Object{
 					Key:  aws.String("home/airflow.cfg"),
-					ETag: aws.String("2"),
+					ETag: aws.String("\"2\""),
 				},
 				&s3.Object{
 					Key:  aws.String("home/config/a.file"),
-					ETag: aws.String("3"),
+					ETag: aws.String("\"3\""),
 				},
 				&s3.Object{
 					Key:  aws.String("home/config/subdir/a.file"),
-					ETag: aws.String("4"),
+					ETag: aws.String("\"4\""),
 				},
 				&s3.Object{
 					Key:  aws.String("home/webserver_config.py"),
-					ETag: aws.String("5"),
+					ETag: aws.String("\"5\""),
 				},
 			},
 		},
