From ff2e1c6e9cc09266ddccc92ef7b52c6fc7093a6c Mon Sep 17 00:00:00 2001
From: Cenk Alti <cenkalti@gmail.com>
Date: Fri, 6 Nov 2020 16:42:30 +0300
Subject: [PATCH] fix drain test

---
 client.go       |  1 +
 drain.go        |  6 ++++++
 drain_test.go   |  1 +
 filereceiver.go | 18 ++++++++++--------
 write.go        |  8 ++++++++
 5 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/client.go b/client.go
index 31a55be..1d325e7 100644
--- a/client.go
+++ b/client.go
@@ -18,6 +18,7 @@ type Client struct {
 	log        log.Logger
 	trackerURL *url.URL
 	httpClient http.Client
+	drainer    bool
 }
 
 // NewClient creates a new Client.
diff --git a/drain.go b/drain.go
index f6d1ae3..33d31b1 100644
--- a/drain.go
+++ b/drain.go
@@ -20,6 +20,8 @@ type Drainer struct {
 	log      log.Logger
 	shutdown chan struct{}
 	stopped  chan struct{}
+
+	stopOnError bool
 }
 
 func NewDrainer(c *Config) (*Drainer, error) {
@@ -42,6 +44,7 @@ func NewDrainer(c *Config) (*Drainer, error) {
 	if err != nil {
 		return nil, err
 	}
+	clt.drainer = true
 	logger := log.NewLogger("drain")
 	d := &Drainer{
 		config:   c,
@@ -91,6 +94,9 @@ func (d *Drainer) Run() error {
 		d.log.Infof("moving fid=%v; %v of %v (%v%%)", fid, i+1, len(fids), ((i+1)*100)/len(fids))
 		if err = d.moveFile(fid); err != nil {
 			d.log.Error(err)
+			if d.stopOnError {
+				return err
+			}
 		}
 	}
 	return nil
diff --git a/drain_test.go b/drain_test.go
index 71f0627..4061629 100644
--- a/drain_test.go
+++ b/drain_test.go
@@ -102,6 +102,7 @@ func TestDrain(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	dr.stopOnError = true
 	err = dr.Run()
 	if err != nil {
 		t.Fatal(err)
diff --git a/filereceiver.go b/filereceiver.go
index 331c385..d4ea122 100644
--- a/filereceiver.go
+++ b/filereceiver.go
@@ -69,14 +69,16 @@ func (f *FileReceiver) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 				return
 			}
 		}
-		ok, err := f.tempfileExists(path)
-		if err != nil {
-			f.internalServerError("cannot check tempfile", err, r, w)
-			return
-		}
-		if !ok {
-			http.Error(w, "tempfile does not exist", http.StatusNotFound)
-			return
+		if r.Header.Get("efes-drain") == "" {
+			ok, err := f.tempfileExists(path)
+			if err != nil {
+				f.internalServerError("cannot check tempfile", err, r, w)
+				return
+			}
+			if !ok {
+				http.Error(w, "tempfile does not exist", http.StatusNotFound)
+				return
+			}
 		}
 		newOffset, digest, err := saveFile(path, offset, length, r.Body, f.log)
 		if oerr, ok := err.(*OffsetMismatchError); ok {
diff --git a/write.go b/write.go
index 429b452..e78bbbf 100644
--- a/write.go
+++ b/write.go
@@ -95,7 +95,11 @@ func (c *Client) sendFile(path string, rs io.ReadSeeker, size int64) (*Checksums
 			}
 		}
 		checksums, err = c.send(path, r, offset, size, bo)
+		if cerr, ok := err.(*ClientError); ok && cerr.Code == 404 {
+			return backoff.Permanent(cerr)
+		}
 		if err != nil {
+			c.log.Errorf("cannot send chunk: %s", err.Error())
 			return err
 		}
 		remoteSha1, err = hex.DecodeString(checksums.Sha1)
@@ -156,10 +160,14 @@ func (c *Client) patch(path string, body io.Reader, offset, size int64) (*http.R
 	if size > -1 {
 		req.Header.Add("efes-file-length", strconv.FormatInt(size, 10))
 	}
+	if c.drainer {
+		req.Header.Add("efes-drain", "true")
+	}
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
+	resp.Body.Close()
 	return resp, checkResponseError(resp)
 }