X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a95f899d7ac84f29b3d019aa410d265bb40833e5..b96d5caa0056472fe67b82bd5305448d85c7d0cd:/services/keep-balance/keep_service.go

diff --git a/services/keep-balance/keep_service.go b/services/keep-balance/keep_service.go
index f65355d0d5..17f8418f62 100644
--- a/services/keep-balance/keep_service.go
+++ b/services/keep-balance/keep_service.go
@@ -1,18 +1,24 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/arvados"
 )
 
 // KeepService represents a keepstore server that is being rebalanced.
 type KeepService struct {
 	arvados.KeepService
+	mounts []*KeepMount
 	*ChangeSet
 }
 
@@ -30,19 +36,19 @@ func (srv *KeepService) URLBase() string {
 
 // CommitPulls sends the current list of pull requests to the storage
 // server (even if the list is empty).
-func (srv *KeepService) CommitPulls(c *arvados.Client) error {
-	return srv.put(c, "pull", srv.ChangeSet.Pulls)
+func (srv *KeepService) CommitPulls(ctx context.Context, c *arvados.Client) error {
+	return srv.put(ctx, c, "pull", srv.ChangeSet.Pulls)
 }
 
 // CommitTrash sends the current list of trash requests to the storage
 // server (even if the list is empty).
-func (srv *KeepService) CommitTrash(c *arvados.Client) error {
-	return srv.put(c, "trash", srv.ChangeSet.Trashes)
+func (srv *KeepService) CommitTrash(ctx context.Context, c *arvados.Client) error {
+	return srv.put(ctx, c, "trash", srv.ChangeSet.Trashes)
 }
 
 // Perform a PUT request at path, with data (as JSON) in the request
 // body.
-func (srv *KeepService) put(c *arvados.Client, path string, data interface{}) error {
+func (srv *KeepService) put(ctx context.Context, c *arvados.Client, path string, data interface{}) error {
 	// We'll start a goroutine to do the JSON encoding, so we can
 	// stream it to the http client through a Pipe, rather than
 	// keeping the entire encoded version in memory.
@@ -59,7 +65,7 @@ func (srv *KeepService) put(c *arvados.Client, path string, data interface{}) er
 	}()
 
 	url := srv.URLBase() + "/" + path
-	req, err := http.NewRequest("PUT", url, ioutil.NopCloser(jsonR))
+	req, err := http.NewRequestWithContext(ctx, "PUT", url, ioutil.NopCloser(jsonR))
 	if err != nil {
 		return fmt.Errorf("building request for %s: %v", url, err)
 	}
@@ -74,3 +80,28 @@ func (srv *KeepService) put(c *arvados.Client, path string, data interface{}) er
 
 	return err
 }
+
+func (srv *KeepService) discoverMounts(c *arvados.Client) error {
+	mounts, err := srv.Mounts(c)
+	if err != nil {
+		return fmt.Errorf("%s: error retrieving mounts: %v", srv, err)
+	}
+	srv.mounts = nil
+	for _, m := range mounts {
+		srv.mounts = append(srv.mounts, &KeepMount{
+			KeepMount:   m,
+			KeepService: srv,
+		})
+	}
+	return nil
+}
+
+type KeepMount struct {
+	arvados.KeepMount
+	KeepService *KeepService
+}
+
+// String implements fmt.Stringer.
+func (mnt *KeepMount) String() string {
+	return fmt.Sprintf("%s (%s) on %s", mnt.UUID, mnt.DeviceID, mnt.KeepService)
+}