3 files changed, 13 insertions, 12 deletions

diff --git a/tools/nixery/server/storage/filesystem.go b/tools/nixery/server/storage/filesystem.go
index 8aca20aac2d0..3fb91bc5e134 100644
--- a/tools/nixery/server/storage/filesystem.go
+++ b/tools/nixery/server/storage/filesystem.go
@@ -2,6 +2,7 @@
 package storage
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -34,7 +35,7 @@ func (b *FSBackend) Name() string {
 	return fmt.Sprintf("Filesystem (%s)", b.path)
 }
 
-func (b *FSBackend) Persist(key string, f func(io.Writer) (string, int64, error)) (string, int64, error) {
+func (b *FSBackend) Persist(ctx context.Context, key string, f Persister) (string, int64, error) {
 	full := path.Join(b.path, key)
 	dir := path.Dir(full)
 	err := os.MkdirAll(dir, 0755)
@@ -53,12 +54,12 @@ func (b *FSBackend) Persist(key string, f func(io.Writer) (string, int64, error)
 	return f(file)
 }
 
-func (b *FSBackend) Fetch(key string) (io.ReadCloser, error) {
+func (b *FSBackend) Fetch(ctx context.Context, key string) (io.ReadCloser, error) {
 	full := path.Join(b.path, key)
 	return os.Open(full)
 }
 
-func (b *FSBackend) Move(old, new string) error {
+func (b *FSBackend) Move(ctx context.Context, old, new string) error {
 	newpath := path.Join(b.path, new)
 	err := os.MkdirAll(path.Dir(newpath), 0755)
 	if err != nil {
diff --git a/tools/nixery/server/storage/gcs.go b/tools/nixery/server/storage/gcs.go
index 749c7ba150e5..b9d70ef20488 100644
--- a/tools/nixery/server/storage/gcs.go
+++ b/tools/nixery/server/storage/gcs.go
@@ -66,8 +66,7 @@ func (b *GCSBackend) Name() string {
 	return "Google Cloud Storage (" + b.bucket + ")"
 }
 
-func (b *GCSBackend) Persist(path string, f func(io.Writer) (string, int64, error)) (string, int64, error) {
-	ctx := context.Background()
+func (b *GCSBackend) Persist(ctx context.Context, path string, f Persister) (string, int64, error) {
 	obj := b.handle.Object(path)
 	w := obj.NewWriter(ctx)
 
@@ -80,8 +79,7 @@ func (b *GCSBackend) Persist(path string, f func(io.Writer) (string, int64, erro
 	return hash, size, w.Close()
 }
 
-func (b *GCSBackend) Fetch(path string) (io.ReadCloser, error) {
-	ctx := context.Background()
+func (b *GCSBackend) Fetch(ctx context.Context, path string) (io.ReadCloser, error) {
 	obj := b.handle.Object(path)
 
 	// Probe whether the file exists before trying to fetch it
@@ -98,8 +96,7 @@ func (b *GCSBackend) Fetch(path string) (io.ReadCloser, error) {
 //
 // The Go API for Cloud Storage does not support renaming objects, but
 // the HTTP API does. The code below makes the relevant call manually.
-func (b *GCSBackend) Move(old, new string) error {
-	ctx := context.Background()
+func (b *GCSBackend) Move(ctx context.Context, old, new string) error {
 	creds, err := google.FindDefaultCredentials(ctx, gcsScope)
 	if err != nil {
 		return err
diff --git a/tools/nixery/server/storage/storage.go b/tools/nixery/server/storage/storage.go
index ad10d682e93a..70095cba4334 100644
--- a/tools/nixery/server/storage/storage.go
+++ b/tools/nixery/server/storage/storage.go
@@ -4,10 +4,13 @@
 package storage
 
 import (
+	"context"
 	"io"
 	"net/http"
 )
 
+type Persister = func(io.Writer) (string, int64, error)
+
 type Backend interface {
 	// Name returns the name of the storage backend, for use in
 	// log messages and such.
@@ -19,14 +22,14 @@ type Backend interface {
 	// It needs to return the SHA256 hash of the data written as
 	// well as the total number of bytes, as those are required
 	// for the image manifest.
-	Persist(string, func(io.Writer) (string, int64, error)) (string, int64, error)
+	Persist(context.Context, string, Persister) (string, int64, error)
 
 	// Fetch retrieves data from the storage backend.
-	Fetch(path string) (io.ReadCloser, error)
+	Fetch(ctx context.Context, path string) (io.ReadCloser, error)
 
 	// Move renames a path inside the storage backend. This is
 	// used for staging uploads while calculating their hashes.
-	Move(old, new string) error
+	Move(ctx context.Context, old, new string) error
 
 	// Serve provides a handler function to serve HTTP requests
 	// for layers in the storage backend.