Add a way to drop operation

Client code executing Processor.Run might want to abort running operation if the operation was not yet run in a batch. Such situation is possible when there is a high request congestion.

From now on Processor.Run will accept new parameter context.Context. This context could be cancelled by the client effectively dropping the operation if it was still waiting to be run.

Example:

```
ctx := context.WithTimeout(context.Background(), 5 * time.Second)
err := processor.Run(ctx, "key", ...)
// err will be OperationCancelled
```

Author: elgopher

README.md

@@ -60,8 +60,8 @@ processor := batch.StartProcessor(
 )
 
 // And use the processor inside http/grpc handler or technology-agnostic service.
-// ResourceKey can be taken from request parameter.
-err := processor.Run(resourceKey, func(r *YourResource) {
+// ctx is a standard context.Context and resourceKey can be taken from request parameter
+err := processor.Run(ctx, resourceKey, func(r *YourResource) {
     // Here you put the code which will executed sequentially inside batch  
 })
 ```

_example/http/http.go

@@ -4,6 +4,7 @@
 package http
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -13,7 +14,7 @@ import (
 )
 
 type TrainService interface {
-	Book(train string, seatNumber int, person string) error
+	Book(ctx context.Context, train string, seatNumber int, person string) error
 }
 
 func ListenAndServe(trainService TrainService) error {
@@ -41,7 +42,7 @@ func bookHandler(trainService TrainService) func(http.ResponseWriter, *http.Requ
 			return
 		}
 
-		err = trainService.Book(trainKey, seat, person)
+		err = trainService.Book(request.Context(), trainKey, seat, person)
 
 		if errors.Is(err, train.ErrValidation("")) {
 			writer.WriteHeader(http.StatusBadRequest)
@@ -51,7 +52,7 @@ func bookHandler(trainService TrainService) func(http.ResponseWriter, *http.Requ
 
 		if err != nil {
 			writer.WriteHeader(http.StatusInternalServerError)
-			fmt.Println("internal server error", err)
+			fmt.Println("internal server error:", err)
 			return
 		}
 

_example/train/service.go

@@ -3,19 +3,21 @@
 
 package train
 
+import "context"
+
 // BatchProcessor is an optional interface to decouple your code from `batch` package.
 type BatchProcessor interface {
-	Run(key string, operation func(*Train)) error
+	Run(ctx context.Context, key string, operation func(*Train)) error
 }
 
 type Service struct {
 	BatchProcessor BatchProcessor
 }
 
-func (s Service) Book(trainName string, seatNumber int, person string) error {
+func (s Service) Book(ctx context.Context, trainName string, seatNumber int, person string) error {
 	var operationError error
 
-	batchError := s.BatchProcessor.Run(trainName, func(train *Train) {
+	batchError := s.BatchProcessor.Run(ctx, trainName, func(train *Train) {
 		operationError = train.Book(seatNumber, person)
 	})
 

batch.go

@@ -32,7 +32,7 @@ type Options[Resource any] struct {
 	// SaveResource saves resource with given key to a database. Returning an error aborts the batch.
 	// This function is called at the end of each batch. Context passed as a first parameter
 	// has a timeout calculated using batch MaxDuration. You can use this information to abort saving resource
-	// if it takes too long.
+	// if it takes too long (thus aborting the entire batch).
 	//
 	// By default, does nothing.
 	SaveResource func(_ context.Context, key string, _ Resource) error
@@ -139,7 +139,10 @@ func (s Options[Resource]) withDefaults() Options[Resource] {
 // Run ends when the entire batch has ended. It returns error when batch is aborted or processor is stopped.
 // Only LoadResource and SaveResource functions can abort the batch by returning an error. If error was reported
 // for a batch all Run calls assigned to this batch will get this error.
-func (p *Processor[Resource]) Run(key string, _operation func(Resource)) error {
+//
+// Operation which is still waiting to be run can be canceled by cancelling ctx. If operation was executed but batch
+// is pending then Run waits until batch ends. When ctx is cancelled then OperationCancelled error is returned.
+func (p *Processor[Resource]) Run(ctx context.Context, key string, _operation func(Resource)) error {
 	select {
 	case <-p.stopped:
 		return ProcessorStopped
@@ -151,13 +154,18 @@ func (p *Processor[Resource]) Run(key string, _operation func(Resource)) error {
 
 	goRoutineNumber := p.options.GoRoutineNumberForKey(key, p.options.GoRoutines)
 
-	p.workerChannels[goRoutineNumber] <- operation[Resource]{
+	o := operation[Resource]{
 		resourceKey: key,
 		run:         _operation,
 		result:      result,
 	}
 
-	return <-result
+	select {
+	case p.workerChannels[goRoutineNumber] <- o:
+		return <-result
+	case <-ctx.Done():
+		return OperationCancelled
+	}
 }
 
 // Stop ends all running batches. No new operations will be accepted.

batch_bench_test.go

@@ -1,6 +1,7 @@
 package batch_test
 
 import (
+	"context"
 	"strconv"
 	"sync"
 	"testing"
@@ -31,7 +32,7 @@ func BenchmarkProcessor_Run(b *testing.B) {
 				key := strconv.Itoa(i % resourceCount)
 				go func() {
 					// when
-					err := processor.Run(key, operation)
+					err := processor.Run(context.Background(), key, operation)
 					require.NoError(b, err)
 					allOperationsFinished.Done()
 				}()

batch_test.go

@@ -18,14 +18,16 @@ import (
 	"github.com/elgopher/batch"
 )
 
+var noTimeout = context.Background()
+
 func TestProcessor_Run(t *testing.T) {
 	t.Run("should run callback on zero-value resource when LoadResource was not provided", func(t *testing.T) {
 		futureValue := FutureValue[*resource]()
 
 		processor := batch.StartProcessor(batch.Options[*resource]{})
 		defer processor.Stop()
 		// when
-		err := processor.Run("key", func(c *resource) {
+		err := processor.Run(noTimeout, "key", func(c *resource) {
 			futureValue.Set(c)
 		})
 		require.NoError(t, err)
@@ -47,7 +49,7 @@ func TestProcessor_Run(t *testing.T) {
 			})
 		defer processor.Stop()
 		// when
-		err := processor.Run(key, func(r *resource) {
+		err := processor.Run(noTimeout, key, func(r *resource) {
 			futureValue.Set(r)
 		})
 		require.NoError(t, err)
@@ -68,7 +70,7 @@ func TestProcessor_Run(t *testing.T) {
 			})
 		defer processor.Stop()
 		// when
-		err := processor.Run(key, func(r *resource) {
+		err := processor.Run(noTimeout, key, func(r *resource) {
 			r.value = 2
 		})
 		require.NoError(t, err)
@@ -90,7 +92,7 @@ func TestProcessor_Run(t *testing.T) {
 
 		started := time.Now()
 		// when
-		err := processor.Run("", func(empty) {})
+		err := processor.Run(noTimeout, "", func(empty) {})
 		require.NoError(t, err)
 
 		elapsed := time.Now().Sub(started)
@@ -103,7 +105,7 @@ func TestProcessor_Run(t *testing.T) {
 
 		started := time.Now()
 		// when
-		err := processor.Run("", func(empty) {})
+		err := processor.Run(noTimeout, "", func(empty) {})
 		require.NoError(t, err)
 
 		elapsed := time.Now().Sub(started)
@@ -126,12 +128,12 @@ func TestProcessor_Run(t *testing.T) {
 
 		key := ""
 
-		err := processor.Run(key, func(empty) {
+		err := processor.Run(noTimeout, key, func(empty) {
 			time.Sleep(100 * time.Millisecond)
 		})
 		require.NoError(t, err)
 
-		err = processor.Run(key, func(empty) {})
+		err = processor.Run(noTimeout, key, func(empty) {})
 		require.NoError(t, err)
 
 		assert.Equal(t, int32(2), atomic.LoadInt32(&batchCount))
@@ -154,7 +156,7 @@ func TestProcessor_Run(t *testing.T) {
 			)
 			defer processor.Stop()
 			// when
-			err := processor.Run(key, func(*resource) {})
+			err := processor.Run(noTimeout, key, func(*resource) {})
 			// then
 			assert.ErrorIs(t, err, customError)
 			// and
@@ -181,7 +183,7 @@ func TestProcessor_Run(t *testing.T) {
 			)
 			defer processor.Stop()
 			// when
-			err := processor.Run(key, func(*resource) {})
+			err := processor.Run(noTimeout, key, func(*resource) {})
 			// then
 			assert.ErrorIs(t, err, timeoutError)
 			// and
@@ -206,7 +208,7 @@ func TestProcessor_Run(t *testing.T) {
 			)
 			defer processor.Stop()
 			// when
-			err := processor.Run(key, func(empty) {})
+			err := processor.Run(noTimeout, key, func(empty) {})
 			// then
 			assert.ErrorIs(t, err, timeoutError)
 		})
@@ -229,7 +231,7 @@ func TestProcessor_Run(t *testing.T) {
 
 		for i := 0; i < iterations; i++ {
 			go func() {
-				err := processor.Run("key", func(r *resource) {
+				err := processor.Run(noTimeout, "key", func(r *resource) {
 					r.value++ // value is not guarded so data race should be reported by `go test`
 				})
 				require.NoError(t, err)
@@ -240,14 +242,43 @@ func TestProcessor_Run(t *testing.T) {
 
 		group.Wait()
 	})
+
+	t.Run("should cancel operation if operation is still waiting to be run", func(t *testing.T) {
+		processor := batch.StartProcessor(batch.Options[empty]{})
+		defer processor.Stop()
+
+		var slowOperationStarted sync.WaitGroup
+		slowOperationStarted.Add(1)
+
+		var slowOperationStopped sync.WaitGroup
+		slowOperationStopped.Add(1)
+
+		key := "key"
+
+		go processor.Run(noTimeout, key, func(empty) {
+			slowOperationStarted.Done()
+			slowOperationStopped.Wait()
+		})
+
+		slowOperationStarted.Wait()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		// when
+		cancel()
+		err := processor.Run(ctx, key, func(empty) {})
+		// then
+		assert.ErrorIs(t, err, batch.OperationCancelled)
+		// cleanup
+		slowOperationStopped.Done()
+	})
 }
 
 func TestProcessor_Stop(t *testing.T) {
 	t.Run("after Stop no new operation can be run", func(t *testing.T) {
 		processor := batch.StartProcessor(batch.Options[empty]{})
 		processor.Stop()
 
-		err := processor.Run("key", func(empty) {})
+		err := processor.Run(noTimeout, "key", func(empty) {})
 		assert.ErrorIs(t, err, batch.ProcessorStopped)
 	})
 
@@ -269,7 +300,7 @@ func TestProcessor_Stop(t *testing.T) {
 		started := time.Now()
 
 		go func() {
-			err := processor.Run("key", func(empty) {
+			err := processor.Run(noTimeout, "key", func(empty) {
 				operationExecuted.Done()
 			})
 			require.NoError(t, err)
@@ -301,7 +332,7 @@ func TestProcessor_Stop(t *testing.T) {
 			},
 		)
 		go func() {
-			_ = processor.Run("key", func(empty) {
+			_ = processor.Run(noTimeout, "key", func(empty) {
 				operationExecuted.Done()
 			})
 		}()

error.go

@@ -6,3 +6,4 @@ package batch
 import "errors"
 
 var ProcessorStopped = errors.New("run failed: processor is stopped")
+var OperationCancelled = errors.New("run failed: operation was canceled before it was run")