Merge branch '17995-filter-by-comparing-attrs'
[arvados.git] / sdk / go / keepclient / keepclient.go
index d600322e6a8463a7a7a66577683aad6c697fa031..68ac886ddd8ede2809a56a577baedfd67604e7d1 100644 (file)
@@ -8,6 +8,7 @@ package keepclient
 
 import (
        "bytes"
+       "context"
        "crypto/md5"
        "errors"
        "fmt"
@@ -21,8 +22,8 @@ import (
        "sync"
        "time"
 
+       "git.arvados.org/arvados.git/sdk/go/arvados"
        "git.arvados.org/arvados.git/sdk/go/arvadosclient"
-       "git.arvados.org/arvados.git/sdk/go/asyncbuf"
        "git.arvados.org/arvados.git/sdk/go/httpserver"
 )
 
@@ -68,11 +69,11 @@ type ErrNotFound struct {
        multipleResponseError
 }
 
-type InsufficientReplicasError error
+type InsufficientReplicasError struct{ error }
 
-type OversizeBlockError error
+type OversizeBlockError struct{ error }
 
-var ErrOversizeBlock = OversizeBlockError(errors.New("Exceeded maximum block size (" + strconv.Itoa(BLOCKSIZE) + ")"))
+var ErrOversizeBlock = OversizeBlockError{error: errors.New("Exceeded maximum block size (" + strconv.Itoa(BLOCKSIZE) + ")")}
 var MissingArvadosApiHost = errors.New("Missing required environment variable ARVADOS_API_HOST")
 var MissingArvadosApiToken = errors.New("Missing required environment variable ARVADOS_API_TOKEN")
 var InvalidLocatorError = errors.New("Invalid locator")
@@ -83,26 +84,31 @@ var ErrNoSuchKeepServer = errors.New("No keep server matching the given UUID is
 // ErrIncompleteIndex is returned when the Index response does not end with a new empty line
 var ErrIncompleteIndex = errors.New("Got incomplete index")
 
-const XKeepDesiredReplicas = "X-Keep-Desired-Replicas"
-const XKeepReplicasStored = "X-Keep-Replicas-Stored"
+const (
+       XKeepDesiredReplicas         = "X-Keep-Desired-Replicas"
+       XKeepReplicasStored          = "X-Keep-Replicas-Stored"
+       XKeepStorageClasses          = "X-Keep-Storage-Classes"
+       XKeepStorageClassesConfirmed = "X-Keep-Storage-Classes-Confirmed"
+)
 
 type HTTPClient interface {
        Do(*http.Request) (*http.Response, error)
 }
 
-// Information about Arvados and Keep servers.
+// KeepClient holds information about Arvados and Keep servers.
 type KeepClient struct {
-       Arvados            *arvadosclient.ArvadosClient
-       Want_replicas      int
-       localRoots         map[string]string
-       writableLocalRoots map[string]string
-       gatewayRoots       map[string]string
-       lock               sync.RWMutex
-       HTTPClient         HTTPClient
-       Retries            int
-       BlockCache         *BlockCache
-       RequestID          string
-       StorageClasses     []string
+       Arvados               *arvadosclient.ArvadosClient
+       Want_replicas         int
+       localRoots            map[string]string
+       writableLocalRoots    map[string]string
+       gatewayRoots          map[string]string
+       lock                  sync.RWMutex
+       HTTPClient            HTTPClient
+       Retries               int
+       BlockCache            *BlockCache
+       RequestID             string
+       StorageClasses        []string
+       DefaultStorageClasses []string // Set by cluster's exported config
 
        // set to 1 if all writable services are of disk type, otherwise 0
        replicasPerService int
@@ -114,7 +120,23 @@ type KeepClient struct {
        disableDiscovery bool
 }
 
-// MakeKeepClient creates a new KeepClient, calls
+func (kc *KeepClient) loadDefaultClasses() error {
+       scData, err := kc.Arvados.ClusterConfig("StorageClasses")
+       if err != nil {
+               return err
+       }
+       classes := scData.(map[string]interface{})
+       for scName := range classes {
+               scConf, _ := classes[scName].(map[string]interface{})
+               isDefault, ok := scConf["Default"].(bool)
+               if ok && isDefault {
+                       kc.DefaultStorageClasses = append(kc.DefaultStorageClasses, scName)
+               }
+       }
+       return nil
+}
+
+// MakeKeepClient creates a new KeepClient, loads default storage classes, calls
 // DiscoverKeepServices(), and returns when the client is ready to
 // use.
 func MakeKeepClient(arv *arvadosclient.ArvadosClient) (*KeepClient, error) {
@@ -133,11 +155,16 @@ func New(arv *arvadosclient.ArvadosClient) *KeepClient {
                        defaultReplicationLevel = int(v)
                }
        }
-       return &KeepClient{
+       kc := &KeepClient{
                Arvados:       arv,
                Want_replicas: defaultReplicationLevel,
                Retries:       2,
        }
+       err = kc.loadDefaultClasses()
+       if err != nil {
+               DebugPrintf("DEBUG: Unable to load the default storage classes cluster config")
+       }
+       return kc
 }
 
 // PutHR puts a block given the block hash, a reader, and the number of bytes
@@ -149,23 +176,12 @@ func New(arv *arvadosclient.ArvadosClient) *KeepClient {
 // Returns an InsufficientReplicasError if 0 <= replicas <
 // kc.Wants_replicas.
 func (kc *KeepClient) PutHR(hash string, r io.Reader, dataBytes int64) (string, int, error) {
-       // Buffer for reads from 'r'
-       var bufsize int
-       if dataBytes > 0 {
-               if dataBytes > BLOCKSIZE {
-                       return "", 0, ErrOversizeBlock
-               }
-               bufsize = int(dataBytes)
-       } else {
-               bufsize = BLOCKSIZE
-       }
-
-       buf := asyncbuf.NewBuffer(make([]byte, 0, bufsize))
-       go func() {
-               _, err := io.Copy(buf, HashCheckingReader{r, md5.New(), hash})
-               buf.CloseWithError(err)
-       }()
-       return kc.putReplicas(hash, buf.NewReader, dataBytes)
+       resp, err := kc.BlockWrite(context.Background(), arvados.BlockWriteOptions{
+               Hash:     hash,
+               Reader:   r,
+               DataSize: int(dataBytes),
+       })
+       return resp.Locator, resp.Replicas, err
 }
 
 // PutHB writes a block to Keep. The hash of the bytes is given in
@@ -173,16 +189,21 @@ func (kc *KeepClient) PutHR(hash string, r io.Reader, dataBytes int64) (string,
 //
 // Return values are the same as for PutHR.
 func (kc *KeepClient) PutHB(hash string, buf []byte) (string, int, error) {
-       newReader := func() io.Reader { return bytes.NewBuffer(buf) }
-       return kc.putReplicas(hash, newReader, int64(len(buf)))
+       resp, err := kc.BlockWrite(context.Background(), arvados.BlockWriteOptions{
+               Hash: hash,
+               Data: buf,
+       })
+       return resp.Locator, resp.Replicas, err
 }
 
 // PutB writes a block to Keep. It computes the hash itself.
 //
 // Return values are the same as for PutHR.
 func (kc *KeepClient) PutB(buffer []byte) (string, int, error) {
-       hash := fmt.Sprintf("%x", md5.Sum(buffer))
-       return kc.PutHB(hash, buffer)
+       resp, err := kc.BlockWrite(context.Background(), arvados.BlockWriteOptions{
+               Data: buffer,
+       })
+       return resp.Locator, resp.Replicas, err
 }
 
 // PutR writes a block to Keep. It first reads all data from r into a buffer
@@ -227,7 +248,7 @@ func (kc *KeepClient) getOrHead(method string, locator string, header http.Heade
        var retryList []string
 
        for triesRemaining > 0 {
-               triesRemaining -= 1
+               triesRemaining--
                retryList = nil
 
                for _, host := range serversToTry {
@@ -333,7 +354,7 @@ func (kc *KeepClient) LocalLocator(locator string) (string, error) {
        return loc, nil
 }
 
-// Get() retrieves a block, given a locator. Returns a reader, the
+// Get retrieves a block, given a locator. Returns a reader, the
 // expected data length, the URL the block is being fetched from, and
 // an error.
 //
@@ -345,13 +366,13 @@ func (kc *KeepClient) Get(locator string) (io.ReadCloser, int64, string, error)
        return rdr, size, url, err
 }
 
-// ReadAt() retrieves a portion of block from the cache if it's
+// ReadAt retrieves a portion of block from the cache if it's
 // present, otherwise from the network.
 func (kc *KeepClient) ReadAt(locator string, p []byte, off int) (int, error) {
        return kc.cache().ReadAt(kc, locator, p, off)
 }
 
-// Ask() verifies that a block with the given hash is available and
+// Ask verifies that a block with the given hash is available and
 // readable, according to at least one Keep service. Unlike Get, it
 // does not retrieve the data or verify that the data content matches
 // the hash specified by the locator.
@@ -416,7 +437,7 @@ func (kc *KeepClient) GetIndex(keepServiceUUID, prefix string) (io.Reader, error
        return bytes.NewReader(respBody[0 : len(respBody)-1]), nil
 }
 
-// LocalRoots() returns the map of local (i.e., disk and proxy) Keep
+// LocalRoots returns the map of local (i.e., disk and proxy) Keep
 // services: uuid -> baseURI.
 func (kc *KeepClient) LocalRoots() map[string]string {
        kc.discoverServices()
@@ -425,7 +446,7 @@ func (kc *KeepClient) LocalRoots() map[string]string {
        return kc.localRoots
 }
 
-// GatewayRoots() returns the map of Keep remote gateway services:
+// GatewayRoots returns the map of Keep remote gateway services:
 // uuid -> baseURI.
 func (kc *KeepClient) GatewayRoots() map[string]string {
        kc.discoverServices()
@@ -434,7 +455,7 @@ func (kc *KeepClient) GatewayRoots() map[string]string {
        return kc.gatewayRoots
 }
 
-// WritableLocalRoots() returns the map of writable local Keep services:
+// WritableLocalRoots returns the map of writable local Keep services:
 // uuid -> baseURI.
 func (kc *KeepClient) WritableLocalRoots() map[string]string {
        kc.discoverServices()
@@ -501,6 +522,11 @@ func (kc *KeepClient) ClearBlockCache() {
        kc.cache().Clear()
 }
 
+func (kc *KeepClient) SetStorageClasses(sc []string) {
+       // make a copy so the caller can't mess with it.
+       kc.StorageClasses = append([]string{}, sc...)
+}
+
 var (
        // There are four global http.Client objects for the four
        // possible permutations of TLS behavior (verify/skip-verify)