diff --git a/pkg/util/contextutil/context.go b/pkg/util/contextutil/context.go index db7dbb03597c..87a139437153 100644 --- a/pkg/util/contextutil/context.go +++ b/pkg/util/contextutil/context.go @@ -79,6 +79,18 @@ func wrap(ctx context.Context, cancel context.CancelFunc) (context.Context, cont } } +// ctxWithStacktrace overrides Err to annotate context.DeadlineExceeded and +// context.Canceled errors with a stacktrace. +// See: https://github.com/cockroachdb/cockroach/issues/95794 +type ctxWithStacktrace struct { + context.Context +} + +// Err implements the context.Context interface. +func (ctx *ctxWithStacktrace) Err() error { + return errors.WithStack(ctx.Context.Err()) +} + // RunWithTimeout runs a function with a timeout, the same way you'd do with // context.WithTimeout. It improves the opaque error messages returned by // WithTimeout by augmenting them with the op string that is passed in. @@ -86,6 +98,7 @@ func RunWithTimeout( ctx context.Context, op string, timeout time.Duration, fn func(ctx context.Context) error, ) error { ctx, cancel := context.WithTimeout(ctx, timeout) + ctx = &ctxWithStacktrace{Context: ctx} defer cancel() start := timeutil.Now() err := fn(ctx) diff --git a/pkg/util/contextutil/context_test.go b/pkg/util/contextutil/context_test.go index 35b89803d5b6..dc3f5f8e8cb1 100644 --- a/pkg/util/contextutil/context_test.go +++ b/pkg/util/contextutil/context_test.go @@ -12,6 +12,7 @@ package contextutil import ( "context" + "fmt" "net" "testing" "time" @@ -68,6 +69,24 @@ func TestRunWithTimeout(t *testing.T) { } } +func testFuncA(ctx context.Context) error { + return testFuncB(ctx) +} + +func testFuncB(ctx context.Context) error { + <-ctx.Done() + return ctx.Err() +} + +func TestRunWithTimeoutCtxWithStacktrace(t *testing.T) { + ctx := context.Background() + err := RunWithTimeout(ctx, "foo", 1, testFuncA) + require.Error(t, err) + stacktrace := fmt.Sprintf("%+v", err) + require.Contains(t, stacktrace, "testFuncB") + require.Contains(t, stacktrace, "testFuncA") +} + // TestRunWithTimeoutWithoutDeadlineExceeded ensures that when a timeout on the // context occurs but the underlying error does not have // context.DeadlineExceeded as its Cause (perhaps due to serialization) the