From 432784dae32aa3f0cb3746b1592b760ce5ef4dd2 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Tue, 27 Mar 2018 11:14:53 -0700
Subject: [PATCH] Fix alloc watcher snapshot streaming

---
 client/alloc_watcher.go              | 13 ++++++++-----
 command/agent/alloc_endpoint_test.go |  2 +-
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/client/alloc_watcher.go b/client/alloc_watcher.go
index e15d4f88c..e84db9ed2 100644
--- a/client/alloc_watcher.go
+++ b/client/alloc_watcher.go
@@ -489,7 +489,7 @@ func (p *remotePrevAlloc) streamAllocDir(ctx context.Context, resp io.ReadCloser
 			// Error snapshotting on the remote side, try to read
 			// the message out of the file and return it.
 			errBuf := make([]byte, int(hdr.Size))
-			if _, err := tr.Read(errBuf); err != nil {
+			if _, err := tr.Read(errBuf); err != nil && err != io.EOF {
 				return fmt.Errorf("error streaming previous alloc %q for new alloc %q; failed reading error message: %v",
 					p.prevAllocID, p.allocID, err)
 			}
@@ -542,6 +542,13 @@ func (p *remotePrevAlloc) streamAllocDir(ctx context.Context, resp io.ReadCloser
 			// is still alive
 			for !canceled() {
 				n, err := tr.Read(buf)
+				if n > 0 && (err == nil || err == io.EOF) {
+					if _, err := f.Write(buf[:n]); err != nil {
+						f.Close()
+						return fmt.Errorf("error writing to file %q: %v", f.Name(), err)
+					}
+				}
+
 				if err != nil {
 					f.Close()
 					if err != io.EOF {
@@ -549,10 +556,6 @@ func (p *remotePrevAlloc) streamAllocDir(ctx context.Context, resp io.ReadCloser
 					}
 					break
 				}
-				if _, err := f.Write(buf[:n]); err != nil {
-					f.Close()
-					return fmt.Errorf("error writing to file %q: %v", f.Name(), err)
-				}
 			}
 
 		}
diff --git a/command/agent/alloc_endpoint_test.go b/command/agent/alloc_endpoint_test.go
index 5a92bf552..10b6f249a 100644
--- a/command/agent/alloc_endpoint_test.go
+++ b/command/agent/alloc_endpoint_test.go
@@ -518,7 +518,7 @@ func TestHTTP_AllocSnapshot_Atomic(t *testing.T) {
 				// Found it!
 				markerFound = true
 				buf := make([]byte, int(header.Size))
-				if _, err := r.Read(buf); err != nil {
+				if _, err := r.Read(buf); err != nil && err != io.EOF {
 					t.Errorf("Unexpected error reading error marker %s: %v", errorFilename, err)
 				} else {
 					markerContents = string(buf)