Better context timeout error messages

This one's aimed at producing somewhat better ergonomics when it comes
to working with context timeouts throughout the project. A sizable
problem that we have right now is that in case of timeout, all that
comes back is the generic error from `context.DeadlineExceeded`,
"context deadline exceeded". You don't know where it came from exactly
as it might've bubble up through many layers of the stack. If multiple
timeouts were in use simultaneously, either one might've caused the
timeout.

Here, introduce a timeout helper that aims to let us easily trace the
origin of a timeout. Traditionally, use of a timeout looked like this:

    func (p *StandardPilot) JobGetAvailable(ctx context.Context, exec riverdriver.Executor, state ProducerState, params *riverdriver.JobGetAvailableParams) ([]*rivertype.JobRow, error) {
        ctx, cancel := context.WithTimeout(ctx, rivercommon.HotOperationTimeout)
        defer cancel()

        return exec.JobGetAvailable(ctx, params)

Here, we add `WithTimeout`/`WithTimeoutV` which are used like this:

    func (p *StandardPilot) JobGetAvailable(ctx context.Context, exec riverdriver.Executor, state ProducerState, params *riverdriver.JobGetAvailableParams) ([]*rivertype.JobRow, error) {
        return timeoututil.WithTimeoutV(ctx, rivercommon.HotOperationTimeout, "StandardPilot.JobGetAvailable", func(ctx context.Context) ([]*rivertype.JobRow, error) {
            return exec.JobGetAvailable(ctx, params)
        })

The helpers apply a timeout with `context.WithTimeout` the same way the
first block works, but when a return timeout error is detected, they
wrap it in a more contextual form of the error. So previously, you'd get
this:

    context deadline exceeded

With the helpers, you get this, containing the name of the operation and
the timeout that'd been applied:

    StandardPilot.JobGetAvailable timed out after 1ns: context deadline exceeded

Unfortunately it has to be a helper function with inner callback because
even using `context.WithTimeoutCause`, `ctx.Err` still returns only a
naked `context.DeadlineExceeded`. You need to call `context.Cause(ctx)`
to get the error that was set as cause.

That said, everything else works. A `context.DeadlineExceeded` that was
wrapped on the way out is still recognized and maintains its wrapper.
Multiple levels of `WithTimeout` can be nested and the right error
message still gets returned, depending on which level timed out.

Author: brandur

rivershared/riverpilot/standard_pilot.go

@@ -7,6 +7,7 @@ import (
 	"github.com/riverqueue/river/internal/rivercommon"
 	"github.com/riverqueue/river/riverdriver"
 	"github.com/riverqueue/river/rivershared/baseservice"
+	"github.com/riverqueue/river/rivershared/util/timeoututil"
 	"github.com/riverqueue/river/rivertype"
 )
 
@@ -21,10 +22,9 @@ func (p *StandardPilot) JobGetAvailable(ctx context.Context, exec riverdriver.Ex
 		return nil, nil
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, rivercommon.HotOperationTimeout)
-	defer cancel()
-
-	return exec.JobGetAvailable(ctx, params)
+	return timeoututil.WithTimeoutV(ctx, rivercommon.HotOperationTimeout, "StandardPilot.JobGetAvailable", func(ctx context.Context) ([]*rivertype.JobRow, error) {
+		return exec.JobGetAvailable(ctx, params)
+	})
 }
 
 func (p *StandardPilot) JobGetStuck(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobGetStuckParams) ([]*rivertype.JobRow, error) {

rivershared/util/timeoututil/main_test.go

@@ -0,0 +1,11 @@
+package timeoututil
+
+import (
+	"testing"
+
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+)
+
+func TestMain(m *testing.M) {
+	riversharedtest.WrapTestMain(m)
+}

rivershared/util/timeoututil/timeout_util.go

@@ -0,0 +1,45 @@
+package timeoututil
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+)
+
+// WithTimeout runs innerFunc with a timeout.
+//
+// If innerFunc returns context.DeadlineExceeded because this helper's local
+// timeout fired, WithTimeout returns an error that includes operation and wraps
+// context.DeadlineExceeded. This makes timeout errors easier to trace back to
+// the specific River operation that introduced the timeout instead of surfacing
+// only the generic "context deadline exceeded" message.
+func WithTimeout(ctx context.Context, timeout time.Duration, operation string, innerFunc func(ctx context.Context) error) error {
+	_, err := WithTimeoutV(ctx, timeout, operation, func(ctx context.Context) (struct{}, error) {
+		return struct{}{}, innerFunc(ctx)
+	})
+	return err
+}
+
+// WithTimeoutV runs innerFunc with a timeout and returns its value.
+//
+// If innerFunc returns context.DeadlineExceeded because this helper's local
+// timeout fired, WithTimeoutV returns an error that includes operation and
+// wraps context.DeadlineExceeded. This makes timeout errors easier to trace
+// back to the specific River operation that introduced the timeout instead of
+// surfacing only the generic "context deadline exceeded" message.
+func WithTimeoutV[T any](ctx context.Context, timeout time.Duration, operation string, innerFunc func(ctx context.Context) (T, error)) (T, error) {
+	// need a specific, local error that we can recognize in case multiple
+	// levels of these helpers are wrapped within one another
+	timeoutErr := fmt.Errorf("timeoututil.WithTimeout: %w", context.DeadlineExceeded)
+
+	ctx, cancel := context.WithTimeoutCause(ctx, timeout, timeoutErr)
+	defer cancel()
+
+	ret, err := innerFunc(ctx)
+	if err != nil && errors.Is(err, context.DeadlineExceeded) && errors.Is(context.Cause(ctx), timeoutErr) {
+		var zero T
+		return zero, fmt.Errorf("%s timed out after %s: %w", operation, timeout, err)
+	}
+	return ret, err
+}

rivershared/util/timeoututil/timeout_util_test.go

@@ -0,0 +1,176 @@
+package timeoututil
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithTimeout(t *testing.T) {
+	t.Parallel()
+
+	t.Run("NestedTimeoutsReturnLocalCause", func(t *testing.T) {
+		t.Parallel()
+
+		err := WithTimeout(context.Background(), time.Hour, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Outer", func(ctx context.Context) error {
+			return WithTimeout(ctx, time.Nanosecond, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Inner", func(ctx context.Context) error {
+				<-ctx.Done()
+				return ctx.Err()
+			})
+		})
+		require.EqualError(t, err, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Inner timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		err = WithTimeout(context.Background(), time.Nanosecond, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Outer", func(ctx context.Context) error {
+			return WithTimeout(ctx, time.Hour, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Inner", func(ctx context.Context) error {
+				<-ctx.Done()
+				return ctx.Err()
+			})
+		})
+		require.EqualError(t, err, "TestWithTimeout.NestedTimeoutsReturnLocalCause.Outer timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("PreservesLocalTimeoutCause", func(t *testing.T) {
+		t.Parallel()
+
+		err := WithTimeout(context.Background(), time.Nanosecond, "TestWithTimeout.PreservesLocalTimeoutCause", func(ctx context.Context) error {
+			<-ctx.Done()
+			return context.Cause(ctx)
+		})
+		require.EqualError(t, err, "TestWithTimeout.PreservesLocalTimeoutCause timed out after 1ns: timeoututil.WithTimeout: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("PreservesWrappedDeadlineExceededFromLocalTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		innerErr := errors.New("inner error")
+
+		err := WithTimeout(context.Background(), time.Nanosecond, "TestWithTimeout.PreservesWrappedDeadlineExceededFromLocalTimeout", func(ctx context.Context) error {
+			<-ctx.Done()
+			return fmt.Errorf("%w: %w", innerErr, ctx.Err())
+		})
+		require.EqualError(t, err, "TestWithTimeout.PreservesWrappedDeadlineExceededFromLocalTimeout timed out after 1ns: inner error: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.ErrorIs(t, err, innerErr)
+	})
+
+	t.Run("ReturnsInnerError", func(t *testing.T) {
+		t.Parallel()
+
+		innerErr := errors.New("inner error")
+
+		err := WithTimeout(context.Background(), time.Hour, "TestWithTimeout.ReturnsInnerError", func(ctx context.Context) error {
+			return innerErr
+		})
+		require.ErrorIs(t, err, innerErr)
+	})
+
+	t.Run("ReturnsLocalTimeoutCause", func(t *testing.T) {
+		t.Parallel()
+
+		err := WithTimeout(context.Background(), time.Nanosecond, "TestWithTimeout.ReturnsLocalTimeoutCause", func(ctx context.Context) error {
+			<-ctx.Done()
+			return ctx.Err()
+		})
+		require.EqualError(t, err, "TestWithTimeout.ReturnsLocalTimeoutCause timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+}
+
+func TestWithTimeoutV(t *testing.T) {
+	t.Parallel()
+
+	t.Run("NestedTimeoutsReturnLocalCause", func(t *testing.T) {
+		t.Parallel()
+
+		ret, err := WithTimeoutV(context.Background(), time.Hour, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Outer", func(ctx context.Context) (int, error) {
+			return WithTimeoutV(ctx, time.Nanosecond, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Inner", func(ctx context.Context) (int, error) {
+				<-ctx.Done()
+				return 9, ctx.Err()
+			})
+		})
+		require.Zero(t, ret)
+		require.EqualError(t, err, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Inner timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		ret, err = WithTimeoutV(context.Background(), time.Nanosecond, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Outer", func(ctx context.Context) (int, error) {
+			return WithTimeoutV(ctx, time.Hour, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Inner", func(ctx context.Context) (int, error) {
+				<-ctx.Done()
+				return 9, ctx.Err()
+			})
+		})
+		require.Zero(t, ret)
+		require.EqualError(t, err, "TestWithTimeoutV.NestedTimeoutsReturnLocalCause.Outer timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("PreservesLocalTimeoutCause", func(t *testing.T) {
+		t.Parallel()
+
+		ret, err := WithTimeoutV(context.Background(), time.Nanosecond, "TestWithTimeoutV.PreservesLocalTimeoutCause", func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			return 9, context.Cause(ctx)
+		})
+		require.Zero(t, ret)
+		require.EqualError(t, err, "TestWithTimeoutV.PreservesLocalTimeoutCause timed out after 1ns: timeoututil.WithTimeout: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	t.Run("PreservesParentCancellation", func(t *testing.T) {
+		t.Parallel()
+
+		parentErr := errors.New("parent cancelled")
+		parentCtx, cancel := context.WithCancelCause(context.Background())
+		cancel(parentErr)
+
+		ret, err := WithTimeoutV(parentCtx, time.Hour, "TestWithTimeoutV.PreservesParentCancellation", func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			return 0, context.Cause(ctx)
+		})
+		require.Zero(t, ret)
+		require.ErrorIs(t, err, parentErr)
+	})
+
+	t.Run("PreservesWrappedDeadlineExceededFromLocalTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		innerErr := errors.New("inner error")
+
+		ret, err := WithTimeoutV(context.Background(), time.Nanosecond, "TestWithTimeoutV.PreservesWrappedDeadlineExceededFromLocalTimeout", func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			return 9, fmt.Errorf("%w: %w", innerErr, ctx.Err())
+		})
+		require.Zero(t, ret)
+		require.EqualError(t, err, "TestWithTimeoutV.PreservesWrappedDeadlineExceededFromLocalTimeout timed out after 1ns: inner error: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.ErrorIs(t, err, innerErr)
+	})
+
+	t.Run("ReturnsInnerValue", func(t *testing.T) {
+		t.Parallel()
+
+		ret, err := WithTimeoutV(context.Background(), time.Hour, "TestWithTimeoutV.ReturnsInnerValue", func(ctx context.Context) (int, error) {
+			return 7, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, 7, ret)
+	})
+
+	t.Run("WrapsDeadlineExceededFromLocalTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		ret, err := WithTimeoutV(context.Background(), time.Nanosecond, "TestWithTimeoutV.WrapsDeadlineExceededFromLocalTimeout", func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			return 9, ctx.Err()
+		})
+		require.Zero(t, ret)
+		require.EqualError(t, err, "TestWithTimeoutV.WrapsDeadlineExceededFromLocalTimeout timed out after 1ns: context deadline exceeded")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+}