X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/40103b455d0df2bc47eaba422a72bf6ec4fd6808..6284bed37f6b2bb3c0671b401aa0d99d5ef7a5b7:/services/keepstore/handlers.go?ds=sidebyside

diff --git a/services/keepstore/handlers.go b/services/keepstore/handlers.go
index 043ab69b17..a6798a9f72 100644
--- a/services/keepstore/handlers.go
+++ b/services/keepstore/handlers.go
@@ -73,24 +73,67 @@ func BadRequestHandler(w http.ResponseWriter, r *http.Request) {
 func GetBlockHandler(resp http.ResponseWriter, req *http.Request) {
 	if enforcePermissions {
 		locator := req.URL.Path[1:] // strip leading slash
-		if err := VerifySignature(locator, GetApiToken(req)); err != nil {
+		if err := VerifySignature(locator, GetAPIToken(req)); err != nil {
 			http.Error(resp, err.Error(), err.(*KeepError).HTTPCode)
 			return
 		}
 	}
 
-	block, err := GetBlock(mux.Vars(req)["hash"])
+	// TODO: Probe volumes to check whether the block _might_
+	// exist. Some volumes/types could support a quick existence
+	// check without causing other operations to suffer. If all
+	// volumes support that, and assure us the block definitely
+	// isn't here, we can return 404 now instead of waiting for a
+	// buffer.
+
+	buf, err := getBufferForResponseWriter(resp, bufs, BlockSize)
+	if err != nil {
+		http.Error(resp, err.Error(), http.StatusServiceUnavailable)
+		return
+	}
+	defer bufs.Put(buf)
+
+	size, err := GetBlock(mux.Vars(req)["hash"], buf, resp)
 	if err != nil {
-		// This type assertion is safe because the only errors
-		// GetBlock can return are DiskHashError or NotFoundError.
-		http.Error(resp, err.Error(), err.(*KeepError).HTTPCode)
+		code := http.StatusInternalServerError
+		if err, ok := err.(*KeepError); ok {
+			code = err.HTTPCode
+		}
+		http.Error(resp, err.Error(), code)
 		return
 	}
-	defer bufs.Put(block)
 
-	resp.Header().Set("Content-Length", strconv.Itoa(len(block)))
+	resp.Header().Set("Content-Length", strconv.Itoa(size))
 	resp.Header().Set("Content-Type", "application/octet-stream")
-	resp.Write(block)
+	resp.Write(buf[:size])
+}
+
+// Get a buffer from the pool -- but give up and return a non-nil
+// error if resp implements http.CloseNotifier and tells us that the
+// client has disconnected before we get a buffer.
+func getBufferForResponseWriter(resp http.ResponseWriter, bufs *bufferPool, bufSize int) ([]byte, error) {
+	var closeNotifier <-chan bool
+	if resp, ok := resp.(http.CloseNotifier); ok {
+		closeNotifier = resp.CloseNotify()
+	}
+	var buf []byte
+	bufReady := make(chan []byte)
+	go func() {
+		bufReady <- bufs.Get(bufSize)
+		close(bufReady)
+	}()
+	select {
+	case buf = <-bufReady:
+		return buf, nil
+	case <-closeNotifier:
+		go func() {
+			// Even if closeNotifier happened first, we
+			// need to keep waiting for our buf so we can
+			// return it to the pool.
+			bufs.Put(<-bufReady)
+		}()
+		return nil, ErrClientDisconnect
+	}
 }
 
 // PutBlockHandler is a HandleFunc to address Put block requests.
@@ -116,8 +159,13 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	buf := bufs.Get(int(req.ContentLength))
-	_, err := io.ReadFull(req.Body, buf)
+	buf, err := getBufferForResponseWriter(resp, bufs, int(req.ContentLength))
+	if err != nil {
+		http.Error(resp, err.Error(), http.StatusServiceUnavailable)
+		return
+	}
+
+	_, err = io.ReadFull(req.Body, buf)
 	if err != nil {
 		http.Error(resp, err.Error(), 500)
 		bufs.Put(buf)
@@ -136,7 +184,7 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
 	// Success; add a size hint, sign the locator if possible, and
 	// return it to the client.
 	returnHash := fmt.Sprintf("%s+%d", hash, req.ContentLength)
-	apiToken := GetApiToken(req)
+	apiToken := GetAPIToken(req)
 	if PermissionSecret != nil && apiToken != "" {
 		expiry := time.Now().Add(blobSignatureTTL)
 		returnHash = SignLocator(returnHash, apiToken, expiry)
@@ -148,7 +196,7 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
 // IndexHandler is a HandleFunc to address /index and /index/{prefix} requests.
 func IndexHandler(resp http.ResponseWriter, req *http.Request) {
 	// Reject unauthorized requests.
-	if !IsDataManagerToken(GetApiToken(req)) {
+	if !IsDataManagerToken(GetAPIToken(req)) {
 		http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
 		return
 	}
@@ -280,7 +328,7 @@ func DeleteHandler(resp http.ResponseWriter, req *http.Request) {
 	hash := mux.Vars(req)["hash"]
 
 	// Confirm that this user is an admin and has a token with unlimited scope.
-	var tok = GetApiToken(req)
+	var tok = GetAPIToken(req)
 	if tok == "" || !CanDelete(tok) {
 		http.Error(resp, PermissionError.Error(), PermissionError.HTTPCode)
 		return
@@ -371,7 +419,7 @@ type PullRequest struct {
 // PullHandler processes "PUT /pull" requests for the data manager.
 func PullHandler(resp http.ResponseWriter, req *http.Request) {
 	// Reject unauthorized requests.
-	if !IsDataManagerToken(GetApiToken(req)) {
+	if !IsDataManagerToken(GetAPIToken(req)) {
 		http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
 		return
 	}
@@ -407,7 +455,7 @@ type TrashRequest struct {
 // TrashHandler processes /trash requests.
 func TrashHandler(resp http.ResponseWriter, req *http.Request) {
 	// Reject unauthorized requests.
-	if !IsDataManagerToken(GetApiToken(req)) {
+	if !IsDataManagerToken(GetAPIToken(req)) {
 		http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
 		return
 	}
@@ -437,7 +485,7 @@ func TrashHandler(resp http.ResponseWriter, req *http.Request) {
 // UntrashHandler processes "PUT /untrash/{hash:[0-9a-f]{32}}" requests for the data manager.
 func UntrashHandler(resp http.ResponseWriter, req *http.Request) {
 	// Reject unauthorized requests.
-	if !IsDataManagerToken(GetApiToken(req)) {
+	if !IsDataManagerToken(GetAPIToken(req)) {
 		http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
 		return
 	}
@@ -481,7 +529,6 @@ func UntrashHandler(resp http.ResponseWriter, req *http.Request) {
 	}
 }
 
-// ==============================
 // GetBlock and PutBlock implement lower-level code for handling
 // blocks by rooting through volumes connected to the local machine.
 // Once the handler has determined that system policy permits the
@@ -492,24 +539,21 @@ func UntrashHandler(resp http.ResponseWriter, req *http.Request) {
 // should be the only part of the code that cares about which volume a
 // block is stored on, so it should be responsible for figuring out
 // which volume to check for fetching blocks, storing blocks, etc.
-// ==============================
 
-// GetBlock fetches and returns the block identified by "hash".
-//
-// On success, GetBlock returns a byte slice with the block data, and
-// a nil error.
+// GetBlock fetches the block identified by "hash" into the provided
+// buf, and returns the data size.
 //
 // If the block cannot be found on any volume, returns NotFoundError.
 //
 // If the block found does not have the correct MD5 hash, returns
 // DiskHashError.
 //
-func GetBlock(hash string) ([]byte, error) {
+func GetBlock(hash string, buf []byte, resp http.ResponseWriter) (int, error) {
 	// Attempt to read the requested hash from a keep volume.
 	errorToCaller := NotFoundError
 
 	for _, vol := range KeepVM.AllReadable() {
-		buf, err := vol.Get(hash)
+		size, err := vol.Get(hash, buf)
 		if err != nil {
 			// IsNotExist is an expected error and may be
 			// ignored. All other errors are logged. In
@@ -523,23 +567,22 @@ func GetBlock(hash string) ([]byte, error) {
 		}
 		// Check the file checksum.
 		//
-		filehash := fmt.Sprintf("%x", md5.Sum(buf))
+		filehash := fmt.Sprintf("%x", md5.Sum(buf[:size]))
 		if filehash != hash {
 			// TODO: Try harder to tell a sysadmin about
 			// this.
 			log.Printf("%s: checksum mismatch for request %s (actual %s)",
 				vol, hash, filehash)
 			errorToCaller = DiskHashError
-			bufs.Put(buf)
 			continue
 		}
 		if errorToCaller == DiskHashError {
 			log.Printf("%s: checksum mismatch for request %s but a good copy was found on another volume and returned",
 				vol, hash)
 		}
-		return buf, nil
+		return size, nil
 	}
-	return nil, errorToCaller
+	return 0, errorToCaller
 }
 
 // PutBlock Stores the BLOCK (identified by the content id HASH) in Keep.
@@ -671,10 +714,10 @@ func IsValidLocator(loc string) bool {
 
 var authRe = regexp.MustCompile(`^OAuth2\s+(.*)`)
 
-// GetApiToken returns the OAuth2 token from the Authorization
+// GetAPIToken returns the OAuth2 token from the Authorization
 // header of a HTTP request, or an empty string if no matching
 // token is found.
-func GetApiToken(req *http.Request) string {
+func GetAPIToken(req *http.Request) string {
 	if auth, ok := req.Header["Authorization"]; ok {
 		if match := authRe.FindStringSubmatch(auth[0]); match != nil {
 			return match[1]