Repository: grafana/dskit (go)

Issue:
Report context cancel cause in `DoBatchWithOptions` instead of `ctx.Err()`.
In https://github.com/grafana/tempo/issues/3957 it was found that analyzing the root cause of `context cancelled` logs can be quite hard to analyze. In the [discussion](https://github.com/grafana/tempo/issues/3957#issuecomment-2286941935) on how to improve the situation, it was found that reporting the cancel cause, introduced by [`WithCancelCause`](https://pkg.go.dev/context#WithCancelCause) in go 1.20, would be more helpful in these situations. 

Note that this won't change the behavior of error reporting when `WithCancel` (without cause) is used.

Pull request description:
Report context cancel cause in DoBatchWithOptions instead of ctx.Err().
**What this PR does**:

Report context cancel cause in `DoBatchWithOptions` instead of `ctx.Err()`.

**Which issue(s) this PR fixes**:

Closes #576

**Checklist**
- [x] Tests updated
- [x] `CHANGELOG.md` updated - the order ox entries should be `[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 023272e..82b8fc5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## Changelog
 
 * [CHANGE] Update the gRPC dependency to v1.66.0 and deprecate the `grpc_server_recv_buffer_pools_enabled` option that is no longer supported by it. #580
+* [CHANGE] `ring.DoBatchWithOptions` (and `ring.DoBatch`) reports the cancelation cause when the context is canceled instead of `context.Canceled`.
 * [CHANGE] Multierror: Implement `Unwrap() []error`. This allows to use `multierror.MultiError` with both `errors.Is()` and `errors.As()`. Previously implemented `Is(error) bool` has been removed. #522
 * [CHANGE] Add a new `grpc_server_recv_buffer_pools_enabled` option that enables recv buffer pools in the gRPC server (assuming `grpc_server_stats_tracking_enabled` is disabled). #510
 * [CHANGE] Add a new `grpc_server_stats_tracking_enabled` option that allows us to disable stats tracking and potentially improve server memory reuse. #507
diff --git a/ring/batch.go b/ring/batch.go
index f982bd6..e107cab 100644
--- a/ring/batch.go
+++ b/ring/batch.go
@@ -131,7 +131,7 @@ func DoBatchWithOptions(ctx context.Context, op Operation, r DoBatchRing, keys [
 		// Get call below takes ~1 microsecond for ~500 instances.
 		// Checking every 10K calls would be every 10ms.
 		if i%10e3 == 0 {
-			if err := ctx.Err(); err != nil {
+			if err := context.Cause(ctx); err != nil {
 				o.Cleanup()
 				return err
 			}
@@ -161,7 +161,7 @@ func DoBatchWithOptions(ctx context.Context, op Operation, r DoBatchRing, keys [
 	}
 
 	// One last check before calling the callbacks: it doesn't make sense if context is canceled.
-	if err := ctx.Err(); err != nil {
+	if err := context.Cause(ctx); err != nil {
 		o.Cleanup()
 		return err
 	}
@@ -196,7 +196,7 @@ func DoBatchWithOptions(ctx context.Context, op Operation, r DoBatchRing, keys [
 	case <-tracker.done:
 		return nil
 	case <-ctx.Done():
-		return ctx.Err()
+		return context.Cause(ctx)
 	}
 }
 


Test patch:
diff --git a/ring/ring_test.go b/ring/ring_test.go
index 7e670a2..909c5d0 100644
--- a/ring/ring_test.go
+++ b/ring/ring_test.go
@@ -3,6 +3,7 @@ package ring
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"math/rand"
@@ -212,41 +213,131 @@ func TestDoBatchWithOptionsContextCancellation(t *testing.T) {
 		numInstances = 100
 		numZones     = 3
 	)
-	keys := make([]uint32, numKeys)
-	generateKeys(rand.New(rand.NewSource(0)), numKeys, keys)
-
-	callback := func(InstanceDesc, []int) error { return nil }
-	desc := &Desc{Ingesters: generateRingInstances(NewRandomTokenGeneratorWithSeed(0), numInstances, numZones, numTokens)}
-	r := newRingForTesting(Config{
-		HeartbeatTimeout:     time.Hour,
-		ZoneAwarenessEnabled: true,
-		SubringCacheDisabled: true,
-		ReplicationFactor:    numZones,
-	}, true)
-	r.setRingStateFromDesc(desc, false, false, false)
-	// Measure how long does it take for a call to succeed.
-	t0 := time.Now()
-	err := DoBatchWithOptions(context.Background(), Write, r, keys, callback, DoBatchOptions{})
-	duration := time.Since(t0)
-	require.NoError(t, err)
-	t.Logf("Call took %s", duration)
+	cancelCause := errors.New("cancel cause")
 
-	// Make a second call cancelling after a hundredth of duration of the first one.
-	// For a 4s first call, this is 40ms: should be enough for this test to not be flaky.
-	ctx, cancel := context.WithTimeout(context.Background(), duration/100)
-	defer cancel()
+	measureDuration := func(r *Ring, keys []uint32) time.Duration {
+		callback := func(InstanceDesc, []int) error { return nil }
+		t0 := time.Now()
+		err := DoBatchWithOptions(context.Background(), Write, r, keys, callback, DoBatchOptions{})
+		duration := time.Since(t0)
+		require.NoError(t, err)
+		t.Logf("Call took %s", duration)
+		return duration
+	}
 
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	err = DoBatchWithOptions(ctx, Write, r, keys, func(_ InstanceDesc, _ []int) error {
+	type callbackFunc = func(InstanceDesc, []int) error
+	never := func(_ InstanceDesc, _ []int) error {
 		t.Errorf("should not be called.")
 		return nil
-	}, DoBatchOptions{Cleanup: wg.Done})
-	require.Error(t, err)
-	require.ErrorIs(t, err, context.DeadlineExceeded)
+	}
+	tests := []struct {
+		name        string
+		setup       func(*Ring, []uint32) (context.Context, callbackFunc)
+		expectedErr error
+	}{
+		{
+			name: "context deadline exceeded",
+			setup: func(r *Ring, keys []uint32) (context.Context, callbackFunc) {
+				duration := measureDuration(r, keys)
+
+				// Make a second call cancelling after a hundredth of duration of the first one.
+				// For a 4s first call, this is 40ms: should be enough for this test to not be flaky.
+				ctx, cancel := context.WithTimeout(context.Background(), duration/100)
+				go func() {
+					<-ctx.Done()
+					cancel()
+				}()
+				return ctx, never
+			},
+			expectedErr: context.DeadlineExceeded,
+		},
+		{
+			name: "context deadline exceeded with cause",
+			setup: func(r *Ring, keys []uint32) (context.Context, callbackFunc) {
+				duration := measureDuration(r, keys)
+
+				// Make a second call cancelling after a hundredth of duration of the first one.
+				// For a 4s first call, this is 40ms: should be enough for this test to not be flaky.
+				ctx, cancel := context.WithTimeoutCause(context.Background(), duration/100, cancelCause)
+				go func() {
+					<-ctx.Done()
+					cancel()
+				}()
+				return ctx, never
+			},
+			expectedErr: cancelCause,
+		},
+		{
+			name: "context initially cancelled without cause",
+			setup: func(_ *Ring, _ []uint32) (context.Context, callbackFunc) {
+				ctx, cancelFunc := context.WithCancel(context.Background())
+				// start batch with cancelled context
+				cancelFunc()
+				return ctx, never
+			},
+			expectedErr: context.Canceled,
+		},
+		{
+			name: "context initially cancelled with cause",
+			setup: func(_ *Ring, _ []uint32) (context.Context, callbackFunc) {
+				ctx, cancelFunc := context.WithCancelCause(context.Background())
+				// start batch with cancelled context
+				cancelFunc(cancelCause)
+				return ctx, never
+			},
+			expectedErr: cancelCause,
+		},
+		{
+			name: "context cancelled during batch processing",
+			setup: func(_ *Ring, _ []uint32) (context.Context, callbackFunc) {
+				ctx, cancel := context.WithCancelCause(context.Background())
+
+				wg := sync.WaitGroup{}
+				wg.Add(numInstances)
 
-	// Wait until cleanup to make sure that callback was never called.
-	wg.Wait()
+				callback := func(_ InstanceDesc, _ []int) error {
+					wg.Done()
+					// let the call to the instance hang until context is cancelled
+					<-ctx.Done()
+					return nil
+				}
+				go func() {
+					// wait until all instances hang, then cancel the context
+					wg.Wait()
+					cancel(cancelCause)
+				}()
+				return ctx, callback
+			},
+			expectedErr: cancelCause,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			keys := make([]uint32, numKeys)
+			generateKeys(rand.New(rand.NewSource(0)), numKeys, keys)
+
+			desc := &Desc{Ingesters: generateRingInstances(NewRandomTokenGeneratorWithSeed(0), numInstances, numZones, numTokens)}
+			r := newRingForTesting(Config{
+				HeartbeatTimeout:     time.Hour,
+				ZoneAwarenessEnabled: true,
+				SubringCacheDisabled: true,
+				ReplicationFactor:    numZones,
+			}, true)
+			r.setRingStateFromDesc(desc, false, false, false)
+
+			ctx, callback := tt.setup(r, keys)
+
+			wg := sync.WaitGroup{}
+			wg.Add(1)
+			err := DoBatchWithOptions(ctx, Write, r, keys, callback, DoBatchOptions{Cleanup: wg.Done})
+			require.Error(t, err)
+			require.ErrorIs(t, err, tt.expectedErr)
+
+			// Wait until cleanup to make sure that callback was never called.
+			wg.Wait()
+		})
+	}
 }
 
 func TestDoBatch_QuorumError(t *testing.T) {