fix(template-manager): handle error gracefully if metadata already exists (#2130)

* fix(template-manager): handle error gracefully if metadata already exists

* fix: handle rate limit

* chore: add comment

* chore: generalize

* chore: generalize

* chore: auto-commit generated changes

* fix: handle error in close as well

* chore: move error handling layer up

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: jakubno

packages/orchestrator/internal/template/build/layer/layer_executor.go

@@ -312,7 +312,16 @@ func (lb *LayerExecutor) PauseAndUpload(
 			},
 		})
 		if err != nil {
-			return fmt.Errorf("error saving UUID to hash mapping: %w", err)
+			// Since the data should be basically identical, this is safe to skip.
+			if !errors.Is(err, storage.ErrObjectRateLimited) {
+				return fmt.Errorf("error saving UUID to hash mapping: %w", err)
+			}
+
+			logger.L().Warn(ctx, "rate limited writing layer metadata to object, skipping",
+				logger.WithBuildID(meta.Template.BuildID),
+				zap.String("hash", hash),
+				zap.Error(err),
+			)
 		}
 
 		userLogger.Debug(ctx, fmt.Sprintf("Saved: %s", meta.Template.BuildID))

packages/shared/pkg/storage/storage.go

@@ -24,6 +24,10 @@ var (
 
 var ErrObjectNotExist = errors.New("object does not exist")
 
+// ErrObjectRateLimited means per-object mutation rate limiting —
+// multiple concurrent writers racing to write the same content-addressed object.
+var ErrObjectRateLimited = errors.New("object access rate limited")
+
 type Provider string
 
 const (

packages/shared/pkg/storage/storage_google.go

@@ -20,6 +20,8 @@ import (
 	"google.golang.org/api/option"
 	"google.golang.org/api/option/internaloption"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
@@ -303,20 +305,43 @@ func (o *gcpObject) ReadAt(ctx context.Context, buff []byte, off int64) (n int,
 	return n, err
 }
 
-func (o *gcpObject) Put(ctx context.Context, data []byte) (e error) {
+func (o *gcpObject) Put(ctx context.Context, data []byte) error {
 	timer := googleWriteTimerFactory.Begin(attribute.String(gcsOperationAttr, gcsOperationAttrWrite))
 
 	w := o.handle.NewWriter(ctx)
-	defer func() {
-		if err := w.Close(); err != nil {
-			e = errors.Join(e, fmt.Errorf("failed to write to %q: %w", o.path, err))
-		}
-	}()
 
 	c, err := io.Copy(w, bytes.NewReader(data))
 	if err != nil && !errors.Is(err, io.EOF) {
+		closeErr := w.Close()
+		if closeErr != nil {
+			logger.L().Warn(ctx, "failed to close GCS writer after copy error",
+				zap.String("object", o.path),
+				zap.NamedError("error_copy", err),
+				zap.Error(closeErr),
+			)
+		}
+
+		timer.Failure(ctx, c)
+
+		// ResourceExhausted from GCS means per-object mutation rate limiting —
+		// multiple concurrent writers racing to write the same content-addressed object.
+		if isResourceExhausted(err) {
+			return ErrObjectRateLimited
+		}
+
+		return fmt.Errorf("failed to write to %q: %w", o.path, err)
+	}
+
+	// For small objects the GCS Writer buffers data in memory during Write()
+	// and performs the actual upload during Close(). ResourceExhausted errors
+	// from per-object mutation rate limiting will surface here.
+	if err := w.Close(); err != nil {
 		timer.Failure(ctx, c)
 
+		if isResourceExhausted(err) {
+			return ErrObjectRateLimited
+		}
+
 		return fmt.Errorf("failed to write to %q: %w", o.path, err)
 	}
 
@@ -470,3 +495,16 @@ func parseServiceAccountBase64(serviceAccount string) (*gcpServiceToken, error)
 
 	return &sa, nil
 }
+
+func isResourceExhausted(err error) bool {
+	type grpcStatusProvider interface {
+		GRPCStatus() *status.Status
+	}
+
+	var se grpcStatusProvider
+	if errors.As(err, &se) {
+		return se.GRPCStatus().Code() == codes.ResourceExhausted
+	}
+
+	return false
+}