clean up implementation of WriteFromFileSystem

- no need to wait for the cache write
- one cache block write failure does not cancel all writes
- wrap reader so we can use io.CopyN instead of writing our own

Author: djeebus

packages/shared/pkg/storage/providers/storage_cache.go

@@ -8,14 +8,14 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
-	"golang.org/x/sync/errgroup"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
@@ -25,8 +25,9 @@ import (
 var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/shared/pkg/storage/providers")
 
 const (
-	cacheFilePermissions = 0o600
-	cacheDirPermissions  = 0o700
+	cacheFilePermissions      = 0o600
+	cacheDirPermissions       = 0o700
+	maxCacheWriterConcurrency = 10
 )
 
 var cacheRootPath = env.GetEnv("SHARED_CHUNK_CACHE_PATH", "")
@@ -185,17 +186,13 @@ func (c *CachedFileObjectProvider) WriteFromFileSystem(ctx context.Context, path
 	// write the file to the disk and the remote system at the same time.
 	// this opens the file twice, but the API makes it difficult to use a MultiWriter
 
-	var eg errgroup.Group
+	go c.createCacheBlocksFromFile(ctx, path)
 
-	eg.Go(func() error {
-		return c.createCacheBlocksFromFile(ctx, path)
-	})
-
-	eg.Go(func() error {
-		return c.inner.WriteFromFileSystem(ctx, path)
-	})
+	if err := c.inner.WriteFromFileSystem(ctx, path); err != nil {
+		return fmt.Errorf("failed to write to remote storage: %w", err)
+	}
 
-	return eg.Wait()
+	return nil
 }
 
 func (c *CachedFileObjectProvider) Write(ctx context.Context, src []byte) (int, error) {
@@ -332,46 +329,68 @@ func (c *CachedFileObjectProvider) Delete(ctx context.Context) error {
 	return c.inner.Delete(ctx)
 }
 
-func (c *CachedFileObjectProvider) createCacheBlocksFromFile(ctx context.Context, inputPath string) error {
-	var err error
+func (c *CachedFileObjectProvider) createCacheBlocksFromFile(ctx context.Context, inputPath string) {
 	ctx, span := tracer.Start(ctx, "CachedFileObjectProvider.createCacheBlocksFromFile")
 	defer span.End()
 
 	input, err := os.Open(inputPath)
 	if err != nil {
-		return fmt.Errorf("failed to open file: %w", err)
+		zap.L().Error("failed to open input file", zap.String("path", inputPath), zap.Error(err))
+		return
 	}
 	defer cleanup("failed to close file", input)
 
 	stat, err := input.Stat()
 	if err != nil {
-		return fmt.Errorf("failed to get file size: %w", err)
+		zap.L().Error("failed to stat input file", zap.String("path", inputPath), zap.Error(err))
+		return
 	}
 
 	totalSize := stat.Size()
-	errs, ctx := errgroup.WithContext(ctx)
-	errs.SetLimit(10)
+	var wg sync.WaitGroup
+	workers := make(chan struct{}, maxCacheWriterConcurrency)
 	for offset := int64(0); offset < totalSize; offset += c.chunkSize {
-		func(offset, totalSize int64) {
-			errs.Go(func() error {
-				if fname, err := c.writeChunkFromFile(ctx, offset, totalSize, input); err != nil {
-					safelyRemoveFile(fname)
-					return err
-				}
-				return nil
-			})
-		}(offset, totalSize)
+		wg.Add(1)
+		go func(offset int64) {
+			defer wg.Done()
+
+			// limit concurrency
+			workers <- struct{}{}
+			defer func() { <-workers }()
+
+			if err := c.writeChunkFromFile(ctx, offset, input); err != nil {
+				zap.L().Error("failed to write chunk file",
+					zap.String("path", inputPath),
+					zap.Int64("offset", offset),
+					zap.Error(err))
+			}
+		}(offset)
 	}
-	return errs.Wait()
+	wg.Wait()
+}
+
+type offsetReader struct {
+	wrapped io.ReaderAt
+	offset  int64
+}
+
+var _ io.Reader = (*offsetReader)(nil)
+
+func (r *offsetReader) Read(p []byte) (n int, err error) {
+	n, err = r.wrapped.ReadAt(p, r.offset)
+	r.offset += int64(n)
+	return
+}
+
+func newOffsetReader(file *os.File, offset int64) *offsetReader {
+	return &offsetReader{file, offset}
 }
 
 // writeChunkFromFile writes a piece of a local file. It does not need to worry about race conditions, as it will only
 // be called when building templates, and templates cannot be built on multiple machines at the same time.x
-func (c *CachedFileObjectProvider) writeChunkFromFile(ctx context.Context, offset int64, totalSize int64, input *os.File) (string, error) {
-	var err error
-	ctx, span := tracer.Start(ctx, "CachedFileObjectProvider.writeChunkFromFile", trace.WithAttributes(
+func (c *CachedFileObjectProvider) writeChunkFromFile(ctx context.Context, offset int64, input *os.File) error {
+	_, span := tracer.Start(ctx, "write chunk-from-file", trace.WithAttributes(
 		attribute.Int64("offset", offset),
-		attribute.Int64("total_size", totalSize),
 	))
 	defer span.End()
 
@@ -380,40 +399,17 @@ func (c *CachedFileObjectProvider) writeChunkFromFile(ctx context.Context, offse
 
 	output, err := os.OpenFile(chunkPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, cacheFilePermissions)
 	if err != nil {
-		return chunkPath, fmt.Errorf("failed to open file %s: %w", chunkPath, err)
+		return fmt.Errorf("failed to open file %s: %w", chunkPath, err)
 	}
 	defer cleanup("failed to close file", output)
 
-	expectedRead := min(c.chunkSize, totalSize-offset)
-	totalRead := int64(0)
-	buffer := make([]byte, min(32*1024, expectedRead))
-	for totalRead < expectedRead {
-		select {
-		case <-ctx.Done():
-			return chunkPath, ctx.Err()
-		default:
-		}
-
-		read, err := input.ReadAt(buffer, offset+totalRead)
-		if ignoreEOF(err) != nil {
-			return chunkPath, fmt.Errorf("failed to write to %q [%d bytes @ %d]: %w",
-				chunkPath, c.chunkSize, offset, err)
-		} else if read == 0 {
-			return chunkPath, fmt.Errorf("empty read at %d+%d", offset, totalRead)
-		}
-
-		if _, err = output.Write(buffer[:read]); err != nil {
-			return chunkPath, fmt.Errorf("failed to write to %q [%d bytes @ %d]: %w",
-				chunkPath, c.chunkSize, offset, err)
-		}
-		totalRead += int64(read)
-
-		if errors.Is(err, io.EOF) {
-			return chunkPath, nil
-		}
+	offsetReader := newOffsetReader(input, offset)
+	if _, err := io.CopyN(output, offsetReader, c.chunkSize); ignoreEOF(err) != nil {
+		safelyRemoveFile(chunkPath)
+		return fmt.Errorf("failed to copy chunk: %w", err)
 	}
 
-	return chunkPath, nil
+	return err // in case err == io.EOF
 }
 
 func ignoreFileMissingError(err error) error {

packages/shared/pkg/storage/providers/storage_cache_test.go

@@ -164,7 +164,11 @@ func TestCachedFileObjectProvider_ReadAt(t *testing.T) {
 	})
 
 	t.Run("WriteFromFileSystem should write to cache", func(t *testing.T) {
-		fakeData := generateBytes(t, 11*1024*1024) // 11 MB
+		const megabyte = 1024 * 1024
+		const fileSize = 11 * megabyte
+		const chunkSize = 2 * megabyte
+
+		fakeData := generateBytes(t, fileSize)
 
 		fakeStorageObjectProvider := storagemocks.NewMockStorageObjectProvider(t)
 		fakeStorageObjectProvider.
@@ -175,7 +179,7 @@ func TestCachedFileObjectProvider_ReadAt(t *testing.T) {
 		tempDir := t.TempDir()
 		c := CachedFileObjectProvider{
 			path:      tempDir,
-			chunkSize: 2 * 1024 * 1024, // 2 MB
+			chunkSize: chunkSize,
 			inner:     fakeStorageObjectProvider,
 		}
 
@@ -188,22 +192,32 @@ func TestCachedFileObjectProvider_ReadAt(t *testing.T) {
 		err = c.WriteFromFileSystem(t.Context(), inputFile)
 		require.NoError(t, err)
 
+		time.Sleep(time.Millisecond * 20)
+
 		// ensure remote is not called
 		c.inner = nil
 
 		// read bytes 4-6 MB
-		buffer := make([]byte, 2*1024*1024)                     // 2 MB buffer
-		read, err := c.ReadAt(t.Context(), buffer, 4*1024*1024) // read 4-6 MB
+		buffer := make([]byte, chunkSize)
+		read, err := c.ReadAt(t.Context(), buffer, 4*megabyte) // read 4-6 MB
 		require.NoError(t, err)
-		assert.Equal(t, fakeData[4*1024*1024:6*1024*1024], buffer)
+		assert.Equal(t, fakeData[4*megabyte:6*megabyte], buffer)
 		assert.Equal(t, len(buffer), read)
 
 		// read bytes 10-11 MB
-		buffer = make([]byte, 1*1024*1024)                      // 2 MB buffer
-		read, err = c.ReadAt(t.Context(), buffer, 10*1024*1024) // read 4-6 MB
-		require.NoError(t, err)
-		assert.Equal(t, fakeData[10*1024*1024:], buffer)
-		assert.Equal(t, len(buffer), read)
+		buffer = make([]byte, chunkSize)
+		read, err = c.ReadAt(t.Context(), buffer, 10*megabyte) // read 10-11 MB
+		require.ErrorIs(t, err, io.EOF)
+		assert.Equal(t, megabyte, read) // short read
+		assert.Equal(t, fakeData[10*megabyte:], buffer[:read])
+
+		// verify all chunk files are len(file) == chunkSize
+		for offset := int64(0); offset < fileSize; offset += chunkSize {
+			fname := c.makeChunkFilename(offset)
+			info, err := os.Stat(fname)
+			require.NoError(t, err)
+			assert.Equal(t, min(chunkSize, fileSize-offset), info.Size())
+		}
 	})
 
 	t.Run("ReadFrom should read from cache", func(t *testing.T) {