feat: support context cancellation in file store

Signed-off-by: Lucas Rodriguez <lucas.rodriguez9616@gmail.com>
oras-project · Aug 13, 2024 · c8b6ec3 · c8b6ec3
1 parent d9e1d43
commit c8b6ec3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/content/file/file.go b/content/file/file.go
@@ -395,7 +395,7 @@ func (s *Store) Predecessors(ctx context.Context, node ocispec.Descriptor) ([]oc
 }
 
 // Add adds a file into the file store.
-func (s *Store) Add(_ context.Context, name, mediaType, path string) (ocispec.Descriptor, error) {
+func (s *Store) Add(ctx context.Context, name, mediaType, path string) (ocispec.Descriptor, error) {
 	if s.isClosedSet() {
 		return ocispec.Descriptor{}, ErrStoreClosed
 	}
@@ -426,7 +426,7 @@ func (s *Store) Add(_ context.Context, name, mediaType, path string) (ocispec.De
 	// generate descriptor
 	var desc ocispec.Descriptor
 	if fi.IsDir() {
-		desc, err = s.descriptorFromDir(name, mediaType, path)
+		desc, err = s.descriptorFromDir(ctx, name, mediaType, path)
 	} else {
 		desc, err = s.descriptorFromFile(fi, mediaType, path)
 	}
@@ -505,7 +505,7 @@ func (s *Store) pushDir(name, target string, expected ocispec.Descriptor, conten
 }
 
 // descriptorFromDir generates descriptor from the given directory.
-func (s *Store) descriptorFromDir(name, mediaType, dir string) (desc ocispec.Descriptor, err error) {
+func (s *Store) descriptorFromDir(ctx context.Context, name, mediaType, dir string) (desc ocispec.Descriptor, err error) {
 	// make a temp file to store the gzip
 	gz, err := s.tempFile()
 	if err != nil {
@@ -532,7 +532,7 @@ func (s *Store) descriptorFromDir(name, mediaType, dir string) (desc ocispec.Des
 	tw := io.MultiWriter(gzw, tarDigester.Hash())
 	buf := bufPool.Get().(*[]byte)
 	defer bufPool.Put(buf)
-	if err := tarDirectory(dir, name, tw, s.TarReproducible, *buf); err != nil {
+	if err := tarDirectory(ctx, dir, name, tw, s.TarReproducible, *buf); err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("failed to tar %s: %w", dir, err)
 	}
 

diff --git a/content/file/utils.go b/content/file/utils.go
@@ -18,6 +18,7 @@ package file
 import (
 	"archive/tar"
 	"compress/gzip"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -31,7 +32,7 @@ import (
 
 // tarDirectory walks the directory specified by path, and tar those files with a new
 // path prefix.
-func tarDirectory(root, prefix string, w io.Writer, removeTimes bool, buf []byte) (err error) {
+func tarDirectory(ctx context.Context, root, prefix string, w io.Writer, removeTimes bool, buf []byte) (err error) {
 	tw := tar.NewWriter(w)
 	defer func() {
 		closeErr := tw.Close()
@@ -45,6 +46,12 @@ func tarDirectory(root, prefix string, w io.Writer, removeTimes bool, buf []byte
 			return err
 		}
 
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		// Rename path
 		name, err := filepath.Rel(root, path)
 		if err != nil {