10467: Abort S3 and release buffer if caller disconnects during S3 PUT request.
[arvados.git] / services / keepstore / azure_blob_volume.go
index 95315643f3536da84991cb10b921b3f07235f938..542a9ca6900894d65e1af67339d8d83639428362 100644 (file)
@@ -2,25 +2,34 @@ package main
 
 import (
        "bytes"
+       "context"
        "errors"
        "flag"
        "fmt"
        "io"
        "io/ioutil"
        "log"
+       "net/http"
        "os"
+       "regexp"
+       "strconv"
        "strings"
+       "sync"
        "time"
 
+       "git.curoverse.com/arvados.git/sdk/go/arvados"
        "github.com/curoverse/azure-sdk-for-go/storage"
 )
 
+const azureDefaultRequestTimeout = arvados.Duration(10 * time.Minute)
+
 var (
+       azureMaxGetBytes           int
        azureStorageAccountName    string
        azureStorageAccountKeyFile string
        azureStorageReplication    int
-       azureWriteRaceInterval     time.Duration = 15 * time.Second
-       azureWriteRacePollTime     time.Duration = time.Second
+       azureWriteRaceInterval     = 15 * time.Second
+       azureWriteRacePollTime     = time.Second
 )
 
 func readKeyFromFile(file string) (string, error) {
@@ -36,37 +45,29 @@ func readKeyFromFile(file string) (string, error) {
 }
 
 type azureVolumeAdder struct {
-       *volumeSet
+       *Config
+}
+
+// String implements flag.Value
+func (s *azureVolumeAdder) String() string {
+       return "-"
 }
 
 func (s *azureVolumeAdder) Set(containerName string) error {
-       if containerName == "" {
-               return errors.New("no container name given")
-       }
-       if azureStorageAccountName == "" || azureStorageAccountKeyFile == "" {
-               return errors.New("-azure-storage-account-name and -azure-storage-account-key-file arguments must given before -azure-storage-container-volume")
-       }
-       accountKey, err := readKeyFromFile(azureStorageAccountKeyFile)
-       if err != nil {
-               return err
-       }
-       azClient, err := storage.NewBasicClient(azureStorageAccountName, accountKey)
-       if err != nil {
-               return errors.New("creating Azure storage client: " + err.Error())
-       }
-       if flagSerializeIO {
-               log.Print("Notice: -serialize is not supported by azure-blob-container volumes.")
-       }
-       v := NewAzureBlobVolume(azClient, containerName, flagReadonly, azureStorageReplication)
-       if err := v.Check(); err != nil {
-               return err
-       }
-       *s.volumeSet = append(*s.volumeSet, v)
+       s.Config.Volumes = append(s.Config.Volumes, &AzureBlobVolume{
+               ContainerName:         containerName,
+               StorageAccountName:    azureStorageAccountName,
+               StorageAccountKeyFile: azureStorageAccountKeyFile,
+               AzureReplication:      azureStorageReplication,
+               ReadOnly:              deprecated.flagReadonly,
+       })
        return nil
 }
 
 func init() {
-       flag.Var(&azureVolumeAdder{&volumes},
+       VolumeTypes = append(VolumeTypes, func() VolumeWithExamples { return &AzureBlobVolume{} })
+
+       flag.Var(&azureVolumeAdder{theConfig},
                "azure-storage-container-volume",
                "Use the given container as a storage volume. Can be given multiple times.")
        flag.StringVar(
@@ -78,51 +79,116 @@ func init() {
                &azureStorageAccountKeyFile,
                "azure-storage-account-key-file",
                "",
-               "File containing the account key used for subsequent --azure-storage-container-volume arguments.")
+               "`File` containing the account key used for subsequent --azure-storage-container-volume arguments.")
        flag.IntVar(
                &azureStorageReplication,
                "azure-storage-replication",
                3,
                "Replication level to report to clients when data is stored in an Azure container.")
+       flag.IntVar(
+               &azureMaxGetBytes,
+               "azure-max-get-bytes",
+               BlockSize,
+               fmt.Sprintf("Maximum bytes to request in a single GET request. If smaller than %d, use multiple concurrent range requests to retrieve a block.", BlockSize))
 }
 
 // An AzureBlobVolume stores and retrieves blocks in an Azure Blob
 // container.
 type AzureBlobVolume struct {
-       azClient      storage.Client
-       bsClient      storage.BlobStorageClient
-       containerName string
-       readonly      bool
-       replication   int
+       StorageAccountName    string
+       StorageAccountKeyFile string
+       ContainerName         string
+       AzureReplication      int
+       ReadOnly              bool
+       RequestTimeout        arvados.Duration
+
+       azClient storage.Client
+       bsClient storage.BlobStorageClient
 }
 
-func NewAzureBlobVolume(client storage.Client, containerName string, readonly bool, replication int) *AzureBlobVolume {
-       return &AzureBlobVolume{
-               azClient:      client,
-               bsClient:      client.GetBlobService(),
-               containerName: containerName,
-               readonly:      readonly,
-               replication:   replication,
+// Examples implements VolumeWithExamples.
+func (*AzureBlobVolume) Examples() []Volume {
+       return []Volume{
+               &AzureBlobVolume{
+                       StorageAccountName:    "example-account-name",
+                       StorageAccountKeyFile: "/etc/azure_storage_account_key.txt",
+                       ContainerName:         "example-container-name",
+                       AzureReplication:      3,
+                       RequestTimeout:        azureDefaultRequestTimeout,
+               },
        }
 }
 
-// Check returns nil if the volume is usable.
-func (v *AzureBlobVolume) Check() error {
-       ok, err := v.bsClient.ContainerExists(v.containerName)
+// Type implements Volume.
+func (v *AzureBlobVolume) Type() string {
+       return "Azure"
+}
+
+// Start implements Volume.
+func (v *AzureBlobVolume) Start() error {
+       if v.ContainerName == "" {
+               return errors.New("no container name given")
+       }
+       if v.StorageAccountName == "" || v.StorageAccountKeyFile == "" {
+               return errors.New("StorageAccountName and StorageAccountKeyFile must be given")
+       }
+       accountKey, err := readKeyFromFile(v.StorageAccountKeyFile)
+       if err != nil {
+               return err
+       }
+       v.azClient, err = storage.NewBasicClient(v.StorageAccountName, accountKey)
+       if err != nil {
+               return fmt.Errorf("creating Azure storage client: %s", err)
+       }
+
+       if v.RequestTimeout == 0 {
+               v.RequestTimeout = azureDefaultRequestTimeout
+       }
+       v.azClient.HTTPClient = &http.Client{
+               Timeout: time.Duration(v.RequestTimeout),
+       }
+       v.bsClient = v.azClient.GetBlobService()
+
+       ok, err := v.bsClient.ContainerExists(v.ContainerName)
        if err != nil {
                return err
        }
        if !ok {
-               return errors.New("container does not exist")
+               return fmt.Errorf("Azure container %q does not exist", v.ContainerName)
        }
        return nil
 }
 
-func (v *AzureBlobVolume) Get(loc string) ([]byte, error) {
+// Return true if expires_at metadata attribute is found on the block
+func (v *AzureBlobVolume) checkTrashed(loc string) (bool, map[string]string, error) {
+       metadata, err := v.bsClient.GetBlobMetadata(v.ContainerName, loc)
+       if err != nil {
+               return false, metadata, v.translateError(err)
+       }
+       if metadata["expires_at"] != "" {
+               return true, metadata, nil
+       }
+       return false, metadata, nil
+}
+
+// Get reads a Keep block that has been stored as a block blob in the
+// container.
+//
+// If the block is younger than azureWriteRaceInterval and is
+// unexpectedly empty, assume a PutBlob operation is in progress, and
+// wait for it to finish writing.
+func (v *AzureBlobVolume) Get(ctx context.Context, loc string, buf []byte) (int, error) {
+       trashed, _, err := v.checkTrashed(loc)
+       if err != nil {
+               return 0, err
+       }
+       if trashed {
+               return 0, os.ErrNotExist
+       }
        var deadline time.Time
        haveDeadline := false
-       buf, err := v.get(loc)
-       for err == nil && len(buf) == 0 && loc != "d41d8cd98f00b204e9800998ecf8427e" {
+       size, err := v.get(loc, buf)
+       for err == nil && size == 0 && loc != "d41d8cd98f00b204e9800998ecf8427e" {
                // Seeing a brand new empty block probably means we're
                // in a race with CreateBlob, which under the hood
                // (apparently) does "CreateEmpty" and "CommitData"
@@ -142,82 +208,150 @@ func (v *AzureBlobVolume) Get(loc string) ([]byte, error) {
                } else if time.Now().After(deadline) {
                        break
                }
-               bufs.Put(buf)
                time.Sleep(azureWriteRacePollTime)
-               buf, err = v.get(loc)
+               size, err = v.get(loc, buf)
        }
        if haveDeadline {
-               log.Printf("Race ended with len(buf)==%d", len(buf))
+               log.Printf("Race ended with size==%d", size)
        }
-       return buf, err
+       return size, err
 }
 
-func (v *AzureBlobVolume) get(loc string) ([]byte, error) {
-       rdr, err := v.bsClient.GetBlob(v.containerName, loc)
-       if err != nil {
-               if strings.Contains(err.Error(), "404 Not Found") {
-                       // "storage: service returned without a response body (404 Not Found)"
-                       return nil, os.ErrNotExist
+func (v *AzureBlobVolume) get(loc string, buf []byte) (int, error) {
+       expectSize := len(buf)
+       if azureMaxGetBytes < BlockSize {
+               // Unfortunately the handler doesn't tell us how long the blob
+               // is expected to be, so we have to ask Azure.
+               props, err := v.bsClient.GetBlobProperties(v.ContainerName, loc)
+               if err != nil {
+                       return 0, v.translateError(err)
+               }
+               if props.ContentLength > int64(BlockSize) || props.ContentLength < 0 {
+                       return 0, fmt.Errorf("block %s invalid size %d (max %d)", loc, props.ContentLength, BlockSize)
                }
-               return nil, err
+               expectSize = int(props.ContentLength)
        }
-       switch err := err.(type) {
-       case nil:
-       default:
-               log.Printf("ERROR IN Get(): %T %#v", err, err)
-               return nil, err
+
+       if expectSize == 0 {
+               return 0, nil
        }
-       defer rdr.Close()
-       buf := bufs.Get(BlockSize)
-       n, err := io.ReadFull(rdr, buf)
-       switch err {
-       case io.EOF, io.ErrUnexpectedEOF:
-               return buf[:n], nil
-       default:
-               bufs.Put(buf)
-               return nil, err
+
+       // We'll update this actualSize if/when we get the last piece.
+       actualSize := -1
+       pieces := (expectSize + azureMaxGetBytes - 1) / azureMaxGetBytes
+       errors := make([]error, pieces)
+       var wg sync.WaitGroup
+       wg.Add(pieces)
+       for p := 0; p < pieces; p++ {
+               go func(p int) {
+                       defer wg.Done()
+                       startPos := p * azureMaxGetBytes
+                       endPos := startPos + azureMaxGetBytes
+                       if endPos > expectSize {
+                               endPos = expectSize
+                       }
+                       var rdr io.ReadCloser
+                       var err error
+                       if startPos == 0 && endPos == expectSize {
+                               rdr, err = v.bsClient.GetBlob(v.ContainerName, loc)
+                       } else {
+                               rdr, err = v.bsClient.GetBlobRange(v.ContainerName, loc, fmt.Sprintf("%d-%d", startPos, endPos-1), nil)
+                       }
+                       if err != nil {
+                               errors[p] = err
+                               return
+                       }
+                       defer rdr.Close()
+                       n, err := io.ReadFull(rdr, buf[startPos:endPos])
+                       if pieces == 1 && (err == io.ErrUnexpectedEOF || err == io.EOF) {
+                               // If we don't know the actual size,
+                               // and just tried reading 64 MiB, it's
+                               // normal to encounter EOF.
+                       } else if err != nil {
+                               errors[p] = err
+                       }
+                       if p == pieces-1 {
+                               actualSize = startPos + n
+                       }
+               }(p)
        }
+       wg.Wait()
+       for _, err := range errors {
+               if err != nil {
+                       return 0, v.translateError(err)
+               }
+       }
+       return actualSize, nil
 }
 
+// Compare the given data with existing stored data.
 func (v *AzureBlobVolume) Compare(loc string, expect []byte) error {
-       rdr, err := v.bsClient.GetBlob(v.containerName, loc)
+       trashed, _, err := v.checkTrashed(loc)
        if err != nil {
                return err
        }
+       if trashed {
+               return os.ErrNotExist
+       }
+       rdr, err := v.bsClient.GetBlob(v.ContainerName, loc)
+       if err != nil {
+               return v.translateError(err)
+       }
        defer rdr.Close()
        return compareReaderWithBuf(rdr, expect, loc[:32])
 }
 
-func (v *AzureBlobVolume) Put(loc string, block []byte) error {
-       if v.readonly {
+// Put stores a Keep block as a block blob in the container.
+func (v *AzureBlobVolume) Put(ctx context.Context, loc string, block []byte) error {
+       if v.ReadOnly {
                return MethodDisabledError
        }
-       return v.bsClient.CreateBlockBlobFromReader(v.containerName, loc, uint64(len(block)), bytes.NewReader(block))
+       return v.bsClient.CreateBlockBlobFromReader(v.ContainerName, loc, uint64(len(block)), bytes.NewReader(block), nil)
 }
 
+// Touch updates the last-modified property of a block blob.
 func (v *AzureBlobVolume) Touch(loc string) error {
-       if v.readonly {
+       if v.ReadOnly {
                return MethodDisabledError
        }
-       return v.bsClient.SetBlobMetadata(v.containerName, loc, map[string]string{
-               "touch": fmt.Sprintf("%d", time.Now()),
-       })
+       trashed, metadata, err := v.checkTrashed(loc)
+       if err != nil {
+               return err
+       }
+       if trashed {
+               return os.ErrNotExist
+       }
+
+       metadata["touch"] = fmt.Sprintf("%d", time.Now())
+       return v.bsClient.SetBlobMetadata(v.ContainerName, loc, metadata, nil)
 }
 
+// Mtime returns the last-modified property of a block blob.
 func (v *AzureBlobVolume) Mtime(loc string) (time.Time, error) {
-       props, err := v.bsClient.GetBlobProperties(v.containerName, loc)
+       trashed, _, err := v.checkTrashed(loc)
+       if err != nil {
+               return time.Time{}, err
+       }
+       if trashed {
+               return time.Time{}, os.ErrNotExist
+       }
+
+       props, err := v.bsClient.GetBlobProperties(v.ContainerName, loc)
        if err != nil {
                return time.Time{}, err
        }
        return time.Parse(time.RFC1123, props.LastModified)
 }
 
+// IndexTo writes a list of Keep blocks that are stored in the
+// container.
 func (v *AzureBlobVolume) IndexTo(prefix string, writer io.Writer) error {
        params := storage.ListBlobsParameters{
-               Prefix: prefix,
+               Prefix:  prefix,
+               Include: "metadata",
        }
        for {
-               resp, err := v.bsClient.ListBlobs(v.containerName, params)
+               resp, err := v.bsClient.ListBlobs(v.ContainerName, params)
                if err != nil {
                        return err
                }
@@ -226,6 +360,9 @@ func (v *AzureBlobVolume) IndexTo(prefix string, writer io.Writer) error {
                        if err != nil {
                                return err
                        }
+                       if !v.isKeepBlock(b.Name) {
+                               continue
+                       }
                        if b.Properties.ContentLength == 0 && t.Add(azureWriteRaceInterval).After(time.Now()) {
                                // A new zero-length blob is probably
                                // just a new non-empty blob that
@@ -234,7 +371,11 @@ func (v *AzureBlobVolume) IndexTo(prefix string, writer io.Writer) error {
                                // value.
                                continue
                        }
-                       fmt.Fprintf(writer, "%s+%d %d\n", b.Name, b.Properties.ContentLength, t.Unix())
+                       if b.Metadata["expires_at"] != "" {
+                               // Trashed blob; exclude it from response
+                               continue
+                       }
+                       fmt.Fprintf(writer, "%s+%d %d\n", b.Name, b.Properties.ContentLength, t.UnixNano())
                }
                if resp.NextMarker == "" {
                        return nil
@@ -243,29 +384,61 @@ func (v *AzureBlobVolume) IndexTo(prefix string, writer io.Writer) error {
        }
 }
 
-func (v *AzureBlobVolume) Delete(loc string) error {
-       if v.readonly {
+// Trash a Keep block.
+func (v *AzureBlobVolume) Trash(loc string) error {
+       if v.ReadOnly {
                return MethodDisabledError
        }
+
        // Ideally we would use If-Unmodified-Since, but that
        // particular condition seems to be ignored by Azure. Instead,
        // we get the Etag before checking Mtime, and use If-Match to
        // ensure we don't delete data if Put() or Touch() happens
        // between our calls to Mtime() and DeleteBlob().
-       props, err := v.bsClient.GetBlobProperties(v.containerName, loc)
+       props, err := v.bsClient.GetBlobProperties(v.ContainerName, loc)
        if err != nil {
                return err
        }
        if t, err := v.Mtime(loc); err != nil {
                return err
-       } else if time.Since(t) < blobSignatureTTL {
+       } else if time.Since(t) < theConfig.BlobSignatureTTL.Duration() {
                return nil
        }
-       return v.bsClient.DeleteBlob(v.containerName, loc, map[string]string{
+
+       // If TrashLifetime == 0, just delete it
+       if theConfig.TrashLifetime == 0 {
+               return v.bsClient.DeleteBlob(v.ContainerName, loc, map[string]string{
+                       "If-Match": props.Etag,
+               })
+       }
+
+       // Otherwise, mark as trash
+       return v.bsClient.SetBlobMetadata(v.ContainerName, loc, map[string]string{
+               "expires_at": fmt.Sprintf("%d", time.Now().Add(theConfig.TrashLifetime.Duration()).Unix()),
+       }, map[string]string{
                "If-Match": props.Etag,
        })
 }
 
+// Untrash a Keep block.
+// Delete the expires_at metadata attribute
+func (v *AzureBlobVolume) Untrash(loc string) error {
+       // if expires_at does not exist, return NotFoundError
+       metadata, err := v.bsClient.GetBlobMetadata(v.ContainerName, loc)
+       if err != nil {
+               return v.translateError(err)
+       }
+       if metadata["expires_at"] == "" {
+               return os.ErrNotExist
+       }
+
+       // reset expires_at metadata attribute
+       metadata["expires_at"] = ""
+       err = v.bsClient.SetBlobMetadata(v.ContainerName, loc, metadata, nil)
+       return v.translateError(err)
+}
+
+// Status returns a VolumeStatus struct with placeholder data.
 func (v *AzureBlobVolume) Status() *VolumeStatus {
        return &VolumeStatus{
                DeviceNum: 1,
@@ -274,14 +447,90 @@ func (v *AzureBlobVolume) Status() *VolumeStatus {
        }
 }
 
+// String returns a volume label, including the container name.
 func (v *AzureBlobVolume) String() string {
-       return fmt.Sprintf("azure-storage-container:%+q", v.containerName)
+       return fmt.Sprintf("azure-storage-container:%+q", v.ContainerName)
 }
 
+// Writable returns true, unless the -readonly flag was on when the
+// volume was added.
 func (v *AzureBlobVolume) Writable() bool {
-       return !v.readonly
+       return !v.ReadOnly
 }
 
+// Replication returns the replication level of the container, as
+// specified by the -azure-storage-replication argument.
 func (v *AzureBlobVolume) Replication() int {
-       return v.replication
+       return v.AzureReplication
+}
+
+// If possible, translate an Azure SDK error to a recognizable error
+// like os.ErrNotExist.
+func (v *AzureBlobVolume) translateError(err error) error {
+       switch {
+       case err == nil:
+               return err
+       case strings.Contains(err.Error(), "Not Found"):
+               // "storage: service returned without a response body (404 Not Found)"
+               return os.ErrNotExist
+       default:
+               return err
+       }
+}
+
+var keepBlockRegexp = regexp.MustCompile(`^[0-9a-f]{32}$`)
+
+func (v *AzureBlobVolume) isKeepBlock(s string) bool {
+       return keepBlockRegexp.MatchString(s)
+}
+
+// EmptyTrash looks for trashed blocks that exceeded TrashLifetime
+// and deletes them from the volume.
+func (v *AzureBlobVolume) EmptyTrash() {
+       var bytesDeleted, bytesInTrash int64
+       var blocksDeleted, blocksInTrash int
+       params := storage.ListBlobsParameters{Include: "metadata"}
+
+       for {
+               resp, err := v.bsClient.ListBlobs(v.ContainerName, params)
+               if err != nil {
+                       log.Printf("EmptyTrash: ListBlobs: %v", err)
+                       break
+               }
+               for _, b := range resp.Blobs {
+                       // Check if the block is expired
+                       if b.Metadata["expires_at"] == "" {
+                               continue
+                       }
+
+                       blocksInTrash++
+                       bytesInTrash += b.Properties.ContentLength
+
+                       expiresAt, err := strconv.ParseInt(b.Metadata["expires_at"], 10, 64)
+                       if err != nil {
+                               log.Printf("EmptyTrash: ParseInt(%v): %v", b.Metadata["expires_at"], err)
+                               continue
+                       }
+
+                       if expiresAt > time.Now().Unix() {
+                               continue
+                       }
+
+                       err = v.bsClient.DeleteBlob(v.ContainerName, b.Name, map[string]string{
+                               "If-Match": b.Properties.Etag,
+                       })
+                       if err != nil {
+                               log.Printf("EmptyTrash: DeleteBlob(%v): %v", b.Name, err)
+                               continue
+                       }
+                       blocksDeleted++
+                       bytesDeleted += b.Properties.ContentLength
+               }
+               if resp.NextMarker == "" {
+                       break
+               }
+               params.Marker = resp.NextMarker
+       }
+
+       log.Printf("EmptyTrash stats for %v: Deleted %v bytes in %v blocks. Remaining in trash: %v bytes in %v blocks.", v.String(), bytesDeleted, blocksDeleted, bytesInTrash-bytesDeleted, blocksInTrash-blocksDeleted)
 }