diff --git a/driver.go b/driver.go index 574624217..2202e6bb0 100644 --- a/driver.go +++ b/driver.go @@ -455,6 +455,10 @@ func (d *Driver) connect(ctx context.Context) error { if d.metaBalancer.balancer == nil { b, err := balancer.New(ctx, d.config, d.pool, d.discoveryOptions...) if err != nil { + // Don't wrap context errors with stack trace - let higher levels handle them + if xerrors.IsContextError(err) { + return err + } return xerrors.WithStackTrace(err) } d.metaBalancer.balancer = b diff --git a/driver_test.go b/driver_test.go new file mode 100644 index 000000000..e5d207e6d --- /dev/null +++ b/driver_test.go @@ -0,0 +1,37 @@ +package ydb + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestOpenWithExpiredContext(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond) + defer cancel() + + time.Sleep(10 * time.Nanosecond) + + _, err := Open(ctx, "grpc://localhost:2136/local") + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Contains(t, err.Error(), "context deadline exceeded") + require.Contains(t, err.Error(), "Open(driver.go:") + // Verify that the error doesn't have multiple stack traces (no "retry failed on attempt") + require.NotContains(t, err.Error(), "retry failed") +} + +func TestOpenWithCanceledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + _, err := Open(ctx, "grpc://localhost:2136/local") + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + require.Contains(t, err.Error(), "context canceled") + require.Contains(t, err.Error(), "Open(driver.go:") + // Verify that the error doesn't have multiple stack traces (no "retry failed on attempt") + require.NotContains(t, err.Error(), "retry failed") +} diff --git a/internal/balancer/balancer.go b/internal/balancer/balancer.go index 1e3414389..ac81a654e 100644 --- a/internal/balancer/balancer.go +++ b/internal/balancer/balancer.go @@ -319,6 +319,10 @@ func New(ctx context.Context, driverConfig *config.Config, pool *conn.Pool, opts } else { // initialization of balancer state if err := b.clusterDiscovery(ctx); err != nil { + // Don't wrap context errors with stack trace - let higher levels handle them + if xerrors.IsContextError(err) { + return nil, err + } return nil, xerrors.WithStackTrace(err) } // run background discovering diff --git a/retry/retry.go b/retry/retry.go index 48f15dd41..807529eb2 100644 --- a/retry/retry.go +++ b/retry/retry.go @@ -267,6 +267,10 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err return nil, nil //nolint:nilnil }, opts...) if err != nil { + // Don't wrap context errors with stack trace - let higher levels handle them + if xerrors.IsContextError(err) { + return err + } return xerrors.WithStackTrace(err) } @@ -329,6 +333,12 @@ func RetryWithResult[T any](ctx context.Context, //nolint:revive,funlen defer func() { onDone(attempts, finalErr) }() + // Check if context is already done before entering retry loop + select { + case <-ctx.Done(): + return zeroValue, ctx.Err() + default: + } for { i++ attempts++