Removed panic and added logging

pkg/snapshot/layered_map.go

@@ -18,10 +18,10 @@ package snapshot
 
 type LayeredMap struct {
 	layers []map[string]string
-	hasher func(string) string
+	hasher func(string) (string, error)
 }
 
-func NewLayeredMap(h func(string) string) *LayeredMap {
+func NewLayeredMap(h func(string) (string, error)) *LayeredMap {
 	l := LayeredMap{
 		hasher: h,
 	}
@@ -42,12 +42,15 @@ func (l *LayeredMap) Get(s string) (string, bool) {
 	return "", false
 }
 
-func (l *LayeredMap) MaybeAdd(s string) bool {
+func (l *LayeredMap) MaybeAdd(s string) (bool, error) {
 	oldV, ok := l.Get(s)
-	newV := l.hasher(s)
+	newV, err := l.hasher(s)
+	if err != nil {
+		return false, err
+	}
 	if ok && newV == oldV {
-		return false
+		return false, nil
 	}
 	l.layers[len(l.layers)-1][s] = newV
-	return true
+	return true, nil
 }

pkg/snapshot/snapshot.go

@@ -32,12 +32,11 @@ import (
 type Snapshotter struct {
 	l         *LayeredMap
 	directory string
-	snapshots []string
 }
 
 // NewSnapshotter creates a new snapshotter rooted at d
 func NewSnapshotter(l *LayeredMap, d string) *Snapshotter {
-	return &Snapshotter{l: l, directory: d, snapshots: []string{}}
+	return &Snapshotter{l: l, directory: d}
 }
 
 // Init initializes a new snapshotter
@@ -48,51 +47,76 @@ func (s *Snapshotter) Init() error {
 	return nil
 }
 
-// TakeSnapshot takes a snapshot of the filesystem, avoiding directories in the whitelist
+// TakeSnapshot takes a snapshot of the filesystem, avoiding directories in the whitelist, and creates
-// It stores changed files in a tar, and returns the contents of this tar at the end
+// a tarball of the changed files. Return contents of the tarball, and whether or not any files were changed
-func (s *Snapshotter) TakeSnapshot() ([]byte, error) {
+func (s *Snapshotter) TakeSnapshot() ([]byte, bool, error) {
-
 	buf := bytes.NewBuffer([]byte{})
-	added, err := s.snapShotFS(buf)
+	filesAdded, err := s.snapShotFS(buf)
 	if err != nil {
-		return nil, err
+		return nil, filesAdded, err
 	}
-	if !added {
+	contents, err := ioutil.ReadAll(buf)
-		logrus.Infof("No files were changed in this command, this layer will not be appended.")
+	if err != nil {
+		return nil, filesAdded, err
+	}
+	return contents, filesAdded, err
+}
+
+// TakeSnapshotOfFiles takes a snapshot of specific files
+// Used for ADD/COPY commands, when we know which files have changed
+func (s *Snapshotter) TakeSnapshotOfFiles(files []string) ([]byte, error) {
+	logrus.Infof("Taking snapshot of files %s", files)
+	s.l.Snapshot()
+	if len(files) == 0 {
+		logrus.Info("No files changed in this command, skipping snapshotting.")
 		return nil, nil
 	}
-	if err != nil {
+	buf := bytes.NewBuffer([]byte{})
-		return nil, err
+	w := tar.NewWriter(buf)
-	}
+	defer w.Close()
-	// Add buffer contents until buffer is empty
+	for _, file := range files {
-	var contents []byte
+		info, err := os.Stat(file)
-	for {
+		if err != nil {
-		next := buf.Next(buf.Len())
+			return nil, err
-		if len(next) == 0 {
+		}
-			break
+		if util.PathInWhitelist(file, s.directory) {
+			logrus.Debugf("Not adding %s to layer, as it is whitelisted", file)
+			continue
+		}
+		// Only add to the tar if we add it to the layeredmap.
+		maybeAdd, err := s.l.MaybeAdd(file)
+		if err != nil {
+			return nil, err
+		}
+		if maybeAdd {
+			util.AddToTar(file, info, w)
 		}
-		contents = append(contents, next...)
 	}
-	return contents, nil
+	return ioutil.ReadAll(buf)
 }
 
 func (s *Snapshotter) snapShotFS(f io.Writer) (bool, error) {
 	s.l.Snapshot()
-	added := false
+	filesAdded := false
 	w := tar.NewWriter(f)
 	defer w.Close()
 
 	err := filepath.Walk(s.directory, func(path string, info os.FileInfo, err error) error {
 		if util.PathInWhitelist(path, s.directory) {
+			logrus.Debugf("Not adding %s to layer, as it's whitelisted", path)
 			return nil
 		}
 
 		// Only add to the tar if we add it to the layeredmap.
-		if s.l.MaybeAdd(path) {
+		maybeAdd, err := s.l.MaybeAdd(path)
-			added = true
+		if err != nil {
+			return err
+		}
+		if maybeAdd {
+			filesAdded = true
 			return util.AddToTar(path, info, w)
 		}
 		return nil
 	})
-	return added, err
+	return filesAdded, err
 }

pkg/snapshot/snapshot_test.go

@@ -38,28 +38,36 @@ func TestSnapshotFileChange(t *testing.T) {
 	// Make some changes to the filesystem
 	newFiles := map[string]string{
 		"foo":           "newbaz1",
+		"bar/bat":       "baz",
 		"workspace/bat": "bat",
 	}
 	if err := testutil.SetupFiles(testDir, newFiles); err != nil {
 		t.Fatalf("Error setting up fs: %s", err)
 	}
 	// Take another snapshot
-	contents, err := snapshotter.TakeSnapshot()
+	contents, filesAdded, err := snapshotter.TakeSnapshot()
 	if err != nil {
 		t.Fatalf("Error taking snapshot of fs: %s", err)
 	}
+	if !filesAdded {
+		t.Fatal("No files added to snapshot.")
+	}
 	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
 	reader := bytes.NewReader(contents)
 	tr := tar.NewReader(reader)
 	fooPath := filepath.Join(testDir, "foo")
+	batPath := filepath.Join(testDir, "bar/bat")
 	snapshotFiles := map[string]string{
 		fooPath: "newbaz1",
+		batPath: "baz",
 	}
+	numFiles := 0
 	for {
 		hdr, err := tr.Next()
 		if err == io.EOF {
 			break
 		}
+		numFiles++
 		if _, isFile := snapshotFiles[hdr.Name]; !isFile {
 			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
 		}
@@ -68,6 +76,9 @@ func TestSnapshotFileChange(t *testing.T) {
 			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, snapshotFiles[hdr.Name], string(contents))
 		}
 	}
+	if numFiles != 2 {
+		t.Fatalf("Incorrect number of files were added, expected: 2, actual: %v", numFiles)
+	}
 }
 
 func TestSnapshotChangePermissions(t *testing.T) {
@@ -82,21 +93,26 @@ func TestSnapshotChangePermissions(t *testing.T) {
 		t.Fatalf("Error changing permissions on %s: %v", batPath, err)
 	}
 	// Take another snapshot
-	contents, err := snapshotter.TakeSnapshot()
+	contents, filesAdded, err := snapshotter.TakeSnapshot()
 	if err != nil {
 		t.Fatalf("Error taking snapshot of fs: %s", err)
 	}
+	if !filesAdded {
+		t.Fatal("No files added to snapshot.")
+	}
 	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
 	reader := bytes.NewReader(contents)
 	tr := tar.NewReader(reader)
 	snapshotFiles := map[string]string{
 		batPath: "baz2",
 	}
+	numFiles := 0
 	for {
 		hdr, err := tr.Next()
 		if err == io.EOF {
 			break
 		}
+		numFiles++
 		if _, isFile := snapshotFiles[hdr.Name]; !isFile {
 			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
 		}
@@ -105,6 +121,57 @@ func TestSnapshotChangePermissions(t *testing.T) {
 			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, snapshotFiles[hdr.Name], string(contents))
 		}
 	}
+	if numFiles != 1 {
+		t.Fatalf("Incorrect number of files were added, expected: 1, got: %v", numFiles)
+	}
+}
+
+func TestSnapshotFiles(t *testing.T) {
+	testDir, snapshotter, err := setUpTestDir()
+	defer os.RemoveAll(testDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Make some changes to the filesystem
+	newFiles := map[string]string{
+		"foo":            "newbaz1",
+		"workspace/file": "bat",
+	}
+	if err := testutil.SetupFiles(testDir, newFiles); err != nil {
+		t.Fatalf("Error setting up fs: %s", err)
+	}
+	filesToSnapshot := []string{
+		filepath.Join(testDir, "foo"),
+		filepath.Join(testDir, "workspace/file"),
+	}
+	contents, err := snapshotter.TakeSnapshotOfFiles(filesToSnapshot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedContents := map[string]string{
+		filepath.Join(testDir, "foo"): "newbaz1",
+	}
+	// Check contents of the snapshot, make sure contents is equivalent to snapshotFiles
+	reader := bytes.NewReader(contents)
+	tr := tar.NewReader(reader)
+	numFiles := 0
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		numFiles = numFiles + 1
+		if _, isFile := expectedContents[hdr.Name]; !isFile {
+			t.Fatalf("File %s unexpectedly in tar", hdr.Name)
+		}
+		contents, _ := ioutil.ReadAll(tr)
+		if string(contents) != expectedContents[hdr.Name] {
+			t.Fatalf("Contents of %s incorrect, expected: %s, actual: %s", hdr.Name, expectedContents[hdr.Name], string(contents))
+		}
+	}
+	if numFiles != 1 {
+		t.Fatalf("%s was not added.", filepath.Join(testDir, "foo"))
+	}
 }
 
 func TestEmptySnapshot(t *testing.T) {
@@ -114,13 +181,13 @@ func TestEmptySnapshot(t *testing.T) {
 		t.Fatal(err)
 	}
 	// Take snapshot with no changes
-	contents, err := snapshotter.TakeSnapshot()
+	_, filesAdded, err := snapshotter.TakeSnapshot()
 	if err != nil {
 		t.Fatalf("Error taking snapshot of fs: %s", err)
 	}
 	// Since we took a snapshot with no changes, contents should be nil
-	if contents != nil {
+	if filesAdded {
-		t.Fatal("Contents should be nil, since no changes to the filesystem were made.")
+		t.Fatal("Files added even though no changes to file system were made.")
 	}
 }
 

pkg/util/util.go

@@ -36,12 +36,12 @@ func SetLogLevel(logLevel string) error {
 }
 
 // Hasher returns a hash function, used in snapshotting to determine if a file has changed
-func Hasher() func(string) string {
+func Hasher() func(string) (string, error) {
-	hasher := func(p string) string {
+	hasher := func(p string) (string, error) {
 		h := md5.New()
 		fi, err := os.Lstat(p)
 		if err != nil {
-			panic(err)
+			return "", err
 		}
 		h.Write([]byte(fi.Mode().String()))
 		h.Write([]byte(fi.ModTime().String()))
@@ -49,15 +49,15 @@ func Hasher() func(string) string {
 		if fi.Mode().IsRegular() {
 			f, err := os.Open(p)
 			if err != nil {
-				panic(err)
+				return "", err
 			}
 			defer f.Close()
 			if _, err := io.Copy(h, f); err != nil {
-				panic(err)
+				return "", err
 			}
 		}
 
-		return hex.EncodeToString(h.Sum(nil))
+		return hex.EncodeToString(h.Sum(nil)), nil
 	}
 	return hasher
 }