Reuse context.Context for the entire batch

This is needed to support database transactions. After LoadResource succeeded the context cannot be canceled, because in the database driver there might be a running go-routine dedicated for transaction which will automatically roll back the transaction once context is canceled.

Author: elgopher

batch_test.go

@@ -271,6 +271,48 @@ func TestProcessor_Run(t *testing.T) {
 		// cleanup
 		slowOperationStopped.Done()
 	})
+
+	t.Run("should use same context for LoadResource and SaveResource", func(t *testing.T) {
+		loadResourceContext := FutureValue[context.Context]()
+		saveResourceContext := FutureValue[context.Context]()
+
+		processor := batch.StartProcessor(batch.Options[empty]{
+			LoadResource: func(ctx context.Context, key string) (empty, error) {
+				loadResourceContext.Set(ctx)
+				return empty{}, nil
+			},
+			SaveResource: func(ctx context.Context, key string, r empty) error {
+				saveResourceContext.Set(ctx)
+				return nil
+			},
+			MinDuration: time.Millisecond,
+		})
+		// when
+		_ = processor.Run(context.Background(), "key", func(empty) {})
+		// then
+		assert.Same(t, loadResourceContext.Get(t), saveResourceContext.Get(t))
+	})
+
+	t.Run("should cancel context passed to LoadResource once Run finished", func(t *testing.T) {
+		loadResourceContext := FutureValue[context.Context]()
+
+		processor := batch.StartProcessor(batch.Options[empty]{
+			LoadResource: func(ctx context.Context, key string) (empty, error) {
+				loadResourceContext.Set(ctx)
+				return empty{}, nil
+			},
+			MinDuration: time.Millisecond,
+			MaxDuration: time.Minute,
+		})
+		// when
+		_ = processor.Run(context.Background(), "key", func(empty) {})
+		// then
+		select {
+		case <-loadResourceContext.Get(t).Done():
+		case <-time.After(time.Second):
+			assert.Fail(t, "Timeout waiting for canceling the context")
+		}
+	})
 }
 
 func TestProcessor_Stop(t *testing.T) {

goroutine.go

@@ -21,21 +21,24 @@ type batch[Resource any] struct {
 }
 
 func (b *batch[Resource]) process() {
+	hardDeadlineContext, cancel := context.WithDeadline(context.Background(), b.hardDeadline)
+	defer cancel()
+
 	softDeadlineReached := time.NewTimer(b.softDeadline.Sub(time.Now()))
 	defer softDeadlineReached.Stop()
 
 	for {
 		select {
 		case <-b.stopped:
-			b.end()
+			b.end(hardDeadlineContext)
 			return
 
 		case <-softDeadlineReached.C:
-			b.end()
+			b.end(hardDeadlineContext)
 			return
 
 		case _operation := <-b.incomingOperations:
-			err := b.load()
+			err := b.load(hardDeadlineContext)
 			if err != nil {
 				_operation.result <- err
 				return
@@ -47,36 +50,30 @@ func (b *batch[Resource]) process() {
 	}
 }
 
-func (b *batch[Resource]) end() {
+func (b *batch[Resource]) end(ctx context.Context) {
 	if b.resource == nil {
 		return
 	}
 
-	err := b.save()
+	err := b.save(ctx)
 	for _, result := range b.results {
 		result <- err
 	}
 }
 
-func (b *batch[Resource]) save() error {
-	ctx, cancel := context.WithDeadline(context.Background(), b.hardDeadline)
-	defer cancel()
-
+func (b *batch[Resource]) save(ctx context.Context) error {
 	if err := b.SaveResource(ctx, b.resourceKey, *b.resource); err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func (b *batch[Resource]) load() error {
+func (b *batch[Resource]) load(ctx context.Context) error {
 	if b.alreadyLoaded() {
 		return nil
 	}
 
-	ctx, cancel := context.WithDeadline(context.Background(), b.hardDeadline)
-	defer cancel()
-
 	resource, err := b.LoadResource(ctx, b.resourceKey)
 	if err != nil {
 		return err