[disk] drain any remaining data before Put returns

io.PipeWriter Write calls block until all the data written in the
call is read from the corresponding io.PipeReader. If we don't read
all that data, then the writing goroutine will block forever.

The PipeWriter used in bytestream Write calls is intended to be
consumed by disk.Put(), but if that returns early then there will
be blocked writes. To un-block them, we should ensure that any
remaining data is drained before disk.Put returns.

Backport of #474

Author: mostynb

cache/disk/disk.go

@@ -223,8 +223,15 @@ func (c *Cache) loadExistingFiles() error {
 
 // Put stores a stream of `size` bytes from `r` into the cache.
 // If `hash` is not the empty string, and the contents don't match it,
-// a non-nil error is returned.
+// a non-nil error is returned. All data will be read from `r` before
+// this function returns.
 func (c *Cache) Put(kind cache.EntryKind, hash string, size int64, r io.Reader) (rErr error) {
+	defer func() {
+		if r != nil {
+			_, _ = io.Copy(ioutil.Discard, r)
+		}
+	}()
+
 	if size < 0 {
 		return fmt.Errorf("Invalid (negative) size: %d", size)
 	}
@@ -237,7 +244,6 @@ func (c *Cache) Put(kind cache.EntryKind, hash string, size int64, r io.Reader)
 	}
 
 	if kind == cache.CAS && size == 0 && hash == emptySha256 {
-		io.Copy(ioutil.Discard, r)
 		return nil
 	}
 
@@ -304,6 +310,8 @@ func (c *Cache) Put(kind cache.EntryKind, hash string, size int64, r io.Reader)
 		return err
 	}
 
+	r = nil // We read all the data from r.
+
 	if c.proxy != nil {
 		rc, err := os.Open(tf.Name())
 		if err != nil {