10467: Use context instead of http.CloseNotifier to interrupt buffer waits.
[arvados.git] / services / keepstore / handlers.go
index 95af1b48707c6b189982dc18762cb517769bd117..210e2b4039fa828e5ef8eb57c9eb480ecc844d43 100644 (file)
@@ -9,6 +9,7 @@ package main
 
 import (
        "container/list"
+       "context"
        "crypto/md5"
        "encoding/json"
        "fmt"
@@ -20,6 +21,7 @@ import (
        "regexp"
        "runtime"
        "strconv"
+       "strings"
        "sync"
        "time"
 )
@@ -53,6 +55,9 @@ func MakeRESTRouter() *mux.Router {
        // Replace the current trash queue.
        rest.HandleFunc(`/trash`, TrashHandler).Methods("PUT")
 
+       // Untrash moves blocks from trash back into store
+       rest.HandleFunc(`/untrash/{hash:[0-9a-f]{32}}`, UntrashHandler).Methods("PUT")
+
        // Any request which does not match any of these routes gets
        // 400 Bad Request.
        rest.NotFoundHandler = http.HandlerFunc(BadRequestHandler)
@@ -67,30 +72,86 @@ 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) {
-       if enforcePermissions {
+       ctx := contextForResponse(context.TODO(), resp)
+
+       if theConfig.RequireSignatures {
                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 := getBufferWithContext(ctx, bufs, BlockSize)
        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)
+               http.Error(resp, err.Error(), http.StatusServiceUnavailable)
                return
        }
-       defer bufs.Put(block)
+       defer bufs.Put(buf)
 
-       resp.Header().Set("Content-Length", strconv.Itoa(len(block)))
+       size, err := GetBlock(ctx, mux.Vars(req)["hash"], buf, resp)
+       if err != nil {
+               code := http.StatusInternalServerError
+               if err, ok := err.(*KeepError); ok {
+                       code = err.HTTPCode
+               }
+               http.Error(resp, err.Error(), code)
+               return
+       }
+
+       resp.Header().Set("Content-Length", strconv.Itoa(size))
        resp.Header().Set("Content-Type", "application/octet-stream")
-       resp.Write(block)
+       resp.Write(buf[:size])
+}
+
+// Return a new context that gets cancelled by resp's
+// CloseNotifier. If resp does not implement http.CloseNotifier,
+// return parent.
+func contextForResponse(parent context.Context, resp http.ResponseWriter) context.Context {
+       cn, ok := resp.(http.CloseNotifier)
+       if !ok {
+               return parent
+       }
+       ctx, cancel := context.WithCancel(parent)
+       go func(c <-chan bool) {
+               <-c
+               cancel()
+       }(cn.CloseNotify())
+       return ctx
+}
+
+// Get a buffer from the pool -- but give up and return a non-nil
+// 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)
+       }()
+       select {
+       case buf := <-bufReady:
+               return buf, nil
+       case <-ctx.Done():
+               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.
 func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
+       ctx := contextForResponse(context.TODO(), resp)
+
        hash := mux.Vars(req)["hash"]
 
        // Detect as many error conditions as possible before reading
@@ -112,15 +173,20 @@ func PutBlockHandler(resp http.ResponseWriter, req *http.Request) {
                return
        }
 
-       buf := bufs.Get(int(req.ContentLength))
-       _, err := io.ReadFull(req.Body, buf)
+       buf, err := getBufferWithContext(ctx, 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)
                return
        }
 
-       replication, err := PutBlock(buf, hash)
+       replication, err := PutBlock(ctx, buf, hash)
        bufs.Put(buf)
 
        if err != nil {
@@ -132,9 +198,9 @@ 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)
-       if PermissionSecret != nil && apiToken != "" {
-               expiry := time.Now().Add(blobSignatureTTL)
+       apiToken := GetAPIToken(req)
+       if theConfig.blobSigningKey != nil && apiToken != "" {
+               expiry := time.Now().Add(theConfig.BlobSignatureTTL.Duration())
                returnHash = SignLocator(returnHash, apiToken, expiry)
        }
        resp.Header().Set("X-Keep-Replicas-Stored", strconv.Itoa(replication))
@@ -144,7 +210,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 !IsSystemAuth(GetAPIToken(req)) {
                http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
                return
        }
@@ -276,13 +342,13 @@ 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
        }
 
-       if neverDelete {
+       if !theConfig.EnableDelete {
                http.Error(resp, MethodDisabledError.Error(), MethodDisabledError.HTTPCode)
                return
        }
@@ -295,7 +361,7 @@ func DeleteHandler(resp http.ResponseWriter, req *http.Request) {
                Failed  int `json:"copies_failed"`
        }
        for _, vol := range KeepVM.AllWritable() {
-               if err := vol.Delete(hash); err == nil {
+               if err := vol.Trash(hash); err == nil {
                        result.Deleted++
                } else if os.IsNotExist(err) {
                        continue
@@ -367,7 +433,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 !IsSystemAuth(GetAPIToken(req)) {
                http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
                return
        }
@@ -403,7 +469,7 @@ type TrashRequest struct {
 // TrashHandler processes /trash requests.
 func TrashHandler(resp http.ResponseWriter, req *http.Request) {
        // Reject unauthorized requests.
-       if !IsDataManagerToken(GetApiToken(req)) {
+       if !IsSystemAuth(GetAPIToken(req)) {
                http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
                return
        }
@@ -430,7 +496,53 @@ func TrashHandler(resp http.ResponseWriter, req *http.Request) {
        trashq.ReplaceQueue(tlist)
 }
 
-// ==============================
+// 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 !IsSystemAuth(GetAPIToken(req)) {
+               http.Error(resp, UnauthorizedError.Error(), UnauthorizedError.HTTPCode)
+               return
+       }
+
+       hash := mux.Vars(req)["hash"]
+
+       if len(KeepVM.AllWritable()) == 0 {
+               http.Error(resp, "No writable volumes", http.StatusNotFound)
+               return
+       }
+
+       var untrashedOn, failedOn []string
+       var numNotFound int
+       for _, vol := range KeepVM.AllWritable() {
+               err := vol.Untrash(hash)
+
+               if os.IsNotExist(err) {
+                       numNotFound++
+               } else if err != nil {
+                       log.Printf("Error untrashing %v on volume %v", hash, vol.String())
+                       failedOn = append(failedOn, vol.String())
+               } else {
+                       log.Printf("Untrashed %v on volume %v", hash, vol.String())
+                       untrashedOn = append(untrashedOn, vol.String())
+               }
+       }
+
+       if numNotFound == len(KeepVM.AllWritable()) {
+               http.Error(resp, "Block not found on any of the writable volumes", http.StatusNotFound)
+               return
+       }
+
+       if len(failedOn) == len(KeepVM.AllWritable()) {
+               http.Error(resp, "Failed to untrash on all writable volumes", http.StatusInternalServerError)
+       } else {
+               respBody := "Successfully untrashed on: " + strings.Join(untrashedOn, ",")
+               if len(failedOn) > 0 {
+                       respBody += "; Failed to untrash on: " + strings.Join(failedOn, ",")
+               }
+               resp.Write([]byte(respBody))
+       }
+}
+
 // 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
@@ -441,24 +553,26 @@ func TrashHandler(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(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() {
-               buf, err := vol.Get(hash)
+               size, err := vol.Get(ctx, hash, buf)
+               select {
+               case <-ctx.Done():
+                       return 0, ctx.Err()
+               default:
+               }
                if err != nil {
                        // IsNotExist is an expected error and may be
                        // ignored. All other errors are logged. In
@@ -472,28 +586,27 @@ 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.
 //
-// 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.
@@ -518,7 +631,7 @@ func GetBlock(hash string) ([]byte, 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 {
@@ -536,7 +649,7 @@ func PutBlock(block []byte, hash string) (int, error) {
        // 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(context.TODO(), hash, block); err == nil {
                        return vol.Replication(), nil // success!
                }
        }
@@ -549,7 +662,12 @@ 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)
+               select {
+               case <-ctx.Done():
+                       return 0, ctx.Err()
+               default:
+               }
                if err == nil {
                        return vol.Replication(), nil // success!
                }
@@ -620,10 +738,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]
@@ -652,7 +770,7 @@ func CanDelete(apiToken string) bool {
        }
        // Blocks may be deleted only when Keep has been configured with a
        // data manager.
-       if IsDataManagerToken(apiToken) {
+       if IsSystemAuth(apiToken) {
                return true
        }
        // TODO(twp): look up apiToken with the API server
@@ -661,8 +779,8 @@ func CanDelete(apiToken string) bool {
        return false
 }
 
-// IsDataManagerToken returns true if apiToken represents the data
-// manager's token.
-func IsDataManagerToken(apiToken string) bool {
-       return dataManagerToken != "" && apiToken == dataManagerToken
+// IsSystemAuth returns true if the given token is allowed to perform
+// system level actions like deleting data.
+func IsSystemAuth(token string) bool {
+       return token != "" && token == theConfig.systemAuthToken
 }