X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/c7399ec7afdf0cfdd0f3177f410f102083a26e15..8d73164d5ca50e7af1b3752b4251eaf9f11a9fc8:/sdk/go/keepclient/keepclient.go

diff --git a/sdk/go/keepclient/keepclient.go b/sdk/go/keepclient/keepclient.go
index 54a4a374b9..169f1457e2 100644
--- a/sdk/go/keepclient/keepclient.go
+++ b/sdk/go/keepclient/keepclient.go
@@ -22,6 +22,7 @@ import (
 
 	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
 	"git.curoverse.com/arvados.git/sdk/go/asyncbuf"
+	"git.curoverse.com/arvados.git/sdk/go/httpserver"
 )
 
 // A Keep "block" is 64MB.
@@ -99,6 +100,8 @@ type KeepClient struct {
 	HTTPClient         HTTPClient
 	Retries            int
 	BlockCache         *BlockCache
+	RequestID          string
+	StorageClasses     []string
 
 	// set to 1 if all writable services are of disk type, otherwise 0
 	replicasPerService int
@@ -200,6 +203,17 @@ func (kc *KeepClient) getOrHead(method string, locator string) (io.ReadCloser, i
 		return ioutil.NopCloser(bytes.NewReader(nil)), 0, "", nil
 	}
 
+	reqid := kc.getRequestID()
+
+	var expectLength int64
+	if parts := strings.SplitN(locator, "+", 3); len(parts) < 2 {
+		expectLength = -1
+	} else if n, err := strconv.ParseInt(parts[1], 10, 64); err != nil {
+		expectLength = -1
+	} else {
+		expectLength = n
+	}
+
 	var errs []string
 
 	tries_remaining := 1 + kc.Retries
@@ -223,14 +237,17 @@ func (kc *KeepClient) getOrHead(method string, locator string) (io.ReadCloser, i
 				errs = append(errs, fmt.Sprintf("%s: %v", url, err))
 				continue
 			}
-			req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", kc.Arvados.ApiToken))
+			req.Header.Add("Authorization", "OAuth2 "+kc.Arvados.ApiToken)
+			req.Header.Add("X-Request-Id", reqid)
 			resp, err := kc.httpClient().Do(req)
 			if err != nil {
 				// Probably a network error, may be transient,
 				// can try again.
 				errs = append(errs, fmt.Sprintf("%s: %v", url, err))
 				retryList = append(retryList, host)
-			} else if resp.StatusCode != http.StatusOK {
+				continue
+			}
+			if resp.StatusCode != http.StatusOK {
 				var respbody []byte
 				respbody, _ = ioutil.ReadAll(&io.LimitedReader{R: resp.Body, N: 4096})
 				resp.Body.Close()
@@ -247,24 +264,29 @@ func (kc *KeepClient) getOrHead(method string, locator string) (io.ReadCloser, i
 				} else if resp.StatusCode == 404 {
 					count404++
 				}
-			} else if resp.ContentLength < 0 {
-				// Missing Content-Length
-				resp.Body.Close()
-				return nil, 0, "", fmt.Errorf("Missing Content-Length of block")
-			} else {
-				// Success.
-				if method == "GET" {
-					return HashCheckingReader{
-						Reader: resp.Body,
-						Hash:   md5.New(),
-						Check:  locator[0:32],
-					}, resp.ContentLength, url, nil
-				} else {
+				continue
+			}
+			if expectLength < 0 {
+				if resp.ContentLength < 0 {
 					resp.Body.Close()
-					return nil, resp.ContentLength, url, nil
+					return nil, 0, "", fmt.Errorf("error reading %q: no size hint, no Content-Length header in response", locator)
 				}
+				expectLength = resp.ContentLength
+			} else if resp.ContentLength >= 0 && expectLength != resp.ContentLength {
+				resp.Body.Close()
+				return nil, 0, "", fmt.Errorf("error reading %q: size hint %d != Content-Length %d", locator, expectLength, resp.ContentLength)
+			}
+			// Success
+			if method == "GET" {
+				return HashCheckingReader{
+					Reader: resp.Body,
+					Hash:   md5.New(),
+					Check:  locator[0:32],
+				}, expectLength, url, nil
+			} else {
+				resp.Body.Close()
+				return nil, expectLength, url, nil
 			}
-
 		}
 		serversToTry = retryList
 	}
@@ -334,7 +356,8 @@ func (kc *KeepClient) GetIndex(keepServiceUUID, prefix string) (io.Reader, error
 		return nil, err
 	}
 
-	req.Header.Add("Authorization", fmt.Sprintf("OAuth2 %s", kc.Arvados.ApiToken))
+	req.Header.Add("Authorization", "OAuth2 "+kc.Arvados.ApiToken)
+	req.Header.Set("X-Request-Id", kc.getRequestID())
 	resp, err := kc.httpClient().Do(req)
 	if err != nil {
 		return nil, err
@@ -523,6 +546,16 @@ func (kc *KeepClient) httpClient() HTTPClient {
 	return c
 }
 
+var reqIDGen = httpserver.IDGenerator{Prefix: "req-"}
+
+func (kc *KeepClient) getRequestID() string {
+	if kc.RequestID != "" {
+		return kc.RequestID
+	} else {
+		return reqIDGen.Next()
+	}
+}
+
 type Locator struct {
 	Hash  string
 	Size  int      // -1 if data size is not known