12447: Require Content-Length in KeepClient.Get

[arvados.git] / sdk / go / keepclient / block_cache.go

diff --git a/sdk/go/keepclient/block_cache.go b/sdk/go/keepclient/block_cache.go
index 23fd331d1f1282446e3d9dda3881916ec80c3cb0..5d8daad60d4be6dfad3a6f3c2003f759f2967956 100644 (file)
--- a/sdk/go/keepclient/block_cache.go
+++ b/sdk/go/keepclient/block_cache.go
@@ -5,7 +5,6 @@
 package keepclient
 
 import (
-       "bytes"
        "io"
        "sort"
        "sync"
@@ -65,17 +64,18 @@ func (c *BlockCache) Get(kc *KeepClient, locator string) ([]byte, error) {
                }
                c.cache[cacheKey] = b
                go func() {
-                       rdr, _, _, err := kc.Get(locator)
-                       data := bytes.NewBuffer(make([]byte, 0, BLOCKSIZE))
+                       rdr, size, _, err := kc.Get(locator)
+                       var data []byte
                        if err == nil {
-                               _, err = io.Copy(data, rdr)
+                               data = make([]byte, size, BLOCKSIZE)
+                               _, err = io.ReadFull(rdr, data)
                                err2 := rdr.Close()
                                if err == nil {
                                        err = err2
                                }
                        }
                        c.mtx.Lock()
-                       b.data, b.err = data.Bytes(), err
+                       b.data, b.err = data, err
                        c.mtx.Unlock()
                        close(b.fetched)
                        go c.Sweep()
@@ -94,13 +94,13 @@ func (c *BlockCache) Get(kc *KeepClient, locator string) ([]byte, error) {
 }
 
 func (c *BlockCache) setup() {
+       c.mtx.Lock()
        c.cache = make(map[string]*cacheBlock)
+       c.mtx.Unlock()
 }
 
 func (c *BlockCache) Clear() {
-       c.mtx.Lock()
        c.setup()
-       c.mtx.Unlock()
 }
 
 type timeSlice []time.Time