diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 5a5ba7fb60..f610915ea1 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1114,9 +1114,20 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces // waiting for an unreachable agent. The timeout scopes // only dialWithLazyValidation, not ensureWorkspaceAgent // or the post-dial binding steps. - dialCtx, dialCancel := context.WithTimeoutCause(ctx, c.server.dialTimeout, errChatDialTimeout) + dialCtx, dialCancelCause := context.WithCancelCause(ctx) + dialTimer := c.server.clock.AfterFunc( + c.server.dialTimeout, + func() { dialCancelCause(errChatDialTimeout) }, + "chatd", + dialTimeoutTimerTag, + ) + dialCancel := func() { + dialTimer.Stop() + dialCancelCause(nil) + } dialResult, err := dialWithLazyValidation( dialCtx, + c.server.clock, agent.ID, chatSnapshot.WorkspaceID.UUID, DialFunc(c.server.agentConnFn), diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 965b6b474e..b8d9766b9b 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -4972,6 +4972,10 @@ func TestGetWorkspaceConn_DialTimeoutDisconnectedRecoveryThreshold(t *testing.T) } clock := quartz.NewMock(t) + timeoutTrap := clock.Trap().AfterFunc("chatd", dialTimeoutTimerTag) + defer timeoutTrap.Close() + delayTrap := clock.Trap().NewTimer("chatd", dialValidationDelayTimerTag) + defer delayTrap.Close() now := clock.Now() disconnectedAgent := database.WorkspaceAgent{ ID: agentID, @@ -5003,7 +5007,10 @@ func TestGetWorkspaceConn_DialTimeoutDisconnectedRecoveryThreshold(t *testing.T) agentInactiveDisconnectTimeout: 30 * time.Second, dialTimeout: 10 * time.Millisecond, } + dialEntered := make(chan struct{}) + var closeDialEntered sync.Once server.agentConnFn = func(ctx context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { + closeDialEntered.Do(func() { close(dialEntered) }) <-ctx.Done() return nil, nil, ctx.Err() } @@ -5019,13 +5026,41 @@ func TestGetWorkspaceConn_DialTimeoutDisconnectedRecoveryThreshold(t *testing.T) defer workspaceCtx.close() ctx := testutil.Context(t, testutil.WaitShort) - gotConn, err := workspaceCtx.getWorkspaceConn(ctx) - require.Nil(t, gotConn) - require.ErrorIs(t, err, tc.wantErr) + type workspaceConnResult struct { + conn workspacesdk.AgentConn + err error + } + resultCh := make(chan workspaceConnResult, 1) + go func() { + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + resultCh <- workspaceConnResult{conn: gotConn, err: err} + }() + + timeoutCall := timeoutTrap.MustWait(ctx) + require.Equal(t, server.dialTimeout, timeoutCall.Duration) + timeoutCall.MustRelease(ctx) + delayCall := delayTrap.MustWait(ctx) + require.Equal(t, workspaceDialValidationDelay, delayCall.Duration) + delayCall.MustRelease(ctx) + select { + case <-dialEntered: + case <-ctx.Done(): + t.Fatal("timed out waiting for dial to start") + } + clock.Advance(server.dialTimeout).MustWait(ctx) + + var result workspaceConnResult + select { + case result = <-resultCh: + case <-ctx.Done(): + t.Fatal("timed out waiting for getWorkspaceConn") + } + require.Nil(t, result.conn) + require.ErrorIs(t, result.err, tc.wantErr) if tc.wantRecovery { - require.ErrorIs(t, err, errChatAgentDisconnected) + require.ErrorIs(t, result.err, errChatAgentDisconnected) } else { - require.NotErrorIs(t, err, errChatAgentDisconnected) + require.NotErrorIs(t, result.err, errChatAgentDisconnected) } workspaceCtx.mu.Lock() diff --git a/coderd/x/chatd/dialvalidation.go b/coderd/x/chatd/dialvalidation.go index 88c035c4c6..37de0cebe4 100644 --- a/coderd/x/chatd/dialvalidation.go +++ b/coderd/x/chatd/dialvalidation.go @@ -8,6 +8,12 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/quartz" +) + +const ( + dialValidationDelayTimerTag = "dial-validation-delay" + dialTimeoutTimerTag = "dial-timeout" ) // DialResult contains the outcome of dialWithLazyValidation. @@ -44,6 +50,7 @@ type dialOut struct { // switches to a different agent or retries the current one once. func dialWithLazyValidation( ctx context.Context, + clock quartz.Clock, agentID uuid.UUID, workspaceID uuid.UUID, dialFn DialFunc, @@ -110,6 +117,9 @@ func dialWithLazyValidation( case result := <-results: drained = true if result.err != nil { + if waitCtx.Err() != nil { + return DialResult{}, waitCtx.Err() + } return DialResult{}, wrapErr(result.err) } return resultForAgent(agentID, result, false), nil @@ -141,7 +151,7 @@ func dialWithLazyValidation( return dialAgent(validatedAgentID, true) } - timer := time.NewTimer(delay) + timer := clock.NewTimer(delay, "chatd", dialValidationDelayTimerTag) defer timer.Stop() select { @@ -150,6 +160,9 @@ func dialWithLazyValidation( if result.err == nil { return resultForAgent(agentID, result, false), nil } + if ctx.Err() != nil { + return DialResult{}, ctx.Err() + } return resolveFastFailure() case <-timer.C: diff --git a/coderd/x/chatd/dialvalidation_internal_test.go b/coderd/x/chatd/dialvalidation_internal_test.go index de8723ee40..cb93129ec1 100644 --- a/coderd/x/chatd/dialvalidation_internal_test.go +++ b/coderd/x/chatd/dialvalidation_internal_test.go @@ -14,11 +14,13 @@ import ( "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" ) func TestDialWithLazyValidation_FastDial(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) agentID := uuid.New() workspaceID := uuid.New() @@ -29,6 +31,7 @@ func TestDialWithLazyValidation_FastDial(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, agentID, workspaceID, func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -60,6 +63,7 @@ func TestDialWithLazyValidation_FastDial(t *testing.T) { func TestDialWithLazyValidation_SlowDialSameAgent(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) agentID := uuid.New() workspaceID := uuid.New() @@ -71,6 +75,7 @@ func TestDialWithLazyValidation_SlowDialSameAgent(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, agentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -111,6 +116,7 @@ func TestDialWithLazyValidation_SlowDialSameAgent(t *testing.T) { func TestDialWithLazyValidation_SlowDialNoCurrentAgent(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) staleAgentID := uuid.New() workspaceID := uuid.New() dialStarted := make(chan struct{}) @@ -122,6 +128,7 @@ func TestDialWithLazyValidation_SlowDialNoCurrentAgent(t *testing.T) { go func() { _, err := dialWithLazyValidation( context.Background(), + clock, staleAgentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -163,6 +170,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { t.Run("LateSuccessReleasesStaleConn", func(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) staleAgentID := uuid.New() currentAgentID := uuid.New() @@ -177,6 +185,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, staleAgentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -225,6 +234,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { t.Run("CanceledFailureDoesNotReleaseStaleConn", func(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) staleAgentID := uuid.New() currentAgentID := uuid.New() @@ -238,6 +248,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, staleAgentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -284,6 +295,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { t.Run("SwitchDoesNotBlock", func(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) staleAgentID := uuid.New() currentAgentID := uuid.New() @@ -310,6 +322,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { go func() { result, err := dialWithLazyValidation( context.Background(), + clock, staleAgentID, workspaceID, func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -375,6 +388,7 @@ func TestDialWithLazyValidation_SlowDialStaleAgent(t *testing.T) { func TestDialWithLazyValidation_FastFailure(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) staleAgentID := uuid.New() currentAgentID := uuid.New() @@ -387,6 +401,7 @@ func TestDialWithLazyValidation_FastFailure(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, staleAgentID, workspaceID, func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -432,6 +447,7 @@ func TestDialWithLazyValidation_FastFailure(t *testing.T) { func TestDialWithLazyValidation_FastFailureSameAgent(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) agentID := uuid.New() workspaceID := uuid.New() @@ -443,6 +459,7 @@ func TestDialWithLazyValidation_FastFailureSameAgent(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, agentID, workspaceID, func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -485,6 +502,7 @@ func TestDialWithLazyValidation_FastFailureSameAgent(t *testing.T) { func TestDialWithLazyValidation_FastFailureSameAgentRetryFails(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) agentID := uuid.New() workspaceID := uuid.New() @@ -493,6 +511,7 @@ func TestDialWithLazyValidation_FastFailureSameAgentRetryFails(t *testing.T) { _, err := dialWithLazyValidation( context.Background(), + clock, agentID, workspaceID, func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -525,6 +544,7 @@ func TestDialWithLazyValidation_FastFailureSameAgentRetryFails(t *testing.T) { func TestDialWithLazyValidation_ValidationError(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctrl := gomock.NewController(t) agentID := uuid.New() workspaceID := uuid.New() @@ -536,6 +556,7 @@ func TestDialWithLazyValidation_ValidationError(t *testing.T) { result, err := dialWithLazyValidation( context.Background(), + clock, agentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -578,6 +599,7 @@ func TestDialWithLazyValidation_ValidationError(t *testing.T) { func TestDialWithLazyValidation_ContextCanceled(t *testing.T) { t.Parallel() + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -588,6 +610,7 @@ func TestDialWithLazyValidation_ContextCanceled(t *testing.T) { _, err := dialWithLazyValidation( ctx, + clock, agentID, workspaceID, func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {