X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/e1902b8a0490aa6b7ffc544c1609d4d57a5110ce..b6bc9bbe45a68a07b1ea8139e1f4d698873739a7:/services/keepstore/handlers.go

diff --git a/services/keepstore/handlers.go b/services/keepstore/handlers.go
index 54b8b485e1..289dce15a0 100644
--- a/services/keepstore/handlers.go
+++ b/services/keepstore/handlers.go
@@ -9,6 +9,7 @@ package main
 
 import (
 	"container/list"
+	"context"
 	"crypto/md5"
 	"encoding/json"
 	"fmt"
@@ -71,6 +72,9 @@ func BadRequestHandler(w http.ResponseWriter, r *http.Request) {
 
 // GetBlockHandler is a HandleFunc to address Get block requests.
 func GetBlockHandler(resp http.ResponseWriter, req *http.Request) {
+	ctx, cancel := contextForResponse(context.TODO(), resp)
+	defer cancel()
+
 	if theConfig.RequireSignatures {
 		locator := req.URL.Path[1:] // strip leading slash
 		if err := VerifySignature(locator, GetAPIToken(req)); err != nil {
@@ -86,14 +90,14 @@ func GetBlockHandler(resp http.ResponseWriter, req *http.Request) {
 	// isn't here, we can return 404 now instead of waiting for a
 	// buffer.
 
-	buf, err := getBufferForResponseWriter(resp, bufs, BlockSize)
+	buf, err := getBufferWithContext(ctx, 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)
+	size, err := GetBlock(ctx, mux.Vars(req)["hash"], buf, resp)
 	if err != nil {
 		code := http.StatusInternalServerError
 		if err, ok := err.(*KeepError); ok {
@@ -108,24 +112,33 @@ func GetBlockHandler(resp http.ResponseWriter, req *http.Request) {
 	resp.Write(buf[:size])
 }
 
+// Return a new context that gets cancelled by resp's CloseNotifier.
+func contextForResponse(parent context.Context, resp http.ResponseWriter) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(parent)
+	if cn, ok := resp.(http.CloseNotifier); ok {
+		go func(c <-chan bool) {
+			select {
+			case <-c:
+				theConfig.debugLogf("cancel context")
+				cancel()
+			case <-ctx.Done():
+			}
+		}(cn.CloseNotify())
+	}
+	return ctx, cancel
+}
+
 // 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
+// error if ctx ends before we get a buffer.
+func getBufferWithContext(ctx context.Context, bufs *bufferPool, bufSize int) ([]byte, error) {
 	bufReady := make(chan []byte)
 	go func() {
 		bufReady <- bufs.Get(bufSize)
-		close(bufReady)
 	}()
 	select {
-	case buf = <-bufReady:
+	case buf := <-bufReady:
 		return buf, nil
-	case <-closeNotifier:
+	case <-ctx.Done():
 		go func() {
 			// Even if closeNotifier happened first, we
 			// need to keep waiting for our buf so we can
@@ -138,6 +151,9 @@ func getBufferForResponseWriter(resp http.ResponseWriter, bufs *bufferPool, bufS
 
 // PutBlockHandler is a HandleFunc to address Put block requests.
 func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
+	ctx, cancel := contextForResponse(context.TODO(), resp)
+	defer cancel()
+
 	hash := mux.Vars(req)["hash"]
 
 	// Detect as many error conditions as possible before reading
@@ -159,7 +175,7 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	buf, err := getBufferForResponseWriter(resp, bufs, int(req.ContentLength))
+	buf, err := getBufferWithContext(ctx, bufs, int(req.ContentLength))
 	if err != nil {
 		http.Error(resp, err.Error(), http.StatusServiceUnavailable)
 		return
@@ -172,12 +188,15 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	replication, err := PutBlock(buf, hash)
+	replication, err := PutBlock(ctx, buf, hash)
 	bufs.Put(buf)
 
 	if err != nil {
-		ke := err.(*KeepError)
-		http.Error(resp, ke.Error(), ke.HTTPCode)
+		code := http.StatusInternalServerError
+		if err, ok := err.(*KeepError); ok {
+			code = err.HTTPCode
+		}
+		http.Error(resp, err.Error(), code)
 		return
 	}
 
@@ -548,12 +567,17 @@ func UntrashHandler(resp http.ResponseWriter, req *http.Request) {
 // If the block found does not have the correct MD5 hash, returns
 // DiskHashError.
 //
-func GetBlock(hash string, buf []byte, resp http.ResponseWriter) (int, error) {
+func GetBlock(ctx context.Context, 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() {
-		size, err := vol.Get(hash, buf)
+		size, err := vol.Get(ctx, hash, buf)
+		select {
+		case <-ctx.Done():
+			return 0, ErrClientDisconnect
+		default:
+		}
 		if err != nil {
 			// IsNotExist is an expected error and may be
 			// ignored. All other errors are logged. In
@@ -587,7 +611,7 @@ func GetBlock(hash string, buf []byte, resp http.ResponseWriter) (int, error) {
 
 // PutBlock Stores the BLOCK (identified by the content id HASH) in Keep.
 //
-// PutBlock(block, hash)
+// PutBlock(ctx, block, hash)
 //   Stores the BLOCK (identified by the content id HASH) in Keep.
 //
 //   The MD5 checksum of the block must be identical to the content id HASH.
@@ -612,7 +636,7 @@ func GetBlock(hash string, buf []byte, resp http.ResponseWriter) (int, error) {
 //          all writes failed). The text of the error message should
 //          provide as much detail as possible.
 //
-func PutBlock(block []byte, hash string) (int, error) {
+func PutBlock(ctx context.Context, block []byte, hash string) (int, error) {
 	// Check that BLOCK's checksum matches HASH.
 	blockhash := fmt.Sprintf("%x", md5.Sum(block))
 	if blockhash != hash {
@@ -623,16 +647,21 @@ func PutBlock(block []byte, hash string) (int, error) {
 	// If we already have this data, it's intact on disk, and we
 	// can update its timestamp, return success. If we have
 	// different data with the same hash, return failure.
-	if n, err := CompareAndTouch(hash, block); err == nil || err == CollisionError {
+	if n, err := CompareAndTouch(ctx, hash, block); err == nil || err == CollisionError {
 		return n, err
+	} else if ctx.Err() != nil {
+		return 0, ErrClientDisconnect
 	}
 
 	// Choose a Keep volume to write to.
 	// If this volume fails, try all of the volumes in order.
 	if vol := KeepVM.NextWritable(); vol != nil {
-		if err := vol.Put(hash, block); err == nil {
+		if err := vol.Put(ctx, hash, block); err == nil {
 			return vol.Replication(), nil // success!
 		}
+		if ctx.Err() != nil {
+			return 0, ErrClientDisconnect
+		}
 	}
 
 	writables := KeepVM.AllWritable()
@@ -643,7 +672,10 @@ func PutBlock(block []byte, hash string) (int, error) {
 
 	allFull := true
 	for _, vol := range writables {
-		err := vol.Put(hash, block)
+		err := vol.Put(ctx, hash, block)
+		if ctx.Err() != nil {
+			return 0, ErrClientDisconnect
+		}
 		if err == nil {
 			return vol.Replication(), nil // success!
 		}
@@ -669,10 +701,13 @@ func PutBlock(block []byte, hash string) (int, error) {
 // the relevant block's modification time in order to protect it from
 // premature garbage collection. Otherwise, it returns a non-nil
 // error.
-func CompareAndTouch(hash string, buf []byte) (int, error) {
+func CompareAndTouch(ctx context.Context, hash string, buf []byte) (int, error) {
 	var bestErr error = NotFoundError
 	for _, vol := range KeepVM.AllWritable() {
-		if err := vol.Compare(hash, buf); err == CollisionError {
+		err := vol.Compare(ctx, hash, buf)
+		if ctx.Err() != nil {
+			return 0, ctx.Err()
+		} else if err == CollisionError {
 			// Stop if we have a block with same hash but
 			// different content. (It will be impossible
 			// to tell which one is wanted if we have