diff --git a/pkg/client/action_listener.go b/pkg/client/action_listener.go new file mode 100644 index 0000000000..c418da1485 --- /dev/null +++ b/pkg/client/action_listener.go @@ -0,0 +1,303 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "time" + + grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry" + "github.com/rs/zerolog" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/timestamppb" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" + "github.com/hatchet-dev/hatchet/pkg/validator" +) + +// ListenerStrategy selects the worker action listen RPC variant. +type ListenerStrategy string + +const ( + // ListenerStrategyV1 uses the legacy Listen RPC. + ListenerStrategyV1 ListenerStrategy = "v1" + // ListenerStrategyV2 uses the ListenV2 RPC. + ListenerStrategyV2 ListenerStrategy = "v2" +) + +type actionListenerImpl struct { + client dispatchercontracts.DispatcherClient + listenClient dispatchercontracts.Dispatcher_ListenClient + v validator.Validator + actionStream *reconnectingStream[dispatchercontracts.Dispatcher_ListenClient] + l *zerolog.Logger + ctx *contextLoader + tenantId string + workerId string + + // listenerStrategy is read and written only from the single action-loop + // goroutine — classifyActionError and connectOnce both execute there. + listenerStrategy ListenerStrategy + + actionStreamOnce sync.Once +} + +func (a *actionListenerImpl) Actions(ctx context.Context) (<-chan *Action, <-chan error, error) { + ch := make(chan *Action) + errCh := make(chan error, 1) + + a.l.Debug().Ctx(ctx).Msg("starting action listener") + + go a.heartbeatLoop(ctx) + go a.actionLoop(ctx, ch, errCh) + + return ch, errCh, nil +} + +func (a *actionListenerImpl) heartbeatLoop(ctx context.Context) { + heartbeatInterval := 4 * time.Second + timer := time.NewTicker(100 * time.Millisecond) + defer timer.Stop() + + lastHeartbeat := time.Now().Add(-5 * time.Second) + firstHeartbeat := true + + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + now := time.Now().UTC() + if lastHeartbeat.Add(heartbeatInterval).After(now) { + continue + } + + a.l.Debug().Ctx(ctx).Str("worker_id", a.workerId).Msg("updating worker heartbeat") + + _, err := a.client.Heartbeat(a.ctx.newContext(ctx), &dispatchercontracts.HeartbeatRequest{ + WorkerId: a.workerId, + HeartbeatAt: timestamppb.New(now), + }) + + if err != nil { + a.l.Error().Ctx(ctx).Err(err).Str("worker_id", a.workerId).Msg("could not update worker heartbeat") + + if status.Code(err) == codes.Unimplemented { + return + } + } + + if !firstHeartbeat { + actualInterval := now.Sub(lastHeartbeat) + if actualInterval > heartbeatInterval+1*time.Second { + a.l.Warn().Ctx(ctx). + Str("worker_id", a.workerId). + Dur("actual_interval", actualInterval.Round(time.Millisecond)). + Dur("expected_interval", heartbeatInterval+1*time.Second). + Msg("worker heartbeat interval delay, possible CPU resource contention") + } + } + + firstHeartbeat = false + lastHeartbeat = now + } + } +} + +func (a *actionListenerImpl) actionLoop(ctx context.Context, ch chan<- *Action, errCh chan<- error) { + defer close(ch) + defer close(errCh) + defer func() { + if err := a.actionStreamCore().Close(); err != nil { + a.l.Error().Ctx(ctx).Err(err).Msg("failed to close action listener stream") + } + }() + + classify := a.classifyActionError(newStreamClassifier(func(ctx context.Context) bool { + return ctx.Err() == nil + })) + + err := listenStream(ctx, a.actionStreamCore(), + func(c dispatchercontracts.Dispatcher_ListenClient) (*dispatchercontracts.AssignedAction, error) { + return c.Recv() + }, + func(assigned *dispatchercontracts.AssignedAction) error { + action, ok := a.actionFromAssigned(ctx, assigned) + if !ok { + return nil + } + select { + case ch <- action: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + classify, + ) + if err != nil && ctx.Err() == nil { + sendListenerError(ctx, errCh, err) + } +} + +func (a *actionListenerImpl) classifyActionError(base streamClassifier) streamClassifier { + return func(ctx context.Context, err error) streamVerdict { + if a.listenerStrategy == ListenerStrategyV2 && status.Code(err) == codes.Unimplemented { + a.l.Debug().Ctx(ctx).Msg("falling back to v1 listener strategy") + a.listenerStrategy = ListenerStrategyV1 + return verdictNoProgress + } + return base(ctx, err) + } +} + +func (a *actionListenerImpl) actionStreamCore() *reconnectingStream[dispatchercontracts.Dispatcher_ListenClient] { + a.actionStreamOnce.Do(func() { + wl := a.l.With().Str("worker_id", a.workerId).Logger() + a.actionStream = newReconnectingStream( + &wl, + "action listener", + a.subscribeActionStream, + func(client dispatchercontracts.Dispatcher_ListenClient) error { + return client.CloseSend() + }, + nil, + ) + + if a.listenClient != nil { + a.actionStream.setInitialClient(a.listenClient) + } + }) + + return a.actionStream +} + +func (a *actionListenerImpl) subscribeActionStream(ctx context.Context) (dispatchercontracts.Dispatcher_ListenClient, error) { + if ctx.Err() != nil { + return nil, ctx.Err() + } + + var listenClient dispatchercontracts.Dispatcher_ListenClient + var err error + + switch a.listenerStrategy { + case ListenerStrategyV1: + listenClient, err = a.client.Listen(a.ctx.newContext(ctx), &dispatchercontracts.WorkerListenRequest{ + WorkerId: a.workerId, + }, grpc_retry.Disable()) + case ListenerStrategyV2: + listenClient, err = a.client.ListenV2(a.ctx.newContext(ctx), &dispatchercontracts.WorkerListenRequest{ + WorkerId: a.workerId, + }, grpc_retry.Disable()) + default: + return nil, fmt.Errorf("unknown listener strategy %s", a.listenerStrategy) + } + + if err != nil { + return nil, err + } + + return listenClient, nil +} + +func (a *actionListenerImpl) actionFromAssigned(ctx context.Context, assignedAction *dispatchercontracts.AssignedAction) (*Action, bool) { + var actionType ActionType + + switch assignedAction.ActionType { + case dispatchercontracts.ActionType_START_STEP_RUN: + actionType = ActionTypeStartStepRun + case dispatchercontracts.ActionType_CANCEL_STEP_RUN: + actionType = ActionTypeCancelStepRun + case dispatchercontracts.ActionType_START_GET_GROUP_KEY: + actionType = ActionTypeStartGetGroupKey + default: + a.l.Error().Ctx(ctx).Str("action_type", string(assignedAction.ActionType)).Msg("unknown action type") + return nil, false + } + + a.l.Debug().Ctx(ctx).Str("action_type", string(actionType)).Str("action_id", assignedAction.ActionId).Msg("received action") + + additionalMetadata, ok := a.parseAdditionalMetadata(ctx, assignedAction) + if !ok { + return nil, false + } + + return &Action{ + TenantId: assignedAction.TenantId, + WorkflowRunId: assignedAction.WorkflowRunId, + GetGroupKeyRunId: assignedAction.GetGroupKeyRunId, + WorkerId: a.workerId, + JobId: assignedAction.JobId, + JobName: assignedAction.JobName, + JobRunId: assignedAction.JobRunId, + StepId: assignedAction.TaskId, + StepName: assignedAction.TaskName, + StepRunId: assignedAction.TaskRunExternalId, + ActionId: assignedAction.ActionId, + ActionType: actionType, + ActionPayload: []byte(assignedAction.ActionPayload), + RetryCount: assignedAction.RetryCount, + AdditionalMetadata: additionalMetadata, + ChildIndex: assignedAction.ChildWorkflowIndex, + ChildKey: assignedAction.ChildWorkflowKey, + ParentWorkflowRunId: assignedAction.ParentWorkflowRunId, + Priority: assignedAction.Priority, + WorkflowId: assignedAction.WorkflowId, + WorkflowVersionId: assignedAction.WorkflowVersionId, + TriggeringEventExternalId: assignedAction.TriggeringEventExternalId, + TriggeringEventKey: assignedAction.TriggeringEventKey, + DurableTaskInvocationCount: assignedAction.DurableTaskInvocationCount, + }, true +} + +func (a *actionListenerImpl) parseAdditionalMetadata(ctx context.Context, assignedAction *dispatchercontracts.AssignedAction) (map[string]string, bool) { + if assignedAction.AdditionalMetadata == nil { + return nil, true + } + + var rawMap map[string]interface{} + if err := json.Unmarshal([]byte(*assignedAction.AdditionalMetadata), &rawMap); err != nil { + a.l.Error().Ctx(ctx).Err(err).Msg("could not unmarshal additional metadata") + return nil, false + } + + additionalMetadata := make(map[string]string) + for k, v := range rawMap { + if strVal, ok := v.(string); ok { + additionalMetadata[k] = strVal + } + } + + return additionalMetadata, true +} + +func streamErrorCode(err error) string { + if err == nil { + return "" + } + + if st, ok := status.FromError(err); ok { + return st.Code().String() + } + + return "unknown" +} + +func (a *actionListenerImpl) Unregister() error { + _, err := a.client.Unsubscribe( + a.ctx.newContext(context.Background()), + &dispatchercontracts.WorkerUnsubscribeRequest{ + WorkerId: a.workerId, + }, + ) + + if err != nil { + return err + } + + return nil +} diff --git a/pkg/client/action_listener_test.go b/pkg/client/action_listener_test.go new file mode 100644 index 0000000000..106a2d688d --- /dev/null +++ b/pkg/client/action_listener_test.go @@ -0,0 +1,329 @@ +package client + +import ( + "context" + "fmt" + "io" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" +) + +func TestWorkerActionsSurvivesMoreThanFiveTransientFailures(t *testing.T) { + recvCalls := atomic.Int32{} + subscribeCalls := atomic.Int32{} + + streamClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + if recvCalls.Add(1) <= DefaultActionListenerRetryCount+1 { + return nil, status.Error(codes.Unavailable, "stream broken") + } + + return &dispatchercontracts.AssignedAction{ + ActionType: dispatchercontracts.ActionType_START_STEP_RUN, + ActionId: "action-1", + }, nil + }, + } + + listener := newTestActionListener(t, streamClient) + listener.client = &mockWorkerDispatcherClient{ + listenV2Fn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + require.NotEmpty(t, opts) + subscribeCalls.Add(1) + return streamClient, nil + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + actions, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case action := <-actions: + require.NotNil(t, action) + assert.Equal(t, "action-1", action.ActionId) + case err := <-errCh: + t.Fatalf("unexpected terminal error: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for action after prolonged transient failures") + } + + assert.Greater(t, recvCalls.Load(), int32(DefaultActionListenerRetryCount)) + assert.Greater(t, subscribeCalls.Load(), int32(0)) +} + +func TestWorkerActionsNoProgressConnectFailuresSurfaceOnErrCh(t *testing.T) { + constructorCalls := atomic.Int32{} + + streamClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return nil, fmt.Errorf("plain recv error") + }, + } + + listener := newTestActionListener(t, streamClient) + listener.client = &mockWorkerDispatcherClient{ + listenV2Fn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + constructorCalls.Add(1) + return nil, fmt.Errorf("plain connect error") + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case err := <-errCh: + require.Error(t, err) + assert.Contains(t, err.Error(), "made no progress") + case <-time.After(10 * time.Second): + t.Fatal("expected terminal error on errCh after no-progress cap") + } + + assert.GreaterOrEqual(t, constructorCalls.Load(), int32(maxConsecutiveStreamNoProgress-2)) +} + +func TestWorkerActionsFallsBackFromV2ToV1(t *testing.T) { + listenCalls := atomic.Int32{} + + v1Client := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return &dispatchercontracts.AssignedAction{ + ActionType: dispatchercontracts.ActionType_START_STEP_RUN, + ActionId: "action-v1", + }, nil + }, + } + + listener := newTestActionListener(t, &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return nil, status.Error(codes.Unimplemented, "v2 not supported") + }, + }) + + listener.client = &mockWorkerDispatcherClient{ + listenFn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) { + require.NotEmpty(t, opts) + listenCalls.Add(1) + return v1Client, nil + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + actions, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case action := <-actions: + require.Equal(t, "action-v1", action.ActionId) + assert.Equal(t, ListenerStrategyV1, listener.listenerStrategy) + case err := <-errCh: + t.Fatalf("unexpected error: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for v1 fallback action") + } + + assert.Equal(t, int32(1), listenCalls.Load()) +} + +func TestWorkerActionsV1UnimplementedIsTerminal(t *testing.T) { + listener := newTestActionListener(t, &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return nil, status.Error(codes.Unimplemented, "v2 not supported") + }, + }) + listener.listenerStrategy = ListenerStrategyV1 + + listener.client = &mockWorkerDispatcherClient{ + listenFn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) { + return nil, status.Error(codes.Unimplemented, "v1 not supported") + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case err := <-errCh: + require.Error(t, err) + assert.Equal(t, codes.Unimplemented, status.Code(err)) + case <-time.After(5 * time.Second): + t.Fatal("expected terminal error for v1 unimplemented") + } +} + +func TestWorkerActionsContextCancelStopsLoop(t *testing.T) { + blockRecv := make(chan struct{}) + + listener := newTestActionListener(t, &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + <-blockRecv + return nil, status.Error(codes.Unavailable, "stream broken") + }, + }) + + ctx, cancel := context.WithCancel(context.Background()) + _, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + cancel() + + select { + case err, ok := <-errCh: + if ok { + t.Fatalf("unexpected terminal error: %v", err) + } + case <-time.After(2 * time.Second): + } +} + +func TestWorkerListenDisablesGrpcRetry(t *testing.T) { + subscribeCalls := atomic.Int32{} + streamClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return nil, io.EOF + }, + } + + listener := newTestActionListener(t, streamClient) + listener.client = &mockWorkerDispatcherClient{ + listenV2Fn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + require.NotEmpty(t, opts) + subscribeCalls.Add(1) + return streamClient, nil + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case err, ok := <-errCh: + if ok && err != nil { + t.Fatalf("unexpected terminal error: %v", err) + } + case <-time.After(2 * time.Second): + } + + assert.GreaterOrEqual(t, subscribeCalls.Load(), int32(1)) +} + +func TestWorkerActionsReconnectsOnEOFWhileContextActive(t *testing.T) { + recvCalls := atomic.Int32{} + subscribeCalls := atomic.Int32{} + + replacementClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return &dispatchercontracts.AssignedAction{ + ActionType: dispatchercontracts.ActionType_START_STEP_RUN, + ActionId: "action-after-eof", + }, nil + }, + } + + streamClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + recvCalls.Add(1) + return nil, io.EOF + }, + } + + listener := newTestActionListener(t, streamClient) + listener.client = &mockWorkerDispatcherClient{ + listenV2Fn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + require.NotEmpty(t, opts) + subscribeCalls.Add(1) + return replacementClient, nil + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + actions, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case action := <-actions: + require.NotNil(t, action) + assert.Equal(t, "action-after-eof", action.ActionId) + case err := <-errCh: + t.Fatalf("unexpected terminal error: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for action after EOF reconnect") + } + + assert.Equal(t, int32(1), recvCalls.Load()) + assert.Equal(t, int32(1), subscribeCalls.Load()) +} + +func TestWorkerActionsRetriesNoProgressBeforeGivingUp(t *testing.T) { + recvCalls := atomic.Int32{} + subscribeCalls := atomic.Int32{} + + replacementClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + return &dispatchercontracts.AssignedAction{ + ActionType: dispatchercontracts.ActionType_START_STEP_RUN, + ActionId: "action-after-unknown", + }, nil + }, + } + + streamClient := &mockListenV2Client{ + recvFn: func() (*dispatchercontracts.AssignedAction, error) { + recvCalls.Add(1) + return nil, status.Error(codes.Unknown, "transient unknown") + }, + } + + listener := newTestActionListener(t, streamClient) + listener.client = &mockWorkerDispatcherClient{ + listenV2Fn: func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + require.NotEmpty(t, opts) + subscribeCalls.Add(1) + return replacementClient, nil + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + actions, errCh, err := listener.Actions(ctx) + require.NoError(t, err) + + select { + case action := <-actions: + require.NotNil(t, action) + assert.Equal(t, "action-after-unknown", action.ActionId) + case err := <-errCh: + t.Fatalf("unexpected terminal error: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for action after no-progress retry") + } + + assert.Equal(t, int32(1), recvCalls.Load()) + assert.Equal(t, int32(1), subscribeCalls.Load()) +} diff --git a/pkg/client/dispatcher.go b/pkg/client/dispatcher.go index 623ef5b903..a4df93464a 100644 --- a/pkg/client/dispatcher.go +++ b/pkg/client/dispatcher.go @@ -13,9 +13,7 @@ import ( grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry" "github.com/rs/zerolog" "google.golang.org/grpc" - "google.golang.org/grpc/codes" _ "google.golang.org/grpc/encoding/gzip" // Register gzip compression codec - "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/timestamppb" dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" @@ -209,31 +207,6 @@ func newDispatcher(conn *grpc.ClientConn, opts *sharedClientOpts, presetWorkerLa } } -type ListenerStrategy string - -const ( - ListenerStrategyV1 ListenerStrategy = "v1" - ListenerStrategyV2 ListenerStrategy = "v2" -) - -type actionListenerImpl struct { - client dispatchercontracts.DispatcherClient - - tenantId string - - listenClient dispatchercontracts.Dispatcher_ListenClient - - workerId string - - l *zerolog.Logger - - v validator.Validator - - ctx *contextLoader - - listenerStrategy ListenerStrategy -} - func (d *dispatcherClientImpl) newActionListener(ctx context.Context, req *GetActionListenerRequest) (*actionListenerImpl, *string, error) { // validate the request if err := d.v.Validate(req); err != nil { @@ -308,7 +281,7 @@ func (d *dispatcherClientImpl) newActionListener(ctx context.Context, req *GetAc // subscribe to the worker listener, err := d.client.ListenV2(d.ctx.newContext(ctx), &dispatchercontracts.WorkerListenRequest{ WorkerId: resp.WorkerId, - }) + }, grpc_retry.Disable()) if err != nil { return nil, nil, fmt.Errorf("could not subscribe to the worker: %w", err) @@ -326,237 +299,6 @@ func (d *dispatcherClientImpl) newActionListener(ctx context.Context, req *GetAc }, &resp.WorkerId, nil } -func (a *actionListenerImpl) Actions(ctx context.Context) (<-chan *Action, <-chan error, error) { - ch := make(chan *Action) - errCh := make(chan error) - - a.l.Debug().Ctx(ctx).Msgf("Starting to listen for actions") - - // update the worker with a last heartbeat time every 4 seconds as long as the worker is connected - go func() { - heartbeatInterval := 4 * time.Second - timer := time.NewTicker(100 * time.Millisecond) - defer timer.Stop() - - // set last heartbeat to 5 seconds ago so that the first heartbeat is sent immediately - lastHeartbeat := time.Now().Add(-5 * time.Second) - firstHeartbeat := true - - for { - select { - case <-ctx.Done(): - return - case <-timer.C: - if now := time.Now().UTC(); lastHeartbeat.Add(heartbeatInterval).Before(now) { - a.l.Debug().Ctx(ctx).Msgf("updating worker %s heartbeat", a.workerId) - - _, err := a.client.Heartbeat(a.ctx.newContext(ctx), &dispatchercontracts.HeartbeatRequest{ - WorkerId: a.workerId, - HeartbeatAt: timestamppb.New(now), - }) - - if err != nil { - a.l.Error().Ctx(ctx).Err(err).Msgf("could not update worker %s heartbeat", a.workerId) - - // if the heartbeat method is unimplemented, don't continue to send heartbeats - if status.Code(err) == codes.Unimplemented { - return - } - } - - // detect heartbeat delays caused by CPU contention or other scheduling issues, - // but skip the first heartbeat since lastHeartbeat is artificially backdated - if !firstHeartbeat { - actualInterval := now.Sub(lastHeartbeat) - // add 1 second to the heartbeat interval to account for the time it takes to send the heartbeat - if actualInterval > heartbeatInterval+1*time.Second { - a.l.Warn().Ctx(ctx).Msgf( - "worker %s heartbeat interval delay (%s >> %s), possible CPU resource contention", - a.workerId, actualInterval.Round(time.Millisecond), heartbeatInterval+1*time.Second, - ) - } - } - - firstHeartbeat = false - lastHeartbeat = now - } - } - } - }() - - go func() { - retries := 0 - - for retries < DefaultActionListenerRetryCount { - assignedAction, err := a.listenClient.Recv() - - if err != nil { - // if context is cancelled, unsubscribe and close the channel - if ctx.Err() != nil { - a.l.Debug().Ctx(ctx).Msgf("Context cancelled, closing channel") - - defer close(ch) - defer close(errCh) - - err := a.listenClient.CloseSend() - - if err != nil { - a.l.Error().Ctx(ctx).Msgf("Failed to close send: %v", err) - } - - return - } - - retries++ - - // if this is an unimplemented error, default to v1 - if a.listenerStrategy == ListenerStrategyV2 && status.Code(err) == codes.Unimplemented { - a.l.Debug().Ctx(ctx).Msgf("Falling back to v1 listener strategy") - a.listenerStrategy = ListenerStrategyV1 - } - - err = a.retrySubscribe(ctx) - - if err != nil { - a.l.Error().Ctx(ctx).Msgf("Failed to resubscribe: %v", err) - errCh <- fmt.Errorf("failed to resubscribe: %w", err) - } - - time.Sleep(DefaultActionListenerRetryInterval) - - continue - } - - retries = 0 - - var actionType ActionType - - switch assignedAction.ActionType { - case dispatchercontracts.ActionType_START_STEP_RUN: - actionType = ActionTypeStartStepRun - case dispatchercontracts.ActionType_CANCEL_STEP_RUN: - actionType = ActionTypeCancelStepRun - case dispatchercontracts.ActionType_START_GET_GROUP_KEY: - actionType = ActionTypeStartGetGroupKey - default: - a.l.Error().Ctx(ctx).Msgf("Unknown action type: %s", assignedAction.ActionType) - continue - } - - a.l.Debug().Ctx(ctx).Msgf("Received action type: %s for action: %s", actionType, assignedAction.ActionId) - - unquoted := assignedAction.ActionPayload - - var additionalMetadata map[string]string - - if assignedAction.AdditionalMetadata != nil { - // Try to unmarshal as map[string]string first - var rawMap map[string]interface{} - if err := json.Unmarshal([]byte(*assignedAction.AdditionalMetadata), &rawMap); err != nil { - // If that fails, try to unmarshal as a single string - a.l.Error().Ctx(ctx).Err(err).Msgf("could not unmarshal additional metadata") - continue - } else { - // Only keep string values from the map - additionalMetadata = make(map[string]string) - for k, v := range rawMap { - if strVal, ok := v.(string); ok { - additionalMetadata[k] = strVal - } - } - } - } - - ch <- &Action{ - TenantId: assignedAction.TenantId, - WorkflowRunId: assignedAction.WorkflowRunId, - GetGroupKeyRunId: assignedAction.GetGroupKeyRunId, - WorkerId: a.workerId, - JobId: assignedAction.JobId, - JobName: assignedAction.JobName, - JobRunId: assignedAction.JobRunId, - StepId: assignedAction.TaskId, - StepName: assignedAction.TaskName, - StepRunId: assignedAction.TaskRunExternalId, - ActionId: assignedAction.ActionId, - ActionType: actionType, - ActionPayload: []byte(unquoted), - RetryCount: assignedAction.RetryCount, - AdditionalMetadata: additionalMetadata, - ChildIndex: assignedAction.ChildWorkflowIndex, - ChildKey: assignedAction.ChildWorkflowKey, - ParentWorkflowRunId: assignedAction.ParentWorkflowRunId, - Priority: assignedAction.Priority, - WorkflowId: assignedAction.WorkflowId, - WorkflowVersionId: assignedAction.WorkflowVersionId, - TriggeringEventExternalId: assignedAction.TriggeringEventExternalId, - TriggeringEventKey: assignedAction.TriggeringEventKey, - DurableTaskInvocationCount: assignedAction.DurableTaskInvocationCount, - } - } - - errCh <- fmt.Errorf("could not subscribe to the worker after %d retries", retries) - - defer close(ch) - defer close(errCh) - - err := a.listenClient.CloseSend() - - if err != nil { - a.l.Error().Ctx(ctx).Msgf("Failed to close send: %v", err) - } - }() - - return ch, errCh, nil -} - -func (a *actionListenerImpl) retrySubscribe(ctx context.Context) error { - retries := 0 - - for retries < DefaultActionListenerRetryCount { - var err error - var listenClient dispatchercontracts.Dispatcher_ListenClient - - if a.listenerStrategy == ListenerStrategyV1 { - listenClient, err = a.client.Listen(a.ctx.newContext(ctx), &dispatchercontracts.WorkerListenRequest{ - WorkerId: a.workerId, - }) - } else if a.listenerStrategy == ListenerStrategyV2 { - listenClient, err = a.client.ListenV2(a.ctx.newContext(ctx), &dispatchercontracts.WorkerListenRequest{ - WorkerId: a.workerId, - }) - } - - if err != nil { - retries++ - a.l.Error().Ctx(ctx).Err(err).Msgf("could not subscribe to the worker") - time.Sleep(DefaultActionListenerRetryInterval) - continue - } - - a.listenClient = listenClient - - return nil - } - - return fmt.Errorf("could not subscribe to the worker after %d retries", retries) -} - -func (a *actionListenerImpl) Unregister() error { - _, err := a.client.Unsubscribe( - a.ctx.newContext(context.Background()), - &dispatchercontracts.WorkerUnsubscribeRequest{ - WorkerId: a.workerId, - }, - ) - - if err != nil { - return err - } - - return nil -} - func (d *dispatcherClientImpl) GetVersion(ctx context.Context) (string, error) { resp, err := d.client.GetVersion(d.ctx.newContext(ctx), &dispatchercontracts.GetVersionRequest{}) if err != nil { diff --git a/pkg/client/durable_listener.go b/pkg/client/durable_listener.go index e7c21030c1..b789502336 100644 --- a/pkg/client/durable_listener.go +++ b/pkg/client/durable_listener.go @@ -4,18 +4,10 @@ package client import ( "context" - "errors" - "fmt" - "io" - "sync" - "time" grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry" "github.com/rs/zerolog" "golang.org/x/sync/errgroup" - "golang.org/x/sync/singleflight" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" contracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1" ) @@ -24,29 +16,43 @@ type DurableEvent *contracts.DurableEvent type DurableEventHandler func(e DurableEvent) error -type DurableEventsListener struct { - constructor func(context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) - - client contracts.V1Dispatcher_ListenForDurableEventClient - clientMu sync.Mutex - generation uint64 - - reconnectGroup singleflight.Group - listenMu sync.Mutex - listening bool - closed bool - - l *zerolog.Logger - - // map of workflow run ids to a list of handlers - handlers sync.Map -} - type listenTuple struct { taskId string signalKey string } +// DurableEventsListener streams durable task signals and dispatches them to +// registered one-shot handlers. +// +// NOTE: flow methods mirror WorkflowRunsListener in listener.go; keep changes +// in sync when adjusting add/ensure/dispatch behavior. +type DurableEventsListener struct { + stream *reconnectingStream[contracts.V1Dispatcher_ListenForDurableEventClient] + l *zerolog.Logger + reg handlerRegistry[listenTuple, DurableEvent] + gate listenGate +} + +func newDurableEventsListener( + l *zerolog.Logger, + constructor func(context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error), +) *DurableEventsListener { + d := &DurableEventsListener{ + reg: newHandlerRegistry[listenTuple, DurableEvent](), + l: l, + } + d.stream = newReconnectingStream( + l, + "durable event listener", + constructor, + func(client contracts.V1Dispatcher_ListenForDurableEventClient) error { + return client.CloseSend() + }, + d.replayHandlers, + ) + return d +} + func (r *subscribeClientImpl) getDurableEventsListener( ctx context.Context, ) (*DurableEventsListener, error) { @@ -57,444 +63,189 @@ func (r *subscribeClientImpl) getDurableEventsListener( return r.durableEventsListener, nil } - constructor := func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + l := newDurableEventsListener(r.l, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { return r.clientv1.ListenForDurableEvent(r.ctx.newContext(ctx), grpc_retry.Disable()) - } - - w := &DurableEventsListener{ - constructor: constructor, - l: r.l, - } - - err := w.retryListen(ctx) + }) - if err != nil { + if err := l.stream.connectSync(ctx); err != nil { return nil, err } - r.durableEventsListener = w - if !w.startListening() { - if closeErr := w.closeStream(); closeErr != nil { - r.l.Error().Ctx(ctx).Err(closeErr).Msg("failed to close durable events listener") + onExit := func() { + r.durableEventsListenerMu.Lock() + if r.durableEventsListener == l { + r.durableEventsListener = nil } + r.durableEventsListenerMu.Unlock() + } - r.durableEventsListener = nil - return nil, errListenerClosed + if err := l.startBackground(onExit); err != nil { + _ = l.stream.closeStream() + return nil, err } - go func() { - defer func() { - r.durableEventsListenerMu.Lock() - if r.durableEventsListener == w { - r.durableEventsListener = nil - } - r.durableEventsListenerMu.Unlock() - }() - defer w.stopListening() - - err := w.listen(ctx) - - if err != nil { - r.l.Error().Ctx(ctx).Err(err).Msg("failed to listen for durable events") - } - }() + r.durableEventsListener = l + return l, nil +} - return w, nil +func (l *DurableEventsListener) AddSignal( + taskId string, + signalKey string, + handler DurableEventHandler, +) error { + return l.add(listenTuple{taskId: taskId, signalKey: signalKey}, "", handler, nil) } -func (w *DurableEventsListener) retryListen(ctx context.Context) error { - if w.isClosed() { +func (l *DurableEventsListener) add( + key listenTuple, + session string, + handler DurableEventHandler, + onError func(error), +) error { + if l.stream.isClosed() { return errListenerClosed } - if err := ctx.Err(); err != nil { + lifecycle := l.stream.lifecycleContext() + if err := l.ensureListening(lifecycle); err != nil { return err } - _, err, _ := w.reconnectGroup.Do("reconnect", func() (interface{}, error) { - return nil, w.doRetryListen(context.Background()) - }) - return err -} - -func (w *DurableEventsListener) doRetryListen(ctx context.Context) error { - if w.isClosed() { - return errListenerClosed - } - - retries := 0 - - for retries < DefaultActionListenerRetryCount { - if retries > 0 { - time.Sleep(DefaultActionListenerRetryInterval) + remove := l.reg.store(key, session, handler, onError) + if err := l.retrySend(key); err != nil { + if err2 := l.ensureListening(lifecycle); err2 != nil { + remove() + return err2 } - - if w.isClosed() { - return errListenerClosed - } - - client, err := w.constructor(ctx) - - if err != nil { - retries++ - w.l.Error().Ctx(ctx).Err(err).Msgf("could not resubscribe to the durable event listener") - continue + if err2 := l.retrySend(key); err2 != nil { + remove() + return err2 } + } - var rangeErr error + if err := l.ensureListening(lifecycle); err != nil { + remove() + return err + } - w.handlers.Range(func(key, value interface{}) bool { - t := key.(listenTuple) + return nil +} - err := client.Send(&contracts.ListenForDurableEventRequest{ +func (l *DurableEventsListener) retrySend(t listenTuple) error { + return l.stream.retrySend(l.stream.lifecycleContext(), + func(c contracts.V1Dispatcher_ListenForDurableEventClient) error { + return c.Send(&contracts.ListenForDurableEventRequest{ TaskId: t.taskId, SignalKey: t.signalKey, }) - - if err != nil { - w.l.Error().Ctx(ctx).Err(err).Msgf("could not listen for durable events on the worker") - rangeErr = err - return false - } - - return true }) +} - if rangeErr != nil { - if closeErr := client.CloseSend(); closeErr != nil { - w.l.Warn().Err(closeErr).Msg("failed to close durable event stream after replay failure") - } - - retries++ - continue - } - - w.clientMu.Lock() - if w.isClosed() { - w.clientMu.Unlock() - if closeErr := client.CloseSend(); closeErr != nil { - w.l.Warn().Err(closeErr).Msg("failed to close durable event stream after listener closed") - } - - return errListenerClosed +func (l *DurableEventsListener) replayHandlers(ctx context.Context, client contracts.V1Dispatcher_ListenForDurableEventClient) error { + for _, key := range l.reg.keys() { + if err := client.Send(&contracts.ListenForDurableEventRequest{ + TaskId: key.taskId, + SignalKey: key.signalKey, + }); err != nil { + return err } - - w.client = client - w.generation++ - w.clientMu.Unlock() - return nil } - - return fmt.Errorf("could not listen for durable events on the worker after %d retries", retries) + return nil } -func (w *DurableEventsListener) getClientSnapshot() (contracts.V1Dispatcher_ListenForDurableEventClient, uint64) { - w.clientMu.Lock() - defer w.clientMu.Unlock() - return w.client, w.generation -} +func (l *DurableEventsListener) dispatch(event *contracts.DurableEvent) error { + key := listenTuple{taskId: event.TaskId, signalKey: event.SignalKey} + regs := l.reg.snapshot(key) + if len(regs) == 0 { + return nil + } -func (w *DurableEventsListener) isListening() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - return w.listening + eg := errgroup.Group{} + for _, r := range regs { + r := r + eg.Go(func() error { return r.handle(event) }) + } + if err := eg.Wait(); err != nil { + return err + } + l.reg.removeRegistrations(key, regs) + return nil } -func (w *DurableEventsListener) isClosed() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - return w.closed +func (l *DurableEventsListener) shouldReconnectOnEOF(ctx context.Context) bool { + return ctx.Err() == nil && !l.stream.isClosed() && l.reg.hasAny() } -func (w *DurableEventsListener) startListening() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - - if w.listening || w.closed { - return false +func (l *DurableEventsListener) runLoop(ctx context.Context) error { + err := listenStream(ctx, l.stream, + func(c contracts.V1Dispatcher_ListenForDurableEventClient) (*contracts.DurableEvent, error) { + return c.Recv() + }, + l.dispatch, + newStreamClassifier(l.shouldReconnectOnEOF), + ) + if err != nil { + n := l.reg.failAll(err) + l.l.Error().Err(err).Str("stream", l.stream.name).Int("handlers", n). + Msg("stream listener terminated; failing registered handlers") } - - w.listening = true - return true -} - -func (w *DurableEventsListener) stopListening() { - w.listenMu.Lock() - w.listening = false - w.listenMu.Unlock() + return err } -func (w *DurableEventsListener) ensureListening(ctx context.Context) error { - if w.isClosed() { +func (l *DurableEventsListener) ensureListening(ctx context.Context) error { + if l.stream.isClosed() { return errListenerClosed } - - if w.isListening() { + if l.gate.active() { return nil } - - if err := w.retryListen(ctx); err != nil { + if err := l.stream.connectSync(ctx); err != nil { return err } - - if !w.startListening() { - if w.isClosed() { + if !l.gate.tryStart(l.stream.isClosed()) { + if l.stream.isClosed() { return errListenerClosed } - return nil } - go func() { - defer w.stopListening() - - if err := w.listen(ctx); err != nil { - w.l.Error().Ctx(ctx).Err(err).Msg("failed to listen for durable events") - } + defer l.gate.stop() + _ = l.runLoop(l.stream.lifecycleContext()) }() - return nil } -type threadSafeDurableEventHandlers struct { - handlers []DurableEventHandler - mu sync.RWMutex -} - -func (l *DurableEventsListener) AddSignal( - taskId string, - signalKey string, - handler DurableEventHandler, -) error { - if l.isClosed() { - return errListenerClosed - } - - t := listenTuple{ - taskId: taskId, - signalKey: signalKey, - } - - if !l.isListening() { - if err := l.ensureListening(context.Background()); err != nil { - return err - } - } - - handlers, _ := l.handlers.LoadOrStore(t, &threadSafeDurableEventHandlers{ - handlers: []DurableEventHandler{}, - }) - - h := handlers.(*threadSafeDurableEventHandlers) - - h.mu.Lock() - h.handlers = append(h.handlers, handler) - l.handlers.Store(t, h) - h.mu.Unlock() - - if err := l.retrySend(t); err != nil { - if !l.isListening() { - if listenErr := l.ensureListening(context.Background()); listenErr != nil { - return listenErr - } - - return l.retrySend(t) - } - - return err - } - - return l.ensureListening(context.Background()) -} - -func (l *DurableEventsListener) retrySend(t listenTuple) error { - for i := 0; i < DefaultActionListenerRetryCount; i++ { - client, genBefore := l.getClientSnapshot() - - if client == nil { - return fmt.Errorf("client is not connected") - } - - err := client.Send(&contracts.ListenForDurableEventRequest{ - TaskId: t.taskId, - SignalKey: t.signalKey, - }) - - if err == nil { - return nil - } - - l.l.Warn().Err(err).Msgf("failed to send durable event subscription, attempt %d/%d", i+1, DefaultActionListenerRetryCount) - - if _, genAfter := l.getClientSnapshot(); genAfter != genBefore { - continue - } - - if retryErr := l.retryListen(context.Background()); retryErr != nil { - l.l.Error().Err(retryErr).Msg("failed to relisten after send failure") +func (l *DurableEventsListener) startBackground(onExit func()) error { + if !l.gate.tryStart(l.stream.isClosed()) { + if l.stream.isClosed() { + return errListenerClosed } - - time.Sleep(DefaultActionListenerRetryInterval) - } - - return fmt.Errorf("could not send to the worker after %d retries", DefaultActionListenerRetryCount) -} - -func (l *DurableEventsListener) Listen(ctx context.Context) error { - if !l.startListening() { return nil } - defer l.stopListening() - - return l.listen(ctx) -} - -func (l *DurableEventsListener) listen(ctx context.Context) error { - consecutiveErrors := 0 - maxConsecutiveErrors := 10 - - client, generation := l.getClientSnapshot() - if client == nil { - return fmt.Errorf("client is not connected") - } - defer func() { - if closeErr := client.CloseSend(); closeErr != nil { - l.l.Warn().Err(closeErr).Msg("failed to close durable event stream after listen exit") - } + go func() { + defer onExit() + defer l.gate.stop() + _ = l.runLoop(l.stream.lifecycleContext()) }() - - for { - event, err := client.Recv() - - if err != nil { - switch { - case errors.Is(err, io.EOF): - if !l.shouldReconnectOnEOF(ctx) { - return nil - } - case status.Code(err) == codes.Canceled: - return nil - case status.Code(err) == codes.Unavailable: - l.l.Warn().Err(err).Msg("dispatcher is unavailable, retrying durable event listener after 1 second") - time.Sleep(1 * time.Second) - } - - consecutiveErrors++ - - if _, genAfter := l.getClientSnapshot(); genAfter != generation { - client, generation = l.getClientSnapshot() - consecutiveErrors = 0 - continue - } - - retryErr := l.retryListen(ctx) - - if retryErr != nil { - if errors.Is(retryErr, errListenerClosed) { - return nil - } - - l.l.Error().Err(retryErr).Msgf("failed to relisten (consecutive errors: %d/%d)", consecutiveErrors, maxConsecutiveErrors) - - if consecutiveErrors >= maxConsecutiveErrors { - return fmt.Errorf("failed to relisten after %d consecutive errors: %w", consecutiveErrors, retryErr) - } - - time.Sleep(DefaultActionListenerRetryInterval) - continue - } - - client, generation = l.getClientSnapshot() - consecutiveErrors = 0 - continue - } - - consecutiveErrors = 0 - - if err := l.handleEvent(event); err != nil { - return err - } - } -} - -func (l *DurableEventsListener) Close() error { - l.listenMu.Lock() - l.closed = true - l.listenMu.Unlock() - - return l.closeStream() + return nil } -func (l *DurableEventsListener) closeStream() error { - l.clientMu.Lock() - client := l.client - l.clientMu.Unlock() - - if client == nil { +func (l *DurableEventsListener) listen(ctx context.Context) error { + if !l.gate.tryStart(l.stream.isClosed()) { return nil } - - return client.CloseSend() + defer l.gate.stop() + return l.runLoop(ctx) } -func (l *DurableEventsListener) handleEvent(e *contracts.DurableEvent) error { - t := listenTuple{ - taskId: e.TaskId, - signalKey: e.SignalKey, - } - - // find all handlers for this workflow run - handlers, ok := l.handlers.Load(t) - - if !ok { - return nil - } - - eg := errgroup.Group{} - - h := handlers.(*threadSafeDurableEventHandlers) - - h.mu.RLock() - - for _, handler := range h.handlers { - handlerCp := handler - - eg.Go(func() error { - return handlerCp(e) - }) - } - - h.mu.RUnlock() - - err := eg.Wait() - - if err == nil { - l.handlers.Delete(t) - } - - return err +func (l *DurableEventsListener) Listen(ctx context.Context) error { + return l.listen(ctx) } -func (l *DurableEventsListener) hasHandlers() bool { - hasHandlers := false - - l.handlers.Range(func(key, value interface{}) bool { - h := value.(*threadSafeDurableEventHandlers) - - h.mu.RLock() - hasHandlers = len(h.handlers) > 0 - h.mu.RUnlock() - - return !hasHandlers - }) - - return hasHandlers +func (l *DurableEventsListener) Close() error { + return l.stream.Close() } -func (l *DurableEventsListener) shouldReconnectOnEOF(ctx context.Context) bool { - if ctx.Err() != nil || l.isClosed() { - return false - } - - return l.hasHandlers() +func (l *DurableEventsListener) isListening() bool { + return l.gate.active() } diff --git a/pkg/client/durable_listener_test.go b/pkg/client/durable_listener_test.go index 211ac51997..cde0ccc234 100644 --- a/pkg/client/durable_listener_test.go +++ b/pkg/client/durable_listener_test.go @@ -9,114 +9,206 @@ import ( "time" "github.com/rs/zerolog" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/codes" - "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" contracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1" + "github.com/hatchet-dev/hatchet/pkg/client/retry" ) -type mockDurableEventClient struct { - sendFn func(req *contracts.ListenForDurableEventRequest) error - recvFn func() (*contracts.DurableEvent, error) - closeSendFn func() error - recvCh chan *contracts.DurableEvent - sendCount atomic.Int32 -} +func TestDurableEventsListenerAddSignalSendsOnceWhenStarting(t *testing.T) { + logger := zerolog.Nop() + recvCh := make(chan *contracts.DurableEvent) -type mockV1DispatcherClient struct { - listenForDurableEventFn func(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) -} + client := &mockDurableEventClient{ + recvCh: recvCh, + } -func (m *mockV1DispatcherClient) DurableTask(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_DurableTaskClient, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + return client, nil + }, client) + + require.NoError(t, listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { + return nil + })) + require.Equal(t, int32(1), client.sendCount.Load()) -func (m *mockV1DispatcherClient) RegisterDurableEvent(ctx context.Context, in *contracts.RegisterDurableEventRequest, opts ...grpc.CallOption) (*contracts.RegisterDurableEventResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") + require.NoError(t, listener.Close()) + close(recvCh) } -func (m *mockV1DispatcherClient) ListenForDurableEvent(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - if m.listenForDurableEventFn != nil { - return m.listenForDurableEventFn(ctx, opts...) +func TestDurableEventsListenerAddSignalRollsBackHandlerWhenSendFails(t *testing.T) { + logger := zerolog.Nop() + client := &mockDurableEventClient{ + sendFn: func(req *contracts.ListenForDurableEventRequest) error { + return status.Error(codes.Unavailable, "send failed") + }, + recvFn: func() (*contracts.DurableEvent, error) { + return nil, io.EOF + }, } - return nil, status.Error(codes.Unimplemented, "not implemented") -} + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + return client, nil + }, client) -func (m *mockDurableEventClient) Send(req *contracts.ListenForDurableEventRequest) error { - m.sendCount.Add(1) - if m.sendFn != nil { - return m.sendFn(req) - } - return nil + err := listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { + return nil + }) + require.Error(t, err) + assert.False(t, listener.reg.hasAny()) + + require.NoError(t, listener.Close()) + require.Eventually(t, func() bool { + return !listener.isListening() + }, time.Second, 10*time.Millisecond) } -func (m *mockDurableEventClient) Recv() (*contracts.DurableEvent, error) { - if m.recvFn != nil { - return m.recvFn() +func TestDurableEventsListenerAddSignalKeepsHandlerDuringRecovery(t *testing.T) { + logger := zerolog.Nop() + oldClient := &mockDurableEventClient{} + recoveredClient := &mockDurableEventClient{ + recvCh: make(chan *contracts.DurableEvent), } - if m.recvCh == nil { - return nil, io.EOF + + var listener *DurableEventsListener + reconnectFailures := atomic.Int32{} + recoveredSends := atomic.Int32{} + missingHandlerDuringRecovery := atomic.Bool{} + + oldClient.sendFn = func(req *contracts.ListenForDurableEventRequest) error { + listener.gate.stop() + return status.Error(codes.Unavailable, "stream broken") } - event, ok := <-m.recvCh - if !ok { - return nil, io.EOF + recoveredClient.sendFn = func(req *contracts.ListenForDurableEventRequest) error { + recoveredSends.Add(1) + if !listener.reg.hasAny() { + missingHandlerDuringRecovery.Store(true) + } + return nil } - return event, nil + + listener = newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + if reconnectFailures.Add(1) <= int32(retry.StreamSyncMaxAttempts*retry.StreamSyncMaxAttempts) { + return nil, status.Error(codes.Unavailable, "engine down") + } + + return recoveredClient, nil + }, oldClient) + require.True(t, listener.gate.tryStart(false)) + + err := listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { + return nil + }) + require.NoError(t, err) + assert.False(t, missingHandlerDuringRecovery.Load()) + assert.True(t, listener.reg.hasAny()) + assert.GreaterOrEqual(t, recoveredSends.Load(), int32(1)) + + require.NoError(t, listener.Close()) + close(recoveredClient.recvCh) + require.Eventually(t, func() bool { + return !listener.isListening() + }, time.Second, 10*time.Millisecond) } -func (m *mockDurableEventClient) CloseSend() error { - if m.closeSendFn != nil { - return m.closeSendFn() +func TestDurableEventsListenerRollbackHandlersByStableID(t *testing.T) { + tuple := listenTuple{taskId: "task-1", signalKey: "signal-1"} + + tests := []struct { + name string + rollbackFirst bool + expectedFirst int32 + expectedSecond int32 + }{ + { + name: "first handler rolls back before second", + rollbackFirst: true, + expectedSecond: 1, + }, + { + name: "second handler rolls back before first", + expectedFirst: 1, + }, } - return nil -} -func (m *mockDurableEventClient) Header() (metadata.MD, error) { - return nil, nil -} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := zerolog.Nop() + listener := newDurableEventsListener(&logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + return &mockDurableEventClient{}, nil + }) + firstCalls := atomic.Int32{} + secondCalls := atomic.Int32{} -func (m *mockDurableEventClient) Trailer() metadata.MD { - return nil -} + rollbackFirst := listener.reg.store(tuple, "", func(e DurableEvent) error { + firstCalls.Add(1) + return nil + }, nil) + rollbackSecond := listener.reg.store(tuple, "", func(e DurableEvent) error { + secondCalls.Add(1) + return nil + }, nil) -func (m *mockDurableEventClient) Context() context.Context { - return context.Background() -} + if tt.rollbackFirst { + rollbackFirst() + } else { + rollbackSecond() + } -func (m *mockDurableEventClient) SendMsg(msg interface{}) error { - return nil -} + require.NoError(t, listener.dispatch(&contracts.DurableEvent{ + TaskId: tuple.taskId, + SignalKey: tuple.signalKey, + })) -func (m *mockDurableEventClient) RecvMsg(msg interface{}) error { - return nil + assert.Equal(t, tt.expectedFirst, firstCalls.Load()) + assert.Equal(t, tt.expectedSecond, secondCalls.Load()) + assert.False(t, listener.reg.hasAny()) + }) + } } -func TestDurableEventsListenerAddSignalSendsOnceWhenStarting(t *testing.T) { +func TestDurableEventsListenerConcurrentStoreAndDispatchDoesNotOrphanHandlers(t *testing.T) { logger := zerolog.Nop() - recvCh := make(chan *contracts.DurableEvent) + listener := newDurableEventsListener(&logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + return &mockDurableEventClient{}, nil + }) + tuple := listenTuple{taskId: "task-1", signalKey: "signal-1"} + event := &contracts.DurableEvent{ + TaskId: tuple.taskId, + SignalKey: tuple.signalKey, + } + + for i := 0; i < 500; i++ { + fired := atomic.Bool{} + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + assert.NoError(t, listener.dispatch(event)) + }() + go func() { + defer wg.Done() + listener.reg.store(tuple, "", func(e DurableEvent) error { + fired.Store(true) + return nil + }, nil) + }() - client := &mockDurableEventClient{ - recvCh: recvCh, - } + wg.Wait() - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - return client, nil - }, - l: &logger, + if !fired.Load() { + require.NoError(t, listener.dispatch(event)) + require.True(t, fired.Load(), "handler stored concurrently with dispatch was orphaned") + } } - require.NoError(t, listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { - return nil - })) - require.Equal(t, int32(1), client.sendCount.Load()) - - require.NoError(t, listener.Close()) - close(recvCh) + assert.False(t, listener.reg.hasAny()) } func TestGetDurableEventsListenerImmediateAddDoesNotOpenSecondStream(t *testing.T) { @@ -270,16 +362,12 @@ func TestDurableEventsListenerReconnectsWhileRetrySendBacksOff(t *testing.T) { constructorCalled := make(chan struct{}) var closeConstructorCalled sync.Once - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - closeConstructorCalled.Do(func() { - close(constructorCalled) - }) - return newClient, nil - }, - client: oldClient, - l: &logger, - } + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + closeConstructorCalled.Do(func() { + close(constructorCalled) + }) + return newClient, nil + }, oldClient) listenErr := make(chan error, 1) go func() { @@ -315,161 +403,6 @@ func TestDurableEventsListenerReconnectsWhileRetrySendBacksOff(t *testing.T) { } } -func TestDurableEventsListenerReconnectDoesNotBlockClientSnapshot(t *testing.T) { - logger := zerolog.Nop() - constructorEntered := make(chan struct{}) - releaseConstructor := make(chan struct{}) - var closeConstructorEntered sync.Once - - oldClient := &mockDurableEventClient{ - recvCh: make(chan *contracts.DurableEvent), - } - newClient := &mockDurableEventClient{ - recvCh: make(chan *contracts.DurableEvent), - } - - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - closeConstructorEntered.Do(func() { - close(constructorEntered) - }) - <-releaseConstructor - return newClient, nil - }, - client: oldClient, - l: &logger, - } - - retryErr := make(chan error, 1) - go func() { - retryErr <- listener.retryListen(context.Background()) - }() - - select { - case <-constructorEntered: - case <-time.After(time.Second): - t.Fatal("reconnect did not reach constructor") - } - - snapshotRead := make(chan struct { - client contracts.V1Dispatcher_ListenForDurableEventClient - generation uint64 - }, 1) - go func() { - client, generation := listener.getClientSnapshot() - snapshotRead <- struct { - client contracts.V1Dispatcher_ListenForDurableEventClient - generation uint64 - }{client: client, generation: generation} - }() - - select { - case snapshot := <-snapshotRead: - require.Same(t, oldClient, snapshot.client) - require.Equal(t, uint64(0), snapshot.generation) - case <-time.After(100 * time.Millisecond): - t.Fatal("getClientSnapshot blocked behind reconnect") - } - - close(releaseConstructor) - - require.NoError(t, <-retryErr) - require.NoError(t, listener.Close()) -} - -// TestDurableEventsListenerDeliversEventAfterReconnectDuringRetryBackoff -// verifies the user-visible symptom: after a disconnect during AddSignal's -// retry window, the listener should reconnect and deliver the durable event -// from the new stream without waiting for retrySend to exhaust all attempts. -func TestDurableEventsListenerDeliversEventAfterReconnectDuringRetryBackoff(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - logger := zerolog.Nop() - sendAttempted := make(chan struct{}) - var closeSendAttempted sync.Once - - // The first stream simulates a dead gRPC stream: sends fail, and the receive - // loop observes the disconnect that should trigger reconnection. - oldClient := &mockDurableEventClient{ - sendFn: func(req *contracts.ListenForDurableEventRequest) error { - closeSendAttempted.Do(func() { - close(sendAttempted) - }) - return status.Error(codes.Unavailable, "stream broken") - }, - recvFn: func() (*contracts.DurableEvent, error) { - <-sendAttempted - return nil, status.Error(codes.Internal, "stream broken") - }, - } - - newClient := &mockDurableEventClient{ - recvCh: make(chan *contracts.DurableEvent, 1), - } - - // Reconnection immediately makes the awaited event available on the new stream. - // If reconnect is blocked by retrySend, this event will not reach the handler. - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - newClient.recvCh <- &contracts.DurableEvent{ - TaskId: "task-1", - SignalKey: "signal-1", - Data: []byte(`{"ok":true}`), - } - return newClient, nil - }, - client: oldClient, - l: &logger, - } - - listenErr := make(chan error, 1) - go func() { - listenErr <- listener.Listen(ctx) - }() - - require.Eventually(t, listener.isListening, time.Second, 10*time.Millisecond) - - received := make(chan DurableEvent, 1) - addErr := make(chan error, 1) - go func() { - addErr <- listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { - received <- e - return nil - }) - }() - - require.Eventually(t, func() bool { - select { - case <-sendAttempted: - return true - default: - return false - } - }, time.Second, 10*time.Millisecond) - - select { - case event := <-received: - require.Equal(t, "task-1", event.TaskId) - require.Equal(t, "signal-1", event.SignalKey) - require.JSONEq(t, `{"ok":true}`, string(event.Data)) - case err := <-addErr: - require.NoError(t, err) - select { - case event := <-received: - require.Equal(t, "task-1", event.TaskId) - require.Equal(t, "signal-1", event.SignalKey) - require.JSONEq(t, `{"ok":true}`, string(event.Data)) - case <-time.After(200 * time.Millisecond): - t.Fatal("AddSignal returned before durable event was delivered after reconnect") - } - case err := <-listenErr: - require.NoError(t, err) - case <-time.After(200 * time.Millisecond): - t.Fatal("expected durable event to be delivered after reconnect during retrySend backoff") - } -} - func TestDurableEventsListenerRestartsAfterListenExits(t *testing.T) { logger := zerolog.Nop() @@ -482,13 +415,9 @@ func TestDurableEventsListenerRestartsAfterListenExits(t *testing.T) { recvCh: make(chan *contracts.DurableEvent, 1), } - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - return replacementClient, nil - }, - client: initialClient, - l: &logger, - } + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + return replacementClient, nil + }, initialClient) listenErr := make(chan error, 1) go func() { @@ -534,24 +463,16 @@ func TestDurableEventsListenerReconnectsOnEOFWithRegisteredHandlers(t *testing.T } constructorCalls := atomic.Int32{} - listener := &DurableEventsListener{ - constructor: func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { - constructorCalls.Add(1) - return replacementClient, nil - }, - client: initialClient, - l: &logger, - } + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + constructorCalls.Add(1) + return replacementClient, nil + }, initialClient) received := make(chan DurableEvent, 1) - listener.handlers.Store(listenTuple{taskId: "task-1", signalKey: "signal-1"}, &threadSafeDurableEventHandlers{ - handlers: []DurableEventHandler{ - func(e DurableEvent) error { - received <- e - return nil - }, - }, - }) + listener.reg.store(listenTuple{taskId: "task-1", signalKey: "signal-1"}, "", func(e DurableEvent) error { + received <- e + return nil + }, nil) listenErr := make(chan error, 1) go func() { @@ -580,3 +501,110 @@ func TestDurableEventsListenerReconnectsOnEOFWithRegisteredHandlers(t *testing.T close(replacementClient.recvCh) require.NoError(t, <-listenErr) } + +func TestAddSignalBoundedWhenEngineDown(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} + + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + constructorCalls.Add(1) + return nil, status.Error(codes.Unavailable, "engine down") + }, nil) + + err := listener.AddSignal("task-1", "signal-1", func(e DurableEvent) error { + return nil + }) + + require.Error(t, err) + assert.LessOrEqual(t, constructorCalls.Load(), int32(retry.StreamSyncMaxAttempts+1)) +} + +func TestListenForDurableEventsRespectsCancelledContext(t *testing.T) { + logger := zerolog.Nop() + + subscriber := &subscribeClientImpl{ + clientv1: &mockV1DispatcherClient{ + listenForDurableEventFn: func(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + t.Fatal("constructor should not run when caller context is already cancelled") + return nil, nil + }, + }, + l: &logger, + ctx: newContextLoader("", nil), + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := subscriber.ListenForDurableEvents(ctx) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestListenForDurableEventsPassesCallerContextToInitialListen(t *testing.T) { + logger := zerolog.Nop() + callerCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + + subscriber := &subscribeClientImpl{ + clientv1: &mockV1DispatcherClient{ + listenForDurableEventFn: func(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + assert.NotEqual(t, context.Background(), ctx) + cancel() + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("constructor context was not cancelled with caller context") + } + return nil, status.Error(codes.Unavailable, "engine down") + }, + }, + l: &logger, + ctx: newContextLoader("", nil), + } + + _, err := subscriber.ListenForDurableEvents(callerCtx) + require.Error(t, err) +} + +func TestListenForDurableEventsRespectsDeadlineContext(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} + + subscriber := &subscribeClientImpl{ + clientv1: &mockV1DispatcherClient{ + listenForDurableEventFn: func(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + constructorCalls.Add(1) + <-ctx.Done() + return nil, ctx.Err() + }, + }, + l: &logger, + ctx: newContextLoader("", nil), + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + + _, err := subscriber.ListenForDurableEvents(ctx) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Equal(t, int32(1), constructorCalls.Load()) +} + +func TestDurableListenEOFWithoutHandlersDoesNotReconnect(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} + + listener := newTestDurableEventsListener(t, &logger, func(ctx context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + constructorCalls.Add(1) + return &mockDurableEventClient{recvCh: make(chan *contracts.DurableEvent)}, nil + }, &mockDurableEventClient{ + recvFn: func() (*contracts.DurableEvent, error) { + return nil, io.EOF + }, + }) + + require.NoError(t, listener.Listen(context.Background())) + assert.Equal(t, int32(0), constructorCalls.Load()) +} diff --git a/pkg/client/durable_task_listener.go b/pkg/client/durable_task_listener.go index 2ffaad8604..3d3eefd338 100644 --- a/pkg/client/durable_task_listener.go +++ b/pkg/client/durable_task_listener.go @@ -13,6 +13,7 @@ import ( "google.golang.org/grpc/status" v1 "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1" + "github.com/hatchet-dev/hatchet/pkg/client/retry" ) const ( @@ -337,7 +338,9 @@ func (l *DurableTaskListener) receiveLoop(ctx context.Context) { return } l.l.Error().Err(err).Msg("DurableTaskListener: connection failed, retrying") - time.Sleep(l.reconnectInterval) + if sleepErr := retry.Sleep(ctx, l.reconnectInterval); sleepErr != nil { + return + } continue } @@ -352,7 +355,9 @@ func (l *DurableTaskListener) receiveLoop(ctx context.Context) { } l.failPendingAcks(fmt.Errorf("connection reset: %w", err)) l.l.Warn().Err(err).Msg("DurableTaskListener: stream ended, reconnecting") - time.Sleep(l.reconnectInterval) + if sleepErr := retry.Sleep(ctx, l.reconnectInterval); sleepErr != nil { + return + } continue } @@ -360,7 +365,9 @@ func (l *DurableTaskListener) receiveLoop(ctx context.Context) { if isCancelled(ctx) { return } - time.Sleep(l.reconnectInterval) + if sleepErr := retry.Sleep(ctx, l.reconnectInterval); sleepErr != nil { + return + } } } diff --git a/pkg/client/durable_task_listener_test.go b/pkg/client/durable_task_listener_test.go index fe4acfef22..b9442324a8 100644 --- a/pkg/client/durable_task_listener_test.go +++ b/pkg/client/durable_task_listener_test.go @@ -491,3 +491,36 @@ func TestBufferedCompletionDeliveredToLateConsumer(t *testing.T) { assert.Equal(t, 0, listener.BufferedCompletionCount()) } + +func TestDurableTaskListenerCancellationDuringSleep(t *testing.T) { + l := zerolog.Nop() + connectCh := make(chan struct{}) + + listener := NewDurableTaskListener( + "test-worker", + func(ctx context.Context) (v1.V1Dispatcher_DurableTaskClient, error) { + close(connectCh) + return nil, status.Error(codes.Unavailable, "down") + }, + &l, + WithReconnectInterval(time.Second), + ) + + ctx, cancel := context.WithCancel(context.Background()) + listener.Start(ctx) + + require.Eventually(t, func() bool { + select { + case <-connectCh: + return true + default: + return false + } + }, 2*time.Second, 10*time.Millisecond) + + cancel() + + require.Eventually(t, func() bool { + return !listener.IsRunning() + }, 2*time.Second, 10*time.Millisecond) +} diff --git a/pkg/client/handler_registry.go b/pkg/client/handler_registry.go new file mode 100644 index 0000000000..bf45fbb2db --- /dev/null +++ b/pkg/client/handler_registry.go @@ -0,0 +1,218 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "fmt" + "sync" +) + +// registration is one stored handler plus its optional permanent-failure +// callback and the registration id that makes stale removal a no-op. +type registration[E any] struct { + handle func(E) error + onError func(error) + regID uint64 +} + +// handlerRegistry is a keyed registry of event handlers shared by the +// subscription listeners. It is a pure container: one RWMutex over nested +// plain maps, with stale-remove protection via registration ids (replaces the +// tombstone protocol). E is the public handler event type. +type handlerRegistry[K comparable, E any] struct { + buckets map[K]map[string]registration[E] + nextReg uint64 + mu sync.RWMutex +} + +func newHandlerRegistry[K comparable, E any]() handlerRegistry[K, E] { + return handlerRegistry[K, E]{ + buckets: make(map[K]map[string]registration[E]), + } +} + +// store registers h under (k, session) and returns an idempotent remove +// closure. An empty session auto-generates a unique one (durable events' +// auto-increment semantics). The closure removes only the exact registration +// it created: if (k, session) was re-registered since, the stale closure is a +// no-op (the ABA guarantee formerly provided by tombstoned buckets). +func (r *handlerRegistry[K, E]) store(k K, session string, h func(E) error, onError func(error)) (remove func()) { + r.mu.Lock() + defer r.mu.Unlock() + + if session == "" { + r.nextReg++ + session = fmt.Sprintf("auto-%d", r.nextReg) + } + + r.nextReg++ + regID := r.nextReg + + bucket := r.buckets[k] + if bucket == nil { + bucket = make(map[string]registration[E]) + r.buckets[k] = bucket + } + + bucket[session] = registration[E]{ + handle: h, + onError: onError, + regID: regID, + } + + return func() { + r.mu.Lock() + defer r.mu.Unlock() + + bucket, ok := r.buckets[k] + if !ok { + return + } + + existing, ok := bucket[session] + if !ok || existing.regID != regID { + return + } + + delete(bucket, session) + if len(bucket) == 0 { + delete(r.buckets, k) + } + } +} + +// removeSession unconditionally removes whatever is registered under +// (k, session). Backs RemoveWorkflowRun. +func (r *handlerRegistry[K, E]) removeSession(k K, session string) { + r.mu.Lock() + defer r.mu.Unlock() + + bucket, ok := r.buckets[k] + if !ok { + return + } + + delete(bucket, session) + if len(bucket) == 0 { + delete(r.buckets, k) + } +} + +// hasAny reports whether any handler is registered (EOF reconnect policy). +func (r *handlerRegistry[K, E]) hasAny() bool { + r.mu.RLock() + defer r.mu.RUnlock() + + for _, bucket := range r.buckets { + if len(bucket) > 0 { + return true + } + } + return false +} + +// keys snapshots the registered keys (replay after reconnect). +func (r *handlerRegistry[K, E]) keys() []K { + r.mu.RLock() + defer r.mu.RUnlock() + + keys := make([]K, 0, len(r.buckets)) + for k := range r.buckets { + keys = append(keys, k) + } + return keys +} + +// snapshot returns the registrations for k without holding the lock during +// dispatch. +func (r *handlerRegistry[K, E]) snapshot(k K) []registration[E] { + r.mu.RLock() + defer r.mu.RUnlock() + + bucket, ok := r.buckets[k] + if !ok || len(bucket) == 0 { + return nil + } + + regs := make([]registration[E], 0, len(bucket)) + for _, reg := range bucket { + regs = append(regs, reg) + } + return regs +} + +// removeRegistrations removes exactly the given registrations (matched by +// regID) under k. Backs one-shot dispatch: a handler stored concurrently with +// a dispatch is never removed. +func (r *handlerRegistry[K, E]) removeRegistrations(k K, regs []registration[E]) { + r.mu.Lock() + defer r.mu.Unlock() + + bucket, ok := r.buckets[k] + if !ok { + return + } + + for _, reg := range regs { + for session, existing := range bucket { + if existing.regID == reg.regID { + delete(bucket, session) + break + } + } + } + + if len(bucket) == 0 { + delete(r.buckets, k) + } +} + +// failAll empties the registry, invokes every non-nil onError with err, and +// returns how many registrations were dropped (for the caller's log line). +func (r *handlerRegistry[K, E]) failAll(err error) int { + r.mu.Lock() + defer r.mu.Unlock() + + count := 0 + for _, bucket := range r.buckets { + for _, reg := range bucket { + count++ + if reg.onError != nil { + reg.onError(err) + } + } + } + + r.buckets = make(map[K]map[string]registration[E]) + return count +} + +// listenGate ensures at most one listen loop runs per listener. +type listenGate struct { + mu sync.Mutex + listening bool +} + +func (g *listenGate) tryStart(closed bool) bool { + g.mu.Lock() + defer g.mu.Unlock() + + if g.listening || closed { + return false + } + + g.listening = true + return true +} + +func (g *listenGate) stop() { + g.mu.Lock() + g.listening = false + g.mu.Unlock() +} + +func (g *listenGate) active() bool { + g.mu.Lock() + defer g.mu.Unlock() + return g.listening +} diff --git a/pkg/client/handler_registry_test.go b/pkg/client/handler_registry_test.go new file mode 100644 index 0000000000..d321228ce6 --- /dev/null +++ b/pkg/client/handler_registry_test.go @@ -0,0 +1,115 @@ +package client + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestHandlerRegistryStaleRemoveIsNoOp(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + remove := reg.store("key", "session-a", func(int) error { return nil }, nil) + reg.removeSession("key", "session-a") + + called := false + reg.store("key", "session-b", func(int) error { + called = true + return nil + }, nil) + + remove() + regs := reg.snapshot("key") + require.Len(t, regs, 1) + require.NoError(t, regs[0].handle(1)) + assert.True(t, called) +} + +func TestHandlerRegistryStaleRemoveSameSessionIsNoOp(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + remove := reg.store("key", "session", func(int) error { return nil }, nil) + reg.removeSession("key", "session") + + called := false + reg.store("key", "session", func(int) error { + called = true + return nil + }, nil) + + remove() + regs := reg.snapshot("key") + require.Len(t, regs, 1) + require.NoError(t, regs[0].handle(1)) + assert.True(t, called) +} + +func TestHandlerRegistryFailAllNotifiesAndEmptiesRegistry(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + var notified atomic.Int32 + reg.store("k1", "s1", func(int) error { return nil }, func(error) { notified.Add(1) }) + reg.store("k2", "s2", func(int) error { return nil }, func(error) { notified.Add(1) }) + reg.store("k3", "s3", func(int) error { return nil }, nil) + + err := errors.New("boom") + count := reg.failAll(err) + + assert.Equal(t, 3, count) + assert.Equal(t, int32(2), notified.Load()) + assert.False(t, reg.hasAny()) +} + +func TestHandlerRegistryRemoveRegistrationsSparesConcurrentStores(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + regs := reg.snapshot("key") + require.Empty(t, regs) + + removeFirst := reg.store("key", "first", func(int) error { return nil }, nil) + firstSnapshot := reg.snapshot("key") + require.Len(t, firstSnapshot, 1) + + reg.store("key", "second", func(int) error { return nil }, nil) + + reg.removeRegistrations("key", firstSnapshot) + removeFirst() + + second := reg.snapshot("key") + require.Len(t, second, 1) +} + +func TestHandlerRegistryAutoSessionsAreUnique(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + reg.store("key", "", func(int) error { return nil }, nil) + reg.store("key", "", func(int) error { return nil }, nil) + + regs := reg.snapshot("key") + assert.Len(t, regs, 2) +} + +func TestHandlerRegistryConcurrentStoreAndSnapshotRace(t *testing.T) { + reg := newHandlerRegistry[string, int]() + + const iterations = 500 + var wg sync.WaitGroup + + for i := 0; i < iterations; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + key := "key" + remove := reg.store(key, "", func(int) error { return nil }, nil) + snap := reg.snapshot(key) + reg.removeRegistrations(key, snap) + remove() + }(i) + } + + wg.Wait() +} diff --git a/pkg/client/listener.go b/pkg/client/listener.go index 294618122c..570e16996c 100644 --- a/pkg/client/listener.go +++ b/pkg/client/listener.go @@ -5,7 +5,6 @@ package client import ( "context" "errors" - "fmt" "io" "sync" "time" @@ -13,18 +12,13 @@ import ( grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry" "github.com/rs/zerolog" "golang.org/x/sync/errgroup" - "golang.org/x/sync/singleflight" "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" sharedcontracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1" "github.com/hatchet-dev/hatchet/pkg/validator" ) -var errListenerClosed = errors.New("listener is closed") - // ResourceType represents the type of resource type ResourceType int32 @@ -97,22 +91,33 @@ type RunHandler func(event WorkflowEvent) error type StreamHandler func(event StreamEvent) error type WorkflowRunEventHandler func(event WorkflowRunEvent) error +// WorkflowRunsListener streams workflow run completion events and dispatches +// them to per-run session handlers. type WorkflowRunsListener struct { - constructor func(context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) - - client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient - clientMu sync.Mutex - generation uint64 - - reconnectGroup singleflight.Group - listenMu sync.Mutex - listening bool - closed bool - - l *zerolog.Logger + stream *reconnectingStream[dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient] + reg handlerRegistry[string, WorkflowRunEvent] + gate listenGate + l *zerolog.Logger +} - // map of workflow run ids to a list of handlers - handlers sync.Map +func newWorkflowRunsListener( + l *zerolog.Logger, + constructor func(context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error), +) *WorkflowRunsListener { + w := &WorkflowRunsListener{ + reg: newHandlerRegistry[string, WorkflowRunEvent](), + l: l, + } + w.stream = newReconnectingStream( + l, + "workflow run listener", + constructor, + func(client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return client.CloseSend() + }, + w.replayHandlers, + ) + return w } func (r *subscribeClientImpl) getWorkflowRunsListener( @@ -125,462 +130,187 @@ func (r *subscribeClientImpl) getWorkflowRunsListener( return r.workflowRunListener, nil } - constructor := func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + l := newWorkflowRunsListener(r.l, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { return r.client.SubscribeToWorkflowRuns(r.ctx.newContext(ctx), grpc_retry.Disable()) - } - - w := &WorkflowRunsListener{ - constructor: constructor, - l: r.l, - } - - err := w.retrySubscribe(ctx) + }) - if err != nil { + if err := l.stream.connectSync(ctx); err != nil { return nil, err } - r.workflowRunListener = w - if !w.startListening() { - if closeErr := w.closeStream(); closeErr != nil { - r.l.Error().Err(closeErr).Msg("failed to close workflow run events listener") + onExit := func() { + r.workflowRunListenerMu.Lock() + if r.workflowRunListener == l { + r.workflowRunListener = nil } - - r.workflowRunListener = nil - return nil, errListenerClosed + r.workflowRunListenerMu.Unlock() } - go func() { - defer func() { - r.workflowRunListenerMu.Lock() - if r.workflowRunListener == w { - r.workflowRunListener = nil - } - r.workflowRunListenerMu.Unlock() - }() - defer w.stopListening() - - err := w.listen(ctx) - - if err != nil { - r.l.Error().Err(err).Msg("failed to listen for workflow run events") - } - }() - - return w, nil -} - -// getClientSnapshot returns the current client and its generation without holding the lock. -func (w *WorkflowRunsListener) getClientSnapshot() (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, uint64) { - w.clientMu.Lock() - defer w.clientMu.Unlock() - return w.client, w.generation -} - -func (w *WorkflowRunsListener) isListening() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - return w.listening -} - -func (w *WorkflowRunsListener) isClosed() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - return w.closed -} - -func (w *WorkflowRunsListener) startListening() bool { - w.listenMu.Lock() - defer w.listenMu.Unlock() - - if w.listening || w.closed { - return false + if err := l.startBackground(onExit); err != nil { + _ = l.stream.closeStream() + return nil, err } - w.listening = true - return true + r.workflowRunListener = l + return l, nil } -func (w *WorkflowRunsListener) stopListening() { - w.listenMu.Lock() - w.listening = false - w.listenMu.Unlock() +func (l *WorkflowRunsListener) AddWorkflowRun( + workflowRunId, sessionId string, + handler WorkflowRunEventHandler, +) error { + return l.addWorkflowRun(workflowRunId, sessionId, handler, nil) } -func (w *WorkflowRunsListener) ensureListening(ctx context.Context) error { - if w.isClosed() { +func (l *WorkflowRunsListener) addWorkflowRun( + workflowRunId, sessionId string, + handler WorkflowRunEventHandler, + onError func(error), +) error { + if l.stream.isClosed() { return errListenerClosed } - if w.isListening() { - return nil - } - - if err := w.retrySubscribe(ctx); err != nil { + lifecycle := l.stream.lifecycleContext() + if err := l.ensureListening(lifecycle); err != nil { return err } - if !w.startListening() { - if w.isClosed() { - return errListenerClosed + remove := l.reg.store(workflowRunId, sessionId, handler, onError) + if err := l.retrySend(workflowRunId); err != nil { + if err2 := l.ensureListening(lifecycle); err2 != nil { + remove() + return err2 } - - return nil - } - - go func() { - defer w.stopListening() - - if err := w.listen(ctx); err != nil { - w.l.Error().Err(err).Msg("failed to listen for workflow run events") + if err2 := l.retrySend(workflowRunId); err2 != nil { + remove() + return err2 } - }() - - return nil -} - -// retrySubscribe coalesces concurrent reconnection attempts via singleflight. -// Multiple goroutines calling this concurrently will share a single reconnection attempt. -func (w *WorkflowRunsListener) retrySubscribe(ctx context.Context) error { - if w.isClosed() { - return errListenerClosed } - if err := ctx.Err(); err != nil { + if err := l.ensureListening(lifecycle); err != nil { + remove() return err } - _, err, _ := w.reconnectGroup.Do("reconnect", func() (interface{}, error) { - return nil, w.doRetrySubscribe(context.Background()) - }) - return err + return nil } -func (w *WorkflowRunsListener) doRetrySubscribe(ctx context.Context) error { - if w.isClosed() { - return errListenerClosed - } - - retries := 0 - - for retries < DefaultActionListenerRetryCount { - if retries > 0 { - time.Sleep(DefaultActionListenerRetryInterval) - } - - if w.isClosed() { - return errListenerClosed - } - - client, err := w.constructor(ctx) - - if err != nil { - retries++ - w.l.Error().Err(err).Msgf("could not resubscribe to the listener") - continue - } - - var rangeErr error - - w.handlers.Range(func(key, value interface{}) bool { - workflowRunId := key.(string) - - err := client.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{ - WorkflowRunId: workflowRunId, - }) - - if err != nil { - w.l.Error().Err(err).Msgf("could not subscribe to the worker for workflow run id %s", workflowRunId) - rangeErr = err - return false - } - - return true - }) - - if rangeErr != nil { - if closeErr := client.CloseSend(); closeErr != nil { - w.l.Warn().Err(closeErr).Msg("failed to close workflow run stream after replay failure") - } - - retries++ - continue - } - - w.clientMu.Lock() - if w.isClosed() { - w.clientMu.Unlock() - if closeErr := client.CloseSend(); closeErr != nil { - w.l.Warn().Err(closeErr).Msg("failed to close workflow run stream after listener closed") - } - - return errListenerClosed - } - - w.client = client - w.generation++ - w.clientMu.Unlock() - return nil - } - - return fmt.Errorf("could not subscribe to the worker after %d retries", retries) +func (l *WorkflowRunsListener) RemoveWorkflowRun(workflowRunId, sessionId string) { + l.reg.removeSession(workflowRunId, sessionId) } -type threadSafeHandlers struct { - // map of session ids to handlers - handlers map[string]WorkflowRunEventHandler - mu sync.RWMutex +func (l *WorkflowRunsListener) retrySend(workflowRunId string) error { + return l.stream.retrySend(l.stream.lifecycleContext(), + func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: workflowRunId}) + }) } -func (l *WorkflowRunsListener) AddWorkflowRun( - workflowRunId, sessionId string, - handler WorkflowRunEventHandler, -) error { - if l.isClosed() { - return errListenerClosed - } - - if !l.isListening() { - if err := l.ensureListening(context.Background()); err != nil { +func (l *WorkflowRunsListener) replayHandlers(ctx context.Context, client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + for _, runId := range l.reg.keys() { + if err := client.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: runId}); err != nil { return err } } - - handlers, _ := l.handlers.LoadOrStore(workflowRunId, &threadSafeHandlers{ - handlers: map[string]WorkflowRunEventHandler{}, - }) - - h := handlers.(*threadSafeHandlers) - - h.mu.Lock() - h.handlers[sessionId] = handler - l.handlers.Store(workflowRunId, h) - h.mu.Unlock() - - if err := l.retrySend(workflowRunId); err != nil { - if !l.isListening() { - if listenErr := l.ensureListening(context.Background()); listenErr != nil { - return listenErr - } - - return l.retrySend(workflowRunId) - } - - return err - } - - return l.ensureListening(context.Background()) + return nil } -func (l *WorkflowRunsListener) RemoveWorkflowRun( - workflowRunId, sessionId string, -) { - handlers, ok := l.handlers.Load(workflowRunId) - - if !ok { - return +func (l *WorkflowRunsListener) dispatch(event *dispatchercontracts.WorkflowRunEvent) error { + regs := l.reg.snapshot(event.WorkflowRunId) + if len(regs) == 0 { + return nil } - h := handlers.(*threadSafeHandlers) - - h.mu.Lock() - defer h.mu.Unlock() - - delete(h.handlers, sessionId) + ev, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + if err != nil { + return err + } - if len(h.handlers) == 0 { - l.handlers.Delete(workflowRunId) + eg := errgroup.Group{} + for _, r := range regs { + r := r + eg.Go(func() error { return r.handle(ev) }) } + return eg.Wait() } -func (l *WorkflowRunsListener) retrySend(workflowRunId string) error { - for i := 0; i < DefaultActionListenerRetryCount; i++ { - client, genBefore := l.getClientSnapshot() - - if client == nil { - return fmt.Errorf("client is not connected") - } - - err := client.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{ - WorkflowRunId: workflowRunId, - }) - - if err == nil { - return nil - } - - l.l.Warn().Err(err).Msgf("failed to send workflow run subscription, attempt %d/%d", i+1, DefaultActionListenerRetryCount) - - // Check if someone else (e.g. Listen) already reconnected while we were sending. - // If so, skip the reconnect and retry the send on the new client immediately. - if _, genAfter := l.getClientSnapshot(); genAfter != genBefore { - continue - } - - if retryErr := l.retrySubscribe(context.Background()); retryErr != nil { - l.l.Error().Err(retryErr).Msg("failed to resubscribe after send failure") - } +func (l *WorkflowRunsListener) shouldReconnectOnEOF(ctx context.Context) bool { + return ctx.Err() == nil && !l.stream.isClosed() && l.reg.hasAny() +} - time.Sleep(DefaultActionListenerRetryInterval) +func (l *WorkflowRunsListener) runLoop(ctx context.Context) error { + err := listenStream(ctx, l.stream, + func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) (*dispatchercontracts.WorkflowRunEvent, error) { + return c.Recv() + }, + l.dispatch, + newStreamClassifier(l.shouldReconnectOnEOF), + ) + if err != nil { + n := l.reg.failAll(err) + l.l.Error().Err(err).Str("stream", l.stream.name).Int("handlers", n). + Msg("stream listener terminated; failing registered handlers") } - - return fmt.Errorf("could not send to the worker after %d retries", DefaultActionListenerRetryCount) + return err } -func (l *WorkflowRunsListener) Listen(ctx context.Context) error { - if !l.startListening() { +func (l *WorkflowRunsListener) ensureListening(ctx context.Context) error { + if l.stream.isClosed() { + return errListenerClosed + } + if l.gate.active() { return nil } - defer l.stopListening() - - return l.listen(ctx) -} - -func (l *WorkflowRunsListener) listen(ctx context.Context) error { - consecutiveErrors := 0 - maxConsecutiveErrors := 10 - - // Take a snapshot of the client so we never hold the lock during a blocking Recv. - client, generation := l.getClientSnapshot() - if client == nil { - return fmt.Errorf("client is not connected") + if err := l.stream.connectSync(ctx); err != nil { + return err } - defer func() { - if closeErr := client.CloseSend(); closeErr != nil { - l.l.Warn().Err(closeErr).Msg("failed to close workflow run stream after listen exit") - } - }() - - for { - event, err := client.Recv() - - if err != nil { - switch { - case errors.Is(err, io.EOF): - if !l.shouldReconnectOnEOF(ctx) { - return nil - } - case status.Code(err) == codes.Canceled: - return nil - case status.Code(err) == codes.Unavailable: - l.l.Warn().Err(err).Msg("dispatcher is unavailable, retrying subscribe after 1 second") - time.Sleep(1 * time.Second) - } - - consecutiveErrors++ - - if _, genAfter := l.getClientSnapshot(); genAfter != generation { - client, generation = l.getClientSnapshot() - consecutiveErrors = 0 - continue - } - - retryErr := l.retrySubscribe(ctx) - - if retryErr != nil { - if errors.Is(retryErr, errListenerClosed) { - return nil - } - - l.l.Error().Err(retryErr).Msgf("failed to resubscribe (consecutive errors: %d/%d)", consecutiveErrors, maxConsecutiveErrors) - - if consecutiveErrors >= maxConsecutiveErrors { - return fmt.Errorf("failed to resubscribe after %d consecutive errors: %w", consecutiveErrors, retryErr) - } - - time.Sleep(DefaultActionListenerRetryInterval) - continue - } - - client, generation = l.getClientSnapshot() - consecutiveErrors = 0 - continue - } - - consecutiveErrors = 0 - - if err := l.handleWorkflowRun(event); err != nil { - return err + if !l.gate.tryStart(l.stream.isClosed()) { + if l.stream.isClosed() { + return errListenerClosed } + return nil } + go func() { + defer l.gate.stop() + _ = l.runLoop(l.stream.lifecycleContext()) + }() + return nil } -func (l *WorkflowRunsListener) Close() error { - l.listenMu.Lock() - l.closed = true - l.listenMu.Unlock() - - return l.closeStream() -} - -func (l *WorkflowRunsListener) closeStream() error { - l.clientMu.Lock() - client := l.client - l.clientMu.Unlock() - - if client == nil { +func (l *WorkflowRunsListener) startBackground(onExit func()) error { + if !l.gate.tryStart(l.stream.isClosed()) { + if l.stream.isClosed() { + return errListenerClosed + } return nil } - - return client.CloseSend() + go func() { + defer onExit() + defer l.gate.stop() + _ = l.runLoop(l.stream.lifecycleContext()) + }() + return nil } -func (l *WorkflowRunsListener) handleWorkflowRun(event *dispatchercontracts.WorkflowRunEvent) error { - // find all handlers for this workflow run - handlers, ok := l.handlers.Load(event.WorkflowRunId) - - if !ok { +func (l *WorkflowRunsListener) listen(ctx context.Context) error { + if !l.gate.tryStart(l.stream.isClosed()) { return nil } - - eg := errgroup.Group{} - - h := handlers.(*threadSafeHandlers) - - h.mu.RLock() - - for _, handler := range h.handlers { - handlerCp := handler - - eg.Go(func() error { - workflowRunEvent, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) - if err != nil { - return err - } - - return handlerCp(workflowRunEvent) - }) - } - - h.mu.RUnlock() - - err := eg.Wait() - - return err + defer l.gate.stop() + return l.runLoop(ctx) } -func (l *WorkflowRunsListener) hasHandlers() bool { - hasHandlers := false - - l.handlers.Range(func(key, value interface{}) bool { - h := value.(*threadSafeHandlers) - - h.mu.RLock() - hasHandlers = len(h.handlers) > 0 - h.mu.RUnlock() - - return !hasHandlers - }) - - return hasHandlers +func (l *WorkflowRunsListener) Listen(ctx context.Context) error { + return l.listen(ctx) } -func (l *WorkflowRunsListener) shouldReconnectOnEOF(ctx context.Context) bool { - if ctx.Err() != nil || l.isClosed() { - return false - } +func (l *WorkflowRunsListener) Close() error { + return l.stream.Close() +} - return l.hasHandlers() +func (l *WorkflowRunsListener) isListening() bool { + return l.gate.active() } type SubscribeClient interface { @@ -694,45 +424,75 @@ func (r *subscribeClientImpl) Stream(ctx context.Context, workflowRunId string, } } -func (r *subscribeClientImpl) StreamByAdditionalMetadata(ctx context.Context, key string, value string, handler StreamHandler) error { - stream, err := r.client.SubscribeToWorkflowEvents(r.ctx.newContext(ctx), &dispatchercontracts.SubscribeToWorkflowEventsRequest{ - AdditionalMetaKey: &key, - AdditionalMetaValue: &value, - }) - - if err != nil { - return err +func (r *subscribeClientImpl) newMetadataStream(ctx context.Context, key, value string) *reconnectingStream[dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient] { + return newReconnectingStreamWithLifecycle( + ctx, + r.l, + "metadata stream listener", + func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { + return r.client.SubscribeToWorkflowEvents(r.ctx.newContext(ctx), &dispatchercontracts.SubscribeToWorkflowEventsRequest{ + AdditionalMetaKey: &key, + AdditionalMetaValue: &value, + }, grpc_retry.Disable()) + }, + func(client dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient) error { + return client.CloseSend() + }, + nil, + ) +} + +func metadataStreamClassifier() streamClassifier { + base := newStreamClassifier(func(ctx context.Context) bool { return ctx.Err() == nil }) + return func(ctx context.Context, err error) streamVerdict { + if v := base(ctx, err); v != verdictNoProgress { + return v + } + return verdictStopError } +} - for { - event, err := stream.Recv() +func runMetadataStream( + ctx context.Context, + stream *reconnectingStream[dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient], + handler StreamHandler, +) error { + return listenStream( + ctx, + stream, + func(client dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient) (*dispatchercontracts.WorkflowEvent, error) { + return client.Recv() + }, + func(event *dispatchercontracts.WorkflowEvent) error { + if event.EventType != dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STREAM { + return ctx.Err() + } - if err != nil { - if errors.Is(err, io.EOF) { - return nil + if err := handler(StreamEvent{ + Message: []byte(event.EventPayload), + }); err != nil { + return err } - return err - } + return ctx.Err() + }, + metadataStreamClassifier(), + ) +} - if event.EventType != dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STREAM { - continue - } +func (r *subscribeClientImpl) StreamByAdditionalMetadata(ctx context.Context, key string, value string, handler StreamHandler) error { + stream := r.newMetadataStream(ctx, key, value) + defer func() { _ = stream.Close() }() - if err := handler(StreamEvent{ - Message: []byte(event.EventPayload), - }); err != nil { - return err - } - } + return runMetadataStream(ctx, stream, handler) } func (r *subscribeClientImpl) SubscribeToWorkflowRunEvents(ctx context.Context) (*WorkflowRunsListener, error) { - return r.getWorkflowRunsListener(context.Background()) + return r.getWorkflowRunsListener(ctx) } func (r *subscribeClientImpl) ListenForDurableEvents(ctx context.Context) (*DurableEventsListener, error) { - return r.getDurableEventsListener(context.Background()) + return r.getDurableEventsListener(ctx) } func workflowEventToDeprecatedWorkflowEvent(event *dispatchercontracts.WorkflowEvent) (*workflowEvent, error) { diff --git a/pkg/client/listener_convert_test.go b/pkg/client/listener_convert_test.go new file mode 100644 index 0000000000..8d5a05da7e --- /dev/null +++ b/pkg/client/listener_convert_test.go @@ -0,0 +1,124 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/timestamppb" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" +) + +func TestWorkflowEventToDeprecatedWorkflowEvent_Success(t *testing.T) { + workflowRunId := "test-workflow-run-123" + stepId := "test-step-456" + eventPayload := "test-payload" + + event := &dispatchercontracts.WorkflowEvent{ + WorkflowRunId: workflowRunId, + ResourceType: dispatchercontracts.ResourceType_RESOURCE_TYPE_STEP_RUN, + ResourceId: stepId, + EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STARTED, + EventPayload: eventPayload, + EventTimestamp: timestamppb.New(time.Unix(1234567890, 123456789)), + } + + deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) + + require.NoError(t, err, "conversion should succeed") + require.NotNil(t, deprecated, "deprecated event should not be nil") + + assert.Equal(t, workflowRunId, deprecated.WorkflowRunId) + assert.Equal(t, stepId, deprecated.ResourceId) + assert.Equal(t, eventPayload, deprecated.EventPayload) +} + +func TestWorkflowEventToDeprecatedWorkflowEvent_WithNilTimestamp(t *testing.T) { + event := &dispatchercontracts.WorkflowEvent{ + WorkflowRunId: "test-run", + ResourceId: "test-resource", + EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_COMPLETED, + EventPayload: "payload", + EventTimestamp: nil, + } + + deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) + + require.NoError(t, err, "conversion should succeed even with nil timestamp") + require.NotNil(t, deprecated, "deprecated event should not be nil") + assert.Equal(t, "test-run", deprecated.WorkflowRunId) +} + +func TestWorkflowEventToDeprecatedWorkflowEvent_EmptyEvent(t *testing.T) { + event := &dispatchercontracts.WorkflowEvent{} + + deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) + + require.NoError(t, err, "conversion should succeed with empty event") + require.NotNil(t, deprecated, "deprecated event should not be nil") +} + +func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_Success(t *testing.T) { + workflowRunId := "test-workflow-run-789" + + event := &dispatchercontracts.WorkflowRunEvent{ + WorkflowRunId: workflowRunId, + EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, + EventTimestamp: timestamppb.New(time.Unix(9876543210, 987654321)), + } + + deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + + require.NoError(t, err, "conversion should succeed") + require.NotNil(t, deprecated, "deprecated event should not be nil") + + assert.Equal(t, workflowRunId, deprecated.WorkflowRunId) + assert.NotNil(t, deprecated.EventTimestamp) +} + +func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_WithNilTimestamp(t *testing.T) { + event := &dispatchercontracts.WorkflowRunEvent{ + WorkflowRunId: "test-run-2", + EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, + EventTimestamp: nil, + } + + deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + + require.NoError(t, err, "conversion should succeed even with nil timestamp") + require.NotNil(t, deprecated, "deprecated event should not be nil") + assert.Equal(t, "test-run-2", deprecated.WorkflowRunId) +} + +func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_EmptyEvent(t *testing.T) { + event := &dispatchercontracts.WorkflowRunEvent{} + + deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + + require.NoError(t, err, "conversion should succeed with empty event") + require.NotNil(t, deprecated, "deprecated event should not be nil") +} + +func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_WithResults(t *testing.T) { + event := &dispatchercontracts.WorkflowRunEvent{ + WorkflowRunId: "test-run-3", + EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, + EventTimestamp: timestamppb.New(time.Now()), + Results: []*dispatchercontracts.StepRunResult{ + { + TaskRunExternalId: "step-1", + }, + }, + } + + deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + + require.NoError(t, err, "conversion should succeed with results") + require.NotNil(t, deprecated, "deprecated event should not be nil") + assert.Equal(t, "test-run-3", deprecated.WorkflowRunId) + assert.NotNil(t, deprecated.Results) +} diff --git a/pkg/client/listener_test.go b/pkg/client/listener_test.go index f343d7167a..84a9d0024d 100644 --- a/pkg/client/listener_test.go +++ b/pkg/client/listener_test.go @@ -4,7 +4,6 @@ package client import ( "context" - "fmt" "io" "sync" "sync/atomic" @@ -16,163 +15,102 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/codes" - "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" - "google.golang.org/protobuf/types/known/timestamppb" dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" + "github.com/hatchet-dev/hatchet/pkg/client/retry" ) -// mockSubscribeClient implements dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient -type mockSubscribeClient struct { - sendErr error - sendCount atomic.Int32 - recvErr error - recvChan chan *dispatchercontracts.WorkflowRunEvent - closeSendFn func() error - sendFn func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error - recvFn func() (*dispatchercontracts.WorkflowRunEvent, error) -} - -type mockDispatcherClient struct { - subscribeToWorkflowRunsFn func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) -} - -func (m *mockDispatcherClient) Register(ctx context.Context, in *dispatchercontracts.WorkerRegisterRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerRegisterResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) Listen(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) ListenV2(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) Heartbeat(ctx context.Context, in *dispatchercontracts.HeartbeatRequest, opts ...grpc.CallOption) (*dispatchercontracts.HeartbeatResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) SubscribeToWorkflowEvents(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} +func TestWorkflowRunsListenerAddWorkflowRunSendsOnceWhenStarting(t *testing.T) { + logger := zerolog.Nop() + recvChan := make(chan *dispatchercontracts.WorkflowRunEvent) -func (m *mockDispatcherClient) SubscribeToWorkflowRuns(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - if m.subscribeToWorkflowRunsFn != nil { - return m.subscribeToWorkflowRunsFn(ctx, opts...) + client := &mockSubscribeClient{ + recvChan: recvChan, } - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) SendStepActionEvent(ctx context.Context, in *dispatchercontracts.StepActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }, client) -func (m *mockDispatcherClient) SendGroupKeyActionEvent(ctx context.Context, in *dispatchercontracts.GroupKeyActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) PutOverridesData(ctx context.Context, in *dispatchercontracts.OverridesData, opts ...grpc.CallOption) (*dispatchercontracts.OverridesDataResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} - -func (m *mockDispatcherClient) Unsubscribe(ctx context.Context, in *dispatchercontracts.WorkerUnsubscribeRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerUnsubscribeResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} + require.NoError(t, listener.AddWorkflowRun("run-1", "session-1", func(event WorkflowRunEvent) error { + return nil + })) + assert.Equal(t, int32(1), client.sendCount.Load()) -func (m *mockDispatcherClient) RefreshTimeout(ctx context.Context, in *dispatchercontracts.RefreshTimeoutRequest, opts ...grpc.CallOption) (*dispatchercontracts.RefreshTimeoutResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") + require.NoError(t, listener.Close()) + close(recvChan) } -func (m *mockDispatcherClient) ReleaseSlot(ctx context.Context, in *dispatchercontracts.ReleaseSlotRequest, opts ...grpc.CallOption) (*dispatchercontracts.ReleaseSlotResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} +func TestWorkflowRunsListenerAddWorkflowRunRollsBackHandlerWhenSendFails(t *testing.T) { + logger := zerolog.Nop() + client := &mockSubscribeClient{ + sendErr: status.Error(codes.Unavailable, "send failed"), + recvErr: io.EOF, + } -func (m *mockDispatcherClient) RestoreEvictedTask(ctx context.Context, in *dispatchercontracts.RestoreEvictedTaskRequest, opts ...grpc.CallOption) (*dispatchercontracts.RestoreEvictedTaskResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }, client) -func (m *mockDispatcherClient) UpsertWorkerLabels(ctx context.Context, in *dispatchercontracts.UpsertWorkerLabelsRequest, opts ...grpc.CallOption) (*dispatchercontracts.UpsertWorkerLabelsResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") -} + err := listener.AddWorkflowRun("run-1", "session-1", func(event WorkflowRunEvent) error { + return nil + }) + require.Error(t, err) + assert.False(t, listener.reg.hasAny()) -func (m *mockDispatcherClient) GetVersion(ctx context.Context, in *dispatchercontracts.GetVersionRequest, opts ...grpc.CallOption) (*dispatchercontracts.GetVersionResponse, error) { - return nil, status.Error(codes.Unimplemented, "not implemented") + require.NoError(t, listener.Close()) + require.Eventually(t, func() bool { + return !listener.isListening() + }, time.Second, 10*time.Millisecond) } -func (m *mockSubscribeClient) Send(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { - m.sendCount.Add(1) - if m.sendFn != nil { - return m.sendFn(req) +func TestWorkflowRunsListenerAddWorkflowRunKeepsHandlerDuringRecovery(t *testing.T) { + logger := zerolog.Nop() + oldClient := &mockSubscribeClient{} + recoveredClient := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), } - return m.sendErr -} -func (m *mockSubscribeClient) Recv() (*dispatchercontracts.WorkflowRunEvent, error) { - if m.recvFn != nil { - return m.recvFn() - } - if m.recvErr != nil { - return nil, m.recvErr - } - event, ok := <-m.recvChan - if !ok { - return nil, io.EOF - } - return event, nil -} + var listener *WorkflowRunsListener + reconnectFailures := atomic.Int32{} + recoveredSends := atomic.Int32{} + missingHandlerDuringRecovery := atomic.Bool{} -func (m *mockSubscribeClient) CloseSend() error { - if m.closeSendFn != nil { - return m.closeSendFn() + oldClient.sendFn = func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + listener.gate.stop() + return status.Error(codes.Unavailable, "stream broken") } - return nil -} - -func (m *mockSubscribeClient) Header() (metadata.MD, error) { - return nil, nil -} - -func (m *mockSubscribeClient) Trailer() metadata.MD { - return nil -} - -func (m *mockSubscribeClient) Context() context.Context { - return context.Background() -} - -func (m *mockSubscribeClient) SendMsg(msg interface{}) error { - return nil -} - -func (m *mockSubscribeClient) RecvMsg(msg interface{}) error { - return nil -} - -func TestWorkflowRunsListenerAddWorkflowRunSendsOnceWhenStarting(t *testing.T) { - logger := zerolog.Nop() - recvChan := make(chan *dispatchercontracts.WorkflowRunEvent) - - client := &mockSubscribeClient{ - recvChan: recvChan, + recoveredClient.sendFn = func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + recoveredSends.Add(1) + if !listener.reg.hasAny() { + missingHandlerDuringRecovery.Store(true) + } + return nil } - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return client, nil - }, - l: &logger, - } + listener = newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + if reconnectFailures.Add(1) <= int32(retry.StreamSyncMaxAttempts*retry.StreamSyncMaxAttempts) { + return nil, status.Error(codes.Unavailable, "engine down") + } - require.NoError(t, listener.AddWorkflowRun("run-1", "session-1", func(event WorkflowRunEvent) error { + return recoveredClient, nil + }, oldClient) + require.True(t, listener.gate.tryStart(false)) + + err := listener.AddWorkflowRun("run-1", "session-1", func(event WorkflowRunEvent) error { return nil - })) - assert.Equal(t, int32(1), client.sendCount.Load()) + }) + require.NoError(t, err) + assert.False(t, missingHandlerDuringRecovery.Load()) + assert.True(t, listener.reg.hasAny()) + assert.GreaterOrEqual(t, recoveredSends.Load(), int32(1)) require.NoError(t, listener.Close()) - close(recvChan) + close(recoveredClient.recvChan) + require.Eventually(t, func() bool { + return !listener.isListening() + }, time.Second, 10*time.Millisecond) } func TestGetWorkflowRunsListenerImmediateAddDoesNotOpenSecondStream(t *testing.T) { @@ -282,243 +220,7 @@ func TestGetWorkflowRunsListenerInternalCleanupDoesNotTerminallyCloseListener(t }, time.Second, 10*time.Millisecond) } -func TestRetrySend_ResubscribesOnSendFailure(t *testing.T) { - // This test verifies that when Send() fails on a broken stream, - // retrySend will call retrySubscribe to establish a new stream - // and then successfully send on the new stream. - - logger := zerolog.Nop() - - // Track how many times the constructor is called - constructorCalls := atomic.Int32{} - - // First client will fail on Send - failingClient := &mockSubscribeClient{ - sendErr: fmt.Errorf("stream broken"), - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - // Second client will succeed - workingClient := &mockSubscribeClient{ - sendErr: nil, - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - calls := constructorCalls.Add(1) - if calls == 1 { - // First call during initial setup - return failingClient, nil - } - // Subsequent calls after resubscribe - return workingClient, nil - }, - client: failingClient, - l: &logger, - } - - // Call retrySend - it should fail on first attempt, resubscribe, then succeed - err := listener.retrySend("test-workflow-run-id") - - require.NoError(t, err, "retrySend should succeed after resubscribing") - - // Verify constructor was called (for resubscribe) - assert.GreaterOrEqual(t, constructorCalls.Load(), int32(1), "constructor should be called for resubscribe") - - // Verify the failing client received at least one send attempt - assert.GreaterOrEqual(t, failingClient.sendCount.Load(), int32(1), "failing client should have received send attempts") - - // Verify the working client received a send - assert.Equal(t, int32(1), workingClient.sendCount.Load(), "working client should have received exactly one send") -} - -func TestWorkflowRunsListenerReconnectDoesNotBlockClientSnapshot(t *testing.T) { - logger := zerolog.Nop() - constructorEntered := make(chan struct{}) - releaseConstructor := make(chan struct{}) - var closeConstructorEntered sync.Once - - oldClient := &mockSubscribeClient{ - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - newClient := &mockSubscribeClient{ - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - closeConstructorEntered.Do(func() { - close(constructorEntered) - }) - <-releaseConstructor - return newClient, nil - }, - client: oldClient, - l: &logger, - } - - retryErr := make(chan error, 1) - go func() { - retryErr <- listener.retrySubscribe(context.Background()) - }() - - select { - case <-constructorEntered: - case <-time.After(time.Second): - t.Fatal("reconnect did not reach constructor") - } - - snapshotRead := make(chan struct { - client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient - generation uint64 - }, 1) - go func() { - client, generation := listener.getClientSnapshot() - snapshotRead <- struct { - client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient - generation uint64 - }{client: client, generation: generation} - }() - - select { - case snapshot := <-snapshotRead: - require.Same(t, oldClient, snapshot.client) - require.Equal(t, uint64(0), snapshot.generation) - case <-time.After(100 * time.Millisecond): - t.Fatal("getClientSnapshot blocked behind reconnect") - } - - close(releaseConstructor) - - require.NoError(t, <-retryErr) - require.NoError(t, listener.Close()) -} - -func TestRetrySend_FailsAfterMaxRetries(t *testing.T) { - // This test verifies that retrySend returns an error after - // exhausting all retry attempts when resubscribe also fails. - - logger := zerolog.Nop() - - // Client that always fails - failingClient := &mockSubscribeClient{ - sendErr: fmt.Errorf("stream broken"), - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - // Constructor also returns a failing client - return failingClient, nil - }, - client: failingClient, - l: &logger, - } - - // Call retrySend - it should fail after all retries - err := listener.retrySend("test-workflow-run-id") - - require.Error(t, err, "retrySend should fail after exhausting retries") - assert.Contains(t, err.Error(), "could not send to the worker after", "error should indicate retry exhaustion") - - // Verify multiple send attempts were made - assert.Equal(t, int32(DefaultActionListenerRetryCount), failingClient.sendCount.Load(), - "should have attempted send DefaultActionListenerRetryCount times") -} - -func TestRetrySend_SucceedsOnFirstAttempt(t *testing.T) { - // This test verifies that retrySend returns immediately - // when Send() succeeds on the first attempt. - - logger := zerolog.Nop() - - workingClient := &mockSubscribeClient{ - sendErr: nil, - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - constructorCalls := atomic.Int32{} - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - constructorCalls.Add(1) - return workingClient, nil - }, - client: workingClient, - l: &logger, - } - - err := listener.retrySend("test-workflow-run-id") - - require.NoError(t, err, "retrySend should succeed on first attempt") - assert.Equal(t, int32(1), workingClient.sendCount.Load(), "should have sent exactly once") - assert.Equal(t, int32(0), constructorCalls.Load(), "constructor should not be called when first send succeeds") -} - -func TestRetrySend_HandlesNilClient(t *testing.T) { - // This test verifies that retrySend returns an error when client is nil - - logger := zerolog.Nop() - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return nil, fmt.Errorf("connection failed") - }, - client: nil, - l: &logger, - } - - err := listener.retrySend("test-workflow-run-id") - - require.Error(t, err, "retrySend should fail when client is nil") - assert.Contains(t, err.Error(), "client is not connected", "error should indicate client not connected") -} - -func TestRetrySend_ConcurrentSafety(t *testing.T) { - // This test verifies that retrySend is safe to call concurrently - - logger := zerolog.Nop() - - workingClient := &mockSubscribeClient{ - sendErr: nil, - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - sendFn: func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { - time.Sleep(10 * time.Millisecond) // Simulate some work - return nil - }, - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return workingClient, nil - }, - client: workingClient, - l: &logger, - } - - // Launch multiple concurrent retrySend calls - var wg sync.WaitGroup - numGoroutines := 10 - - for i := 0; i < numGoroutines; i++ { - wg.Add(1) - go func(id int) { - defer wg.Done() - err := listener.retrySend(fmt.Sprintf("workflow-run-%d", id)) - assert.NoError(t, err, "concurrent retrySend should succeed") - }(i) - } - - wg.Wait() - - // All sends should have completed - assert.Equal(t, int32(numGoroutines), workingClient.sendCount.Load(), "all concurrent sends should complete") -} - func TestListen_DispatchesEventsToHandlers(t *testing.T) { - // Verifies that Listen receives events and dispatches them to registered handlers. - logger := zerolog.Nop() recvChan := make(chan *dispatchercontracts.WorkflowRunEvent, 1) @@ -526,31 +228,21 @@ func TestListen_DispatchesEventsToHandlers(t *testing.T) { recvChan: recvChan, } - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return client, nil - }, - client: client, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }, client) var receivedEvent atomic.Value - listener.handlers.Store("run-1", &threadSafeHandlers{ - handlers: map[string]WorkflowRunEventHandler{ - "session-1": func(event WorkflowRunEvent) error { - receivedEvent.Store(event) - return nil - }, - }, - }) + listener.reg.store("run-1", "session-1", func(event WorkflowRunEvent) error { + receivedEvent.Store(event) + return nil + }, nil) - // Start Listen in a goroutine listenErr := make(chan error, 1) go func() { listenErr <- listener.Listen(context.Background()) }() - // Send an event recvChan <- &dispatchercontracts.WorkflowRunEvent{ WorkflowRunId: "run-1", } @@ -558,62 +250,18 @@ func TestListen_DispatchesEventsToHandlers(t *testing.T) { require.Eventually(t, func() bool { return receivedEvent.Load() != nil }, time.Second, 10*time.Millisecond) - listener.handlers.Delete("run-1") + listener.reg.removeSession("run-1", "session-1") - // Close the channel to trigger EOF and cleanly end Listen close(recvChan) err := <-listenErr assert.NoError(t, err, "Listen should exit cleanly on EOF") - // Verify handler was called ev := receivedEvent.Load() require.NotNil(t, ev, "handler should have been called") assert.Equal(t, "run-1", ev.(WorkflowRunEvent).WorkflowRunId) } -func TestListen_ExitsOnEOF(t *testing.T) { - // Verifies that Listen returns nil (clean exit) on io.EOF. - - logger := zerolog.Nop() - - client := &mockSubscribeClient{ - recvErr: io.EOF, - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return client, nil - }, - client: client, - l: &logger, - } - - err := listener.Listen(context.Background()) - assert.NoError(t, err, "Listen should return nil on EOF") -} - -func TestListen_ExitsOnCanceled(t *testing.T) { - // Verifies that Listen returns nil on gRPC Canceled status. - - logger := zerolog.Nop() - - client := &mockSubscribeClient{ - recvErr: status.Error(codes.Canceled, "context canceled"), - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return client, nil - }, - client: client, - l: &logger, - } - - err := listener.Listen(context.Background()) - assert.NoError(t, err, "Listen should return nil on Canceled") -} - func TestWorkflowRunsListenerRestartsAfterListenExits(t *testing.T) { logger := zerolog.Nop() @@ -624,13 +272,9 @@ func TestWorkflowRunsListenerRestartsAfterListenExits(t *testing.T) { recvChan: make(chan *dispatchercontracts.WorkflowRunEvent, 1), } - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return replacementClient, nil - }, - client: initialClient, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return replacementClient, nil + }, initialClient) listenErr := make(chan error, 1) go func() { @@ -674,24 +318,16 @@ func TestWorkflowRunsListenerReconnectsOnEOFWithRegisteredHandlers(t *testing.T) } constructorCalls := atomic.Int32{} - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - constructorCalls.Add(1) - return replacementClient, nil - }, - client: initialClient, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return replacementClient, nil + }, initialClient) received := make(chan WorkflowRunEvent, 1) - listener.handlers.Store("run-1", &threadSafeHandlers{ - handlers: map[string]WorkflowRunEventHandler{ - "session-1": func(event WorkflowRunEvent) error { - received <- event - return nil - }, - }, - }) + listener.reg.store("run-1", "session-1", func(event WorkflowRunEvent) error { + received <- event + return nil + }, nil) listenErr := make(chan error, 1) go func() { @@ -726,24 +362,16 @@ func TestWorkflowRunsListenerClosePreventsEOFReconnectAndAdd(t *testing.T) { recvErr: io.EOF, } - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - constructorCalls.Add(1) - return &mockSubscribeClient{ - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - }, nil - }, - client: initialClient, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + }, nil + }, initialClient) - listener.handlers.Store("run-1", &threadSafeHandlers{ - handlers: map[string]WorkflowRunEventHandler{ - "session-1": func(event WorkflowRunEvent) error { - return nil - }, - }, - }) + listener.reg.store("run-1", "session-1", func(event WorkflowRunEvent) error { + return nil + }, nil) require.NoError(t, listener.Close()) require.NoError(t, listener.listen(context.Background())) @@ -755,89 +383,18 @@ func TestWorkflowRunsListenerClosePreventsEOFReconnectAndAdd(t *testing.T) { require.ErrorIs(t, err, errListenerClosed) } -func TestListen_ReconnectsAndUsesNewClient(t *testing.T) { - // Verifies that after a Recv error, Listen reconnects and reads from the new client. - - logger := zerolog.Nop() - - newRecvChan := make(chan *dispatchercontracts.WorkflowRunEvent, 1) - - // First client errors immediately on Recv - brokenClient := &mockSubscribeClient{ - recvErr: fmt.Errorf("stream broken"), - } - - // Second client delivers an event then closes cleanly (EOF) - workingClient := &mockSubscribeClient{ - recvChan: newRecvChan, - } - - constructorCalls := atomic.Int32{} - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - constructorCalls.Add(1) - return workingClient, nil - }, - client: brokenClient, - l: &logger, - } - - var receivedEvent atomic.Value - listener.handlers.Store("run-1", &threadSafeHandlers{ - handlers: map[string]WorkflowRunEventHandler{ - "session-1": func(event WorkflowRunEvent) error { - receivedEvent.Store(event) - return nil - }, - }, - }) - - listenErr := make(chan error, 1) - go func() { - listenErr <- listener.Listen(context.Background()) - }() - - // Send an event on the new client, then close to trigger EOF (clean exit) - newRecvChan <- &dispatchercontracts.WorkflowRunEvent{ - WorkflowRunId: "run-1", - } - - require.Eventually(t, func() bool { - return receivedEvent.Load() != nil - }, time.Second, 10*time.Millisecond) - listener.handlers.Delete("run-1") - close(newRecvChan) - - err := <-listenErr - assert.NoError(t, err, "Listen should exit cleanly on EOF after reconnection") - - // Verify the constructor was called for reconnection - assert.GreaterOrEqual(t, constructorCalls.Load(), int32(1), "constructor should have been called to reconnect") - - // Verify handler received the event from the new client - ev := receivedEvent.Load() - require.NotNil(t, ev, "handler should have received event from new client") - assert.Equal(t, "run-1", ev.(WorkflowRunEvent).WorkflowRunId) -} - func TestClose_NilClient(t *testing.T) { - // Verifies that Close returns nil when client is nil. - logger := zerolog.Nop() - listener := &WorkflowRunsListener{ - client: nil, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return nil, nil + }, nil) err := listener.Close() assert.NoError(t, err, "Close should return nil for nil client") } func TestClose_CallsCloseSend(t *testing.T) { - // Verifies that Close calls CloseSend on the client. - logger := zerolog.Nop() closeCalled := atomic.Bool{} @@ -849,229 +406,252 @@ func TestClose_CallsCloseSend(t *testing.T) { }, } - listener := &WorkflowRunsListener{ - client: client, - l: &logger, - } + listener := newTestWorkflowRunsListener(t, &logger, nil, client) err := listener.Close() assert.NoError(t, err) assert.True(t, closeCalled.Load(), "CloseSend should have been called") } -func TestRetrySubscribe_SingleflightCoalescesConcurrentCalls(t *testing.T) { - // Verifies that concurrent retrySubscribe calls are coalesced into one - // actual reconnection via singleflight. - +func TestAddWorkflowRunBoundedWhenEngineDown(t *testing.T) { logger := zerolog.Nop() constructorCalls := atomic.Int32{} - client := &mockSubscribeClient{ - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return nil, status.Error(codes.Unavailable, "engine down") + }, nil) - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - constructorCalls.Add(1) - // Simulate some latency so concurrent calls overlap - time.Sleep(50 * time.Millisecond) - return client, nil - }, - client: client, - l: &logger, - } + err := listener.AddWorkflowRun("run-1", "session-1", func(event WorkflowRunEvent) error { + return nil + }) - numGoroutines := 10 - var wg sync.WaitGroup + require.Error(t, err) + assert.LessOrEqual(t, constructorCalls.Load(), int32(retry.StreamSyncMaxAttempts+1)) +} - for i := 0; i < numGoroutines; i++ { - wg.Add(1) - go func() { - defer wg.Done() - err := listener.retrySubscribe(context.Background()) - assert.NoError(t, err) - }() - } +func TestListenEOFWithoutHandlersDoesNotReconnect(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} - wg.Wait() + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return &mockSubscribeClient{recvChan: make(chan *dispatchercontracts.WorkflowRunEvent)}, nil + }, &mockSubscribeClient{recvErr: io.EOF}) - // singleflight should coalesce all concurrent calls into 1 constructor invocation - assert.Equal(t, int32(1), constructorCalls.Load(), - "concurrent retrySubscribe calls should be coalesced into a single reconnection") + require.NoError(t, listener.Listen(context.Background())) + assert.Equal(t, int32(0), constructorCalls.Load()) } -func TestRetrySubscribe_GenerationIncrements(t *testing.T) { - // Verifies that the generation counter increments on each successful reconnection. - +func TestRetrySubscribeSyncReplaysHandlers(t *testing.T) { logger := zerolog.Nop() + sendCount := atomic.Int32{} client := &mockSubscribeClient{ - recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), - } - - listener := &WorkflowRunsListener{ - constructor: func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { - return client, nil + sendFn: func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + sendCount.Add(1) + return nil }, - client: client, - l: &logger, + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), } - _, gen0 := listener.getClientSnapshot() - assert.Equal(t, uint64(0), gen0, "initial generation should be 0") - - err := listener.retrySubscribe(context.Background()) - require.NoError(t, err) - - _, gen1 := listener.getClientSnapshot() - assert.Equal(t, uint64(1), gen1, "generation should be 1 after first reconnect") + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }, client) - err = listener.retrySubscribe(context.Background()) - require.NoError(t, err) + listener.reg.store("run-1", "session-1", func(event WorkflowRunEvent) error { return nil }, nil) - _, gen2 := listener.getClientSnapshot() - assert.Equal(t, uint64(2), gen2, "generation should be 2 after second reconnect") + require.NoError(t, listener.stream.connectSync(context.Background())) + assert.Equal(t, int32(1), sendCount.Load()) } -func TestGetClientSnapshot_ReturnsCurrentClient(t *testing.T) { - // Verifies that getClientSnapshot returns the current client and generation. - +func TestSubscribeToWorkflowRunsDisablesGrpcRetry(t *testing.T) { logger := zerolog.Nop() - client1 := &mockSubscribeClient{recvChan: make(chan *dispatchercontracts.WorkflowRunEvent)} - - listener := &WorkflowRunsListener{ - client: client1, - l: &logger, + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowRunsFn: func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + require.NotEmpty(t, opts) + return &mockSubscribeClient{recvChan: make(chan *dispatchercontracts.WorkflowRunEvent)}, nil + }, + }, + l: &logger, + ctx: newContextLoader("", nil), } - got, gen := listener.getClientSnapshot() - assert.Equal(t, client1, got) - assert.Equal(t, uint64(0), gen) + listener, err := subscriber.getWorkflowRunsListener(context.Background()) + require.NoError(t, err) + require.NoError(t, listener.Close()) } -func TestWorkflowEventToDeprecatedWorkflowEvent_Success(t *testing.T) { - // Verifies successful conversion of WorkflowEvent to deprecated WorkflowEvent +func TestStreamByAdditionalMetadataReconnects(t *testing.T) { + logger := zerolog.Nop() + establishCalls := atomic.Int32{} + recvCh := make(chan *dispatchercontracts.WorkflowEvent, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - workflowRunId := "test-workflow-run-123" - stepId := "test-step-456" - eventPayload := "test-payload" + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowEventsFn: func(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { + require.NotEmpty(t, opts) + call := establishCalls.Add(1) + if call == 1 { + return &mockWorkflowEventsClient{ + recvFn: func() (*dispatchercontracts.WorkflowEvent, error) { + return nil, status.Error(codes.Unavailable, "stream broken") + }, + }, nil + } - event := &dispatchercontracts.WorkflowEvent{ - WorkflowRunId: workflowRunId, - ResourceType: dispatchercontracts.ResourceType_RESOURCE_TYPE_STEP_RUN, - ResourceId: stepId, - EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STARTED, - EventPayload: eventPayload, - EventTimestamp: timestamppb.New(time.Unix(1234567890, 123456789)), + return &mockWorkflowEventsClient{recvCh: recvCh}, nil + }, + }, + l: &logger, + ctx: newContextLoader("", nil), } - deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) - - require.NoError(t, err, "conversion should succeed") - require.NotNil(t, deprecated, "deprecated event should not be nil") - - // Verify key fields are preserved - assert.Equal(t, workflowRunId, deprecated.WorkflowRunId) - assert.Equal(t, stepId, deprecated.ResourceId) - assert.Equal(t, eventPayload, deprecated.EventPayload) -} - -func TestWorkflowEventToDeprecatedWorkflowEvent_WithNilTimestamp(t *testing.T) { - // Verifies conversion works when timestamp is nil + done := make(chan error, 1) + go func() { + stream := subscriber.newMetadataStream(ctx, "k", "v") + disableStreamBackoff(t, stream) + defer func() { _ = stream.Close() }() + done <- runMetadataStream(ctx, stream, func(event StreamEvent) error { + cancel() + return nil + }) + }() - event := &dispatchercontracts.WorkflowEvent{ - WorkflowRunId: "test-run", - ResourceId: "test-resource", - EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_COMPLETED, - EventPayload: "payload", - EventTimestamp: nil, + recvCh <- &dispatchercontracts.WorkflowEvent{ + EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STREAM, + EventPayload: "hello", } - deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) - - require.NoError(t, err, "conversion should succeed even with nil timestamp") - require.NotNil(t, deprecated, "deprecated event should not be nil") - assert.Equal(t, "test-run", deprecated.WorkflowRunId) + err := <-done + require.ErrorIs(t, err, context.Canceled) + assert.GreaterOrEqual(t, establishCalls.Load(), int32(2)) } -func TestWorkflowEventToDeprecatedWorkflowEvent_EmptyEvent(t *testing.T) { - // Verifies conversion works with an empty event - - event := &dispatchercontracts.WorkflowEvent{} - - deprecated, err := workflowEventToDeprecatedWorkflowEvent(event) - - require.NoError(t, err, "conversion should succeed with empty event") - require.NotNil(t, deprecated, "deprecated event should not be nil") -} - -func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_Success(t *testing.T) { - // Verifies successful conversion of WorkflowRunEvent to deprecated WorkflowRunEvent +func TestStreamByAdditionalMetadataReconnectsOnEOF(t *testing.T) { + logger := zerolog.Nop() + establishCalls := atomic.Int32{} + handlerCalls := atomic.Int32{} + secondRecvCh := make(chan *dispatchercontracts.WorkflowEvent, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() - workflowRunId := "test-workflow-run-789" + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowEventsFn: func(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { + require.NotEmpty(t, opts) + call := establishCalls.Add(1) + if call == 1 { + return &mockWorkflowEventsClient{ + recvFn: func() (*dispatchercontracts.WorkflowEvent, error) { + return nil, io.EOF + }, + }, nil + } - event := &dispatchercontracts.WorkflowRunEvent{ - WorkflowRunId: workflowRunId, - EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, - EventTimestamp: timestamppb.New(time.Unix(9876543210, 987654321)), + return &mockWorkflowEventsClient{recvCh: secondRecvCh}, nil + }, + }, + l: &logger, + ctx: newContextLoader("", nil), } - deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + done := make(chan error, 1) + go func() { + stream := subscriber.newMetadataStream(ctx, "k", "v") + disableStreamBackoff(t, stream) + defer func() { _ = stream.Close() }() + done <- runMetadataStream(ctx, stream, func(event StreamEvent) error { + handlerCalls.Add(1) + cancel() + return nil + }) + }() - require.NoError(t, err, "conversion should succeed") - require.NotNil(t, deprecated, "deprecated event should not be nil") + secondRecvCh <- &dispatchercontracts.WorkflowEvent{ + EventType: dispatchercontracts.ResourceEventType_RESOURCE_EVENT_TYPE_STREAM, + EventPayload: "hello", + } - // Verify key fields are preserved - assert.Equal(t, workflowRunId, deprecated.WorkflowRunId) - assert.NotNil(t, deprecated.EventTimestamp) + err := <-done + require.ErrorIs(t, err, context.Canceled) + assert.Equal(t, int32(1), handlerCalls.Load()) + assert.GreaterOrEqual(t, establishCalls.Load(), int32(2)) } -func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_WithNilTimestamp(t *testing.T) { - // Verifies conversion works when timestamp is nil +func TestSubscribeToWorkflowRunEventsRespectsCancelledContext(t *testing.T) { + logger := zerolog.Nop() - event := &dispatchercontracts.WorkflowRunEvent{ - WorkflowRunId: "test-run-2", - EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, - EventTimestamp: nil, + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowRunsFn: func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + t.Fatal("constructor should not run when caller context is already cancelled") + return nil, nil + }, + }, + l: &logger, + ctx: newContextLoader("", nil), } - deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + ctx, cancel := context.WithCancel(context.Background()) + cancel() - require.NoError(t, err, "conversion should succeed even with nil timestamp") - require.NotNil(t, deprecated, "deprecated event should not be nil") - assert.Equal(t, "test-run-2", deprecated.WorkflowRunId) + _, err := subscriber.SubscribeToWorkflowRunEvents(ctx) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) } -func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_EmptyEvent(t *testing.T) { - // Verifies conversion works with an empty event - - event := &dispatchercontracts.WorkflowRunEvent{} +func TestSubscribeToWorkflowRunEventsPassesCallerContextToInitialSubscribe(t *testing.T) { + logger := zerolog.Nop() + callerCtx, cancel := context.WithCancel(context.Background()) + defer cancel() - deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowRunsFn: func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + assert.NotEqual(t, context.Background(), ctx) + cancel() + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("constructor context was not cancelled with caller context") + } + return nil, status.Error(codes.Unavailable, "engine down") + }, + }, + l: &logger, + ctx: newContextLoader("", nil), + } - require.NoError(t, err, "conversion should succeed with empty event") - require.NotNil(t, deprecated, "deprecated event should not be nil") + _, err := subscriber.SubscribeToWorkflowRunEvents(callerCtx) + require.Error(t, err) } -func TestWorkflowRunEventToDeprecatedWorkflowRunEvent_WithResults(t *testing.T) { - // Verifies conversion works with step run results +func TestSubscribeToWorkflowRunEventsRespectsDeadlineContext(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} - event := &dispatchercontracts.WorkflowRunEvent{ - WorkflowRunId: "test-run-3", - EventType: dispatchercontracts.WorkflowRunEventType_WORKFLOW_RUN_EVENT_TYPE_FINISHED, - EventTimestamp: timestamppb.New(time.Now()), - Results: []*dispatchercontracts.StepRunResult{ - { - TaskRunExternalId: "step-1", + subscriber := &subscribeClientImpl{ + client: &mockDispatcherClient{ + subscribeToWorkflowRunsFn: func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + <-ctx.Done() + return nil, ctx.Err() }, }, + l: &logger, + ctx: newContextLoader("", nil), } - deprecated, err := workflowRunEventToDeprecatedWorkflowRunEvent(event) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() - require.NoError(t, err, "conversion should succeed with results") - require.NotNil(t, deprecated, "deprecated event should not be nil") - assert.Equal(t, "test-run-3", deprecated.WorkflowRunId) - assert.NotNil(t, deprecated.Results) + _, err := subscriber.SubscribeToWorkflowRunEvents(ctx) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Equal(t, int32(1), constructorCalls.Load()) } diff --git a/pkg/client/mocks_test.go b/pkg/client/mocks_test.go new file mode 100644 index 0000000000..58aa3a178e --- /dev/null +++ b/pkg/client/mocks_test.go @@ -0,0 +1,467 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "io" + "sync/atomic" + "testing" + + "github.com/rs/zerolog" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" + contracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1" +) + +func disableStreamBackoff[C any](t *testing.T, s *reconnectingStream[C]) { + t.Helper() + s.sleep = func(context.Context, int) error { return nil } +} + +func newTestWorkflowStream( + t *testing.T, + client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, + constructor func(context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error), +) *reconnectingStream[dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient] { + t.Helper() + + logger := zerolog.Nop() + stream := newReconnectingStream( + &logger, + "workflow run listener", + constructor, + func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.CloseSend() + }, + nil, + ) + disableStreamBackoff(t, stream) + if client != nil { + stream.setInitialClient(client) + } + return stream +} + +type testListenEvent struct { + value string +} + +type testListenClient struct { + recvFn func() (testListenEvent, error) + closeSendFn func() error + closeCalled atomic.Bool +} + +func (c *testListenClient) Recv() (testListenEvent, error) { + if c.recvFn != nil { + return c.recvFn() + } + return testListenEvent{}, io.EOF +} + +func (c *testListenClient) CloseSend() error { + c.closeCalled.Store(true) + if c.closeSendFn != nil { + return c.closeSendFn() + } + return nil +} + +func newTestListenStream( + t *testing.T, + initial *testListenClient, + constructor func(context.Context) (*testListenClient, error), +) *reconnectingStream[*testListenClient] { + t.Helper() + + logger := zerolog.Nop() + stream := newReconnectingStream( + &logger, + "test listener", + constructor, + func(client *testListenClient) error { + return client.CloseSend() + }, + nil, + ) + disableStreamBackoff(t, stream) + if initial != nil { + stream.setInitialClient(initial) + } + return stream +} + +// mockSubscribeClient implements dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient +type mockSubscribeClient struct { + sendErr error + sendCount atomic.Int32 + recvErr error + recvChan chan *dispatchercontracts.WorkflowRunEvent + closeSendFn func() error + sendFn func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error + recvFn func() (*dispatchercontracts.WorkflowRunEvent, error) +} + +type mockDispatcherClient struct { + subscribeToWorkflowRunsFn func(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) + subscribeToWorkflowEventsFn func(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) +} + +func (m *mockDispatcherClient) Register(ctx context.Context, in *dispatchercontracts.WorkerRegisterRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerRegisterResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) Listen(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) ListenV2(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) Heartbeat(ctx context.Context, in *dispatchercontracts.HeartbeatRequest, opts ...grpc.CallOption) (*dispatchercontracts.HeartbeatResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) SubscribeToWorkflowEvents(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { + if m.subscribeToWorkflowEventsFn != nil { + return m.subscribeToWorkflowEventsFn(ctx, in, opts...) + } + + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) SubscribeToWorkflowRuns(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + if m.subscribeToWorkflowRunsFn != nil { + return m.subscribeToWorkflowRunsFn(ctx, opts...) + } + + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) SendStepActionEvent(ctx context.Context, in *dispatchercontracts.StepActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) SendGroupKeyActionEvent(ctx context.Context, in *dispatchercontracts.GroupKeyActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) PutOverridesData(ctx context.Context, in *dispatchercontracts.OverridesData, opts ...grpc.CallOption) (*dispatchercontracts.OverridesDataResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) Unsubscribe(ctx context.Context, in *dispatchercontracts.WorkerUnsubscribeRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerUnsubscribeResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) RefreshTimeout(ctx context.Context, in *dispatchercontracts.RefreshTimeoutRequest, opts ...grpc.CallOption) (*dispatchercontracts.RefreshTimeoutResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) ReleaseSlot(ctx context.Context, in *dispatchercontracts.ReleaseSlotRequest, opts ...grpc.CallOption) (*dispatchercontracts.ReleaseSlotResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) RestoreEvictedTask(ctx context.Context, in *dispatchercontracts.RestoreEvictedTaskRequest, opts ...grpc.CallOption) (*dispatchercontracts.RestoreEvictedTaskResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) UpsertWorkerLabels(ctx context.Context, in *dispatchercontracts.UpsertWorkerLabelsRequest, opts ...grpc.CallOption) (*dispatchercontracts.UpsertWorkerLabelsResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDispatcherClient) GetVersion(ctx context.Context, in *dispatchercontracts.GetVersionRequest, opts ...grpc.CallOption) (*dispatchercontracts.GetVersionResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockSubscribeClient) Send(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + m.sendCount.Add(1) + if m.sendFn != nil { + return m.sendFn(req) + } + return m.sendErr +} + +func (m *mockSubscribeClient) Recv() (*dispatchercontracts.WorkflowRunEvent, error) { + if m.recvFn != nil { + return m.recvFn() + } + if m.recvErr != nil { + return nil, m.recvErr + } + event, ok := <-m.recvChan + if !ok { + return nil, io.EOF + } + return event, nil +} + +func (m *mockSubscribeClient) CloseSend() error { + if m.closeSendFn != nil { + return m.closeSendFn() + } + return nil +} + +func (m *mockSubscribeClient) Header() (metadata.MD, error) { + return nil, nil +} + +func (m *mockSubscribeClient) Trailer() metadata.MD { + return nil +} + +func (m *mockSubscribeClient) Context() context.Context { + return context.Background() +} + +func (m *mockSubscribeClient) SendMsg(msg interface{}) error { + return nil +} + +func (m *mockSubscribeClient) RecvMsg(msg interface{}) error { + return nil +} + +type mockWorkflowEventsClient struct { + recvFn func() (*dispatchercontracts.WorkflowEvent, error) + recvCh chan *dispatchercontracts.WorkflowEvent + establishFn func(ctx context.Context, opts ...grpc.CallOption) error +} + +func (m *mockWorkflowEventsClient) Recv() (*dispatchercontracts.WorkflowEvent, error) { + if m.recvFn != nil { + return m.recvFn() + } + event, ok := <-m.recvCh + if !ok { + return nil, io.EOF + } + return event, nil +} + +func (m *mockWorkflowEventsClient) CloseSend() error { return nil } +func (m *mockWorkflowEventsClient) Header() (metadata.MD, error) { return nil, nil } +func (m *mockWorkflowEventsClient) Trailer() metadata.MD { return nil } +func (m *mockWorkflowEventsClient) Context() context.Context { return context.Background() } +func (m *mockWorkflowEventsClient) SendMsg(msg interface{}) error { return nil } +func (m *mockWorkflowEventsClient) RecvMsg(msg interface{}) error { return nil } +func newTestWorkflowRunsListener( + t *testing.T, + logger *zerolog.Logger, + constructor func(context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error), + client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, +) *WorkflowRunsListener { + t.Helper() + + l := newWorkflowRunsListener(logger, constructor) + if client != nil { + l.stream.setInitialClient(client) + } + disableStreamBackoff(t, l.stream) + return l +} + +type mockDurableEventClient struct { + sendFn func(req *contracts.ListenForDurableEventRequest) error + recvFn func() (*contracts.DurableEvent, error) + closeSendFn func() error + recvCh chan *contracts.DurableEvent + sendCount atomic.Int32 +} + +type mockV1DispatcherClient struct { + listenForDurableEventFn func(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) +} + +func (m *mockV1DispatcherClient) DurableTask(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_DurableTaskClient, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockV1DispatcherClient) RegisterDurableEvent(ctx context.Context, in *contracts.RegisterDurableEventRequest, opts ...grpc.CallOption) (*contracts.RegisterDurableEventResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockV1DispatcherClient) ListenForDurableEvent(ctx context.Context, opts ...grpc.CallOption) (contracts.V1Dispatcher_ListenForDurableEventClient, error) { + if m.listenForDurableEventFn != nil { + return m.listenForDurableEventFn(ctx, opts...) + } + + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockDurableEventClient) Send(req *contracts.ListenForDurableEventRequest) error { + m.sendCount.Add(1) + if m.sendFn != nil { + return m.sendFn(req) + } + return nil +} + +func (m *mockDurableEventClient) Recv() (*contracts.DurableEvent, error) { + if m.recvFn != nil { + return m.recvFn() + } + if m.recvCh == nil { + return nil, io.EOF + } + event, ok := <-m.recvCh + if !ok { + return nil, io.EOF + } + return event, nil +} + +func (m *mockDurableEventClient) CloseSend() error { + if m.closeSendFn != nil { + return m.closeSendFn() + } + return nil +} + +func (m *mockDurableEventClient) Header() (metadata.MD, error) { + return nil, nil +} + +func (m *mockDurableEventClient) Trailer() metadata.MD { + return nil +} + +func (m *mockDurableEventClient) Context() context.Context { + return context.Background() +} + +func (m *mockDurableEventClient) SendMsg(msg interface{}) error { + return nil +} + +func (m *mockDurableEventClient) RecvMsg(msg interface{}) error { + return nil +} + +func newTestDurableEventsListener( + t *testing.T, + logger *zerolog.Logger, + constructor func(context.Context) (contracts.V1Dispatcher_ListenForDurableEventClient, error), + client contracts.V1Dispatcher_ListenForDurableEventClient, +) *DurableEventsListener { + t.Helper() + + l := newDurableEventsListener(logger, constructor) + if client != nil { + l.stream.setInitialClient(client) + } + disableStreamBackoff(t, l.stream) + return l +} + +type mockListenV2Client struct { + recvFn func() (*dispatchercontracts.AssignedAction, error) +} + +func (m *mockListenV2Client) Recv() (*dispatchercontracts.AssignedAction, error) { + if m.recvFn != nil { + return m.recvFn() + } + return nil, io.EOF +} + +func (m *mockListenV2Client) CloseSend() error { return nil } +func (m *mockListenV2Client) Header() (metadata.MD, error) { return nil, nil } +func (m *mockListenV2Client) Trailer() metadata.MD { return nil } +func (m *mockListenV2Client) Context() context.Context { return context.Background() } +func (m *mockListenV2Client) SendMsg(msg interface{}) error { return nil } +func (m *mockListenV2Client) RecvMsg(msg interface{}) error { return nil } + +type mockWorkerDispatcherClient struct { + listenV2Fn func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) + listenFn func(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) +} + +func (m *mockWorkerDispatcherClient) Register(ctx context.Context, in *dispatchercontracts.WorkerRegisterRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerRegisterResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) Listen(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenClient, error) { + if m.listenFn != nil { + return m.listenFn(ctx, in, opts...) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) ListenV2(ctx context.Context, in *dispatchercontracts.WorkerListenRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_ListenV2Client, error) { + if m.listenV2Fn != nil { + return m.listenV2Fn(ctx, in, opts...) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) Heartbeat(ctx context.Context, in *dispatchercontracts.HeartbeatRequest, opts ...grpc.CallOption) (*dispatchercontracts.HeartbeatResponse, error) { + return &dispatchercontracts.HeartbeatResponse{}, nil +} + +func (m *mockWorkerDispatcherClient) SubscribeToWorkflowEvents(ctx context.Context, in *dispatchercontracts.SubscribeToWorkflowEventsRequest, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowEventsClient, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) SubscribeToWorkflowRuns(ctx context.Context, opts ...grpc.CallOption) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) SendStepActionEvent(ctx context.Context, in *dispatchercontracts.StepActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) SendGroupKeyActionEvent(ctx context.Context, in *dispatchercontracts.GroupKeyActionEvent, opts ...grpc.CallOption) (*dispatchercontracts.ActionEventResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) PutOverridesData(ctx context.Context, in *dispatchercontracts.OverridesData, opts ...grpc.CallOption) (*dispatchercontracts.OverridesDataResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) Unsubscribe(ctx context.Context, in *dispatchercontracts.WorkerUnsubscribeRequest, opts ...grpc.CallOption) (*dispatchercontracts.WorkerUnsubscribeResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) RefreshTimeout(ctx context.Context, in *dispatchercontracts.RefreshTimeoutRequest, opts ...grpc.CallOption) (*dispatchercontracts.RefreshTimeoutResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) ReleaseSlot(ctx context.Context, in *dispatchercontracts.ReleaseSlotRequest, opts ...grpc.CallOption) (*dispatchercontracts.ReleaseSlotResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) RestoreEvictedTask(ctx context.Context, in *dispatchercontracts.RestoreEvictedTaskRequest, opts ...grpc.CallOption) (*dispatchercontracts.RestoreEvictedTaskResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) UpsertWorkerLabels(ctx context.Context, in *dispatchercontracts.UpsertWorkerLabelsRequest, opts ...grpc.CallOption) (*dispatchercontracts.UpsertWorkerLabelsResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockWorkerDispatcherClient) GetVersion(ctx context.Context, in *dispatchercontracts.GetVersionRequest, opts ...grpc.CallOption) (*dispatchercontracts.GetVersionResponse, error) { + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func newTestActionListener(t *testing.T, listenClient dispatchercontracts.Dispatcher_ListenV2Client) *actionListenerImpl { + t.Helper() + + logger := zerolog.Nop() + listener := &actionListenerImpl{ + client: &mockWorkerDispatcherClient{}, + listenClient: listenClient, + workerId: "worker-1", + l: &logger, + listenerStrategy: ListenerStrategyV2, + ctx: newContextLoader("", nil), + } + disableStreamBackoff(t, listener.actionStreamCore()) + return listener +} diff --git a/pkg/client/reconnecting_stream.go b/pkg/client/reconnecting_stream.go new file mode 100644 index 0000000000..5de2a44567 --- /dev/null +++ b/pkg/client/reconnecting_stream.go @@ -0,0 +1,265 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "errors" + "fmt" + "sync" + + "github.com/rs/zerolog" + "golang.org/x/sync/singleflight" + + "github.com/hatchet-dev/hatchet/pkg/client/retry" +) + +var errStreamNotConnected = errors.New("client is not connected") + +// reconnectingStream owns one logical grpc stream that survives reconnects. +// All mutable state sits behind a single mutex; connect attempts are +// single-shot, coalesced through one singleflight group, and never hold the +// mutex (snapshots must not block behind an in-flight connect). +// The lifecycle context is cancelled only by Close. +// NOTE: field order follows govet fieldalignment (enforced by the pre-commit +// autofixer); mu guards client, generation, hasClient, and closed. +type reconnectingStream[C any] struct { + lifecycleCtx context.Context + client C + connectGroup singleflight.Group + + // sleep waits out the backoff delay for the given attempt. Tests inject a + // no-op; production uses retry.SleepStreamBackoff. + sleep func(ctx context.Context, attempt int) error + + constructor func(context.Context) (C, error) + lifecycleCancel context.CancelFunc + replay func(context.Context, C) error + closeSend func(C) error + l *zerolog.Logger + + // name identifies the stream in log messages ("workflow run listener", …). + name string + + generation uint64 + mu sync.Mutex + hasClient bool + closed bool +} + +func newReconnectingStream[C any]( + l *zerolog.Logger, + name string, + constructor func(context.Context) (C, error), + closeSend func(C) error, + replay func(context.Context, C) error, +) *reconnectingStream[C] { + return newReconnectingStreamWithLifecycle(context.Background(), l, name, constructor, closeSend, replay) +} + +func newReconnectingStreamWithLifecycle[C any]( + parent context.Context, + l *zerolog.Logger, + name string, + constructor func(context.Context) (C, error), + closeSend func(C) error, + replay func(context.Context, C) error, +) *reconnectingStream[C] { + lifecycleCtx, lifecycleCancel := context.WithCancel(parent) // nolint: gosec // lifecycleCancel is stored on the struct and called by Close + + return &reconnectingStream[C]{ + constructor: constructor, + closeSend: closeSend, + replay: replay, + sleep: retry.SleepStreamBackoff, + lifecycleCtx: lifecycleCtx, + lifecycleCancel: lifecycleCancel, + name: name, + l: l, + } +} + +func (s *reconnectingStream[C]) snapshot() (client C, generation uint64, ok bool) { + s.mu.Lock() + defer s.mu.Unlock() + return s.client, s.generation, s.hasClient +} + +func (s *reconnectingStream[C]) setInitialClient(client C) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.hasClient { + return + } + + s.client = client + s.hasClient = true +} + +func (s *reconnectingStream[C]) installClient(client C) error { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return errListenerClosed + } + + oldClient := s.client + hadOldClient := s.hasClient + s.client = client + s.hasClient = true + s.generation++ + s.mu.Unlock() + + if hadOldClient && s.closeSend != nil { + if err := s.closeSend(oldClient); err != nil { + s.l.Warn().Err(err).Str("stream", s.name).Msg("failed to close replaced stream client") + } + } + + return nil +} + +func (s *reconnectingStream[C]) isClosed() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.closed +} + +func (s *reconnectingStream[C]) lifecycleContext() context.Context { + return s.lifecycleCtx +} + +func (s *reconnectingStream[C]) connectOnce(ctx context.Context) error { + _, err, _ := s.connectGroup.Do("connect", func() (interface{}, error) { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return nil, errListenerClosed + } + s.mu.Unlock() + + client, err := s.constructor(ctx) + if err != nil { + return nil, err + } + + if s.replay != nil { + if err := s.replay(ctx, client); err != nil { + if s.closeSend != nil { + _ = s.closeSend(client) + } + return nil, err + } + } + + if err := s.installClient(client); err != nil { + if s.closeSend != nil { + _ = s.closeSend(client) + } + return nil, err + } + + return nil, nil + }) + return err +} + +func (s *reconnectingStream[C]) connectSync(ctx context.Context) error { + if s.isClosed() { + return errListenerClosed + } + + if err := ctx.Err(); err != nil { + return err + } + + var lastErr error + for attempt := 0; attempt < retry.StreamSyncMaxAttempts; attempt++ { + if attempt > 0 { + if err := s.sleep(ctx, attempt-1); err != nil { + return err + } + } + + if s.isClosed() { + return errListenerClosed + } + + if err := ctx.Err(); err != nil { + return err + } + + err := s.connectOnce(ctx) + if err == nil { + return nil + } + + lastErr = err + + if errors.Is(err, errListenerClosed) || retry.ClassifyStreamError(ctx, err) == retry.StreamDecisionStop { + return err + } + + s.l.Error().Err(err).Str("stream", s.name).Int("attempt", attempt+1). + Msg("stream connect attempt failed") + } + + return fmt.Errorf("could not connect to %s after %d attempts: %w", s.name, retry.StreamSyncMaxAttempts, lastErr) +} + +func (s *reconnectingStream[C]) retrySend(ctx context.Context, send func(C) error) error { + var lastErr error + for attempt := 0; attempt < retry.StreamSyncMaxAttempts; attempt++ { + client, gen, ok := s.snapshot() + if !ok { + return errStreamNotConnected + } + + err := send(client) + if err == nil { + return nil + } + + lastErr = err + s.l.Warn().Err(err).Str("stream", s.name).Int("attempt", attempt+1).Msg("stream send failed") + + if _, genAfter, _ := s.snapshot(); genAfter != gen { + continue + } + + if rerr := s.connectSync(ctx); rerr != nil { + if errors.Is(rerr, errListenerClosed) || retry.ClassifyStreamError(ctx, rerr) == retry.StreamDecisionStop { + return fmt.Errorf("could not reconnect %s to retry send: %w", s.name, rerr) + } + s.l.Error().Err(rerr).Str("stream", s.name).Msg("stream reconnect after send failure failed") + } + + if serr := s.sleep(ctx, attempt); serr != nil { + return serr + } + } + + return fmt.Errorf("could not send to %s after %d attempts: %w", s.name, retry.StreamSyncMaxAttempts, lastErr) +} + +func (s *reconnectingStream[C]) closeStream() error { + client, _, ok := s.snapshot() + if !ok || s.closeSend == nil { + return nil + } + + return s.closeSend(client) +} + +func (s *reconnectingStream[C]) Close() error { + s.mu.Lock() + s.closed = true + s.mu.Unlock() + + if s.lifecycleCancel != nil { + s.lifecycleCancel() + } + + return s.closeStream() +} diff --git a/pkg/client/reconnecting_stream_test.go b/pkg/client/reconnecting_stream_test.go new file mode 100644 index 0000000000..c7256b59b0 --- /dev/null +++ b/pkg/client/reconnecting_stream_test.go @@ -0,0 +1,384 @@ +package client + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" + "github.com/hatchet-dev/hatchet/pkg/client/retry" +) + +func TestRetrySend_ResubscribesOnSendFailure(t *testing.T) { + constructorCalls := atomic.Int32{} + + failingClient := &mockSubscribeClient{ + sendErr: fmt.Errorf("stream broken"), + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + workingClient := &mockSubscribeClient{ + sendErr: nil, + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, failingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + calls := constructorCalls.Add(1) + if calls == 1 { + return failingClient, nil + } + return workingClient, nil + }) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "test-workflow-run-id"}) + }) + + require.NoError(t, err) + assert.GreaterOrEqual(t, constructorCalls.Load(), int32(1)) + assert.GreaterOrEqual(t, failingClient.sendCount.Load(), int32(1)) + assert.Equal(t, int32(1), workingClient.sendCount.Load()) +} + +func TestReconnectingStreamSnapshotNotBlockedByConnect(t *testing.T) { + constructorEntered := make(chan struct{}) + releaseConstructor := make(chan struct{}) + var closeConstructorEntered sync.Once + + oldClient := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + newClient := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, oldClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + closeConstructorEntered.Do(func() { + close(constructorEntered) + }) + <-releaseConstructor + return newClient, nil + }) + + retryErr := make(chan error, 1) + go func() { + retryErr <- stream.connectOnce(context.Background()) + }() + + select { + case <-constructorEntered: + case <-time.After(time.Second): + t.Fatal("reconnect did not reach constructor") + } + + snapshotRead := make(chan struct { + client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient + generation uint64 + }, 1) + go func() { + client, generation, _ := stream.snapshot() + snapshotRead <- struct { + client dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient + generation uint64 + }{client: client, generation: generation} + }() + + select { + case snapshot := <-snapshotRead: + require.Same(t, oldClient, snapshot.client) + require.Equal(t, uint64(0), snapshot.generation) + case <-time.After(100 * time.Millisecond): + t.Fatal("snapshot blocked behind reconnect") + } + + close(releaseConstructor) + + require.NoError(t, <-retryErr) + require.NoError(t, stream.Close()) +} + +func TestRetrySend_FailsAfterMaxRetries(t *testing.T) { + failingClient := &mockSubscribeClient{ + sendErr: fmt.Errorf("stream broken"), + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, failingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return failingClient, nil + }) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "test-workflow-run-id"}) + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "could not send to") + assert.Equal(t, int32(retry.StreamSyncMaxAttempts), failingClient.sendCount.Load()) +} + +func TestRetrySend_SucceedsOnFirstAttempt(t *testing.T) { + workingClient := &mockSubscribeClient{ + sendErr: nil, + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + constructorCalls := atomic.Int32{} + + stream := newTestWorkflowStream(t, workingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return workingClient, nil + }) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "test-workflow-run-id"}) + }) + + require.NoError(t, err) + assert.Equal(t, int32(1), workingClient.sendCount.Load()) + assert.Equal(t, int32(0), constructorCalls.Load()) +} + +func TestRetrySend_HandlesNilClient(t *testing.T) { + stream := newTestWorkflowStream(t, nil, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return nil, fmt.Errorf("connection failed") + }) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "test-workflow-run-id"}) + }) + + require.Error(t, err) + assert.ErrorIs(t, err, errStreamNotConnected) +} + +func TestRetrySend_ConcurrentSafety(t *testing.T) { + workingClient := &mockSubscribeClient{ + sendErr: nil, + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + sendFn: func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + time.Sleep(10 * time.Millisecond) + return nil + }, + } + + stream := newTestWorkflowStream(t, workingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return workingClient, nil + }) + + var wg sync.WaitGroup + numGoroutines := 10 + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: fmt.Sprintf("workflow-run-%d", id)}) + }) + assert.NoError(t, err) + }(i) + } + + wg.Wait() + assert.Equal(t, int32(numGoroutines), workingClient.sendCount.Load()) +} + +func TestRetrySubscribe_SingleflightCoalescesConcurrentCalls(t *testing.T) { + constructorCalls := atomic.Int32{} + + client := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, client, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + time.Sleep(50 * time.Millisecond) + return client, nil + }) + + numGoroutines := 10 + var wg sync.WaitGroup + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := stream.connectSync(context.Background()) + assert.NoError(t, err) + }() + } + + wg.Wait() + assert.Equal(t, int32(1), constructorCalls.Load()) +} + +func TestRetrySubscribe_GenerationIncrements(t *testing.T) { + client := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, client, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }) + + _, gen0, _ := stream.snapshot() + assert.Equal(t, uint64(0), gen0) + + require.NoError(t, stream.connectSync(context.Background())) + + _, gen1, _ := stream.snapshot() + assert.Equal(t, uint64(1), gen1) + + require.NoError(t, stream.connectSync(context.Background())) + + _, gen2, _ := stream.snapshot() + assert.Equal(t, uint64(2), gen2) +} + +func TestGetClientSnapshot_ReturnsCurrentClient(t *testing.T) { + client1 := &mockSubscribeClient{recvChan: make(chan *dispatchercontracts.WorkflowRunEvent)} + + stream := newTestWorkflowStream(t, client1, nil) + + got, gen, ok := stream.snapshot() + assert.True(t, ok) + assert.Equal(t, client1, got) + assert.Equal(t, uint64(0), gen) +} + +func TestRetrySubscribeSyncStopsAtStreamSyncMaxAttempts(t *testing.T) { + constructorCalls := atomic.Int32{} + + stream := newTestWorkflowStream(t, nil, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return nil, status.Error(codes.Unavailable, "still down") + }) + + err := stream.connectSync(context.Background()) + require.Error(t, err) + assert.Equal(t, int32(retry.StreamSyncMaxAttempts), constructorCalls.Load()) +} + +func TestRetrySendStaleGenerationSkipsReconnect(t *testing.T) { + constructorCalls := atomic.Int32{} + + workingClient := &mockSubscribeClient{recvChan: make(chan *dispatchercontracts.WorkflowRunEvent)} + failingClient := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, failingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return workingClient, nil + }) + + failingClient.sendFn = func(req *dispatchercontracts.SubscribeToWorkflowRunsRequest) error { + require.NoError(t, stream.installClient(workingClient)) + return status.Error(codes.Unavailable, "send failed") + } + + require.NoError(t, stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "run-1"}) + })) + assert.Equal(t, int32(0), constructorCalls.Load()) +} + +func TestReconnectingStreamConnectsAreCoalesced(t *testing.T) { + releaseConstructor := make(chan struct{}) + var concurrentConnect atomic.Int32 + var maxConcurrent atomic.Int32 + + client := &mockSubscribeClient{ + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, client, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + current := concurrentConnect.Add(1) + for { + prev := maxConcurrent.Load() + if current <= prev { + break + } + if maxConcurrent.CompareAndSwap(prev, current) { + break + } + } + defer concurrentConnect.Add(-1) + + <-releaseConstructor + return client, nil + }) + + syncDone := make(chan error, 1) + backgroundDone := make(chan error, 1) + + go func() { + syncDone <- stream.connectSync(context.Background()) + }() + go func() { + backgroundDone <- stream.connectOnce(context.Background()) + }() + + require.Eventually(t, func() bool { + return concurrentConnect.Load() >= 1 + }, time.Second, 10*time.Millisecond) + + assert.Equal(t, int32(1), concurrentConnect.Load()) + assert.Equal(t, int32(1), maxConcurrent.Load()) + + close(releaseConstructor) + + require.NoError(t, <-syncDone) + require.NoError(t, <-backgroundDone) +} + +func TestRetrySendShortCircuitsOnPermanentReconnectError(t *testing.T) { + t.Run("unauthenticated", func(t *testing.T) { + failingClient := &mockSubscribeClient{ + sendErr: status.Error(codes.Unavailable, "send failed"), + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + constructorCalls := atomic.Int32{} + stream := newTestWorkflowStream(t, failingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return nil, status.Error(codes.Unauthenticated, "auth failed") + }) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "run-1"}) + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "could not reconnect") + assert.Less(t, failingClient.sendCount.Load(), int32(retry.StreamSyncMaxAttempts)) + }) + + t.Run("closed listener", func(t *testing.T) { + failingClient := &mockSubscribeClient{ + sendErr: status.Error(codes.Unavailable, "send failed"), + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + stream := newTestWorkflowStream(t, failingClient, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return nil, fmt.Errorf("should not be called") + }) + + require.NoError(t, stream.Close()) + + err := stream.retrySend(context.Background(), func(c dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient) error { + return c.Send(&dispatchercontracts.SubscribeToWorkflowRunsRequest{WorkflowRunId: "run-1"}) + }) + + require.Error(t, err) + assert.ErrorIs(t, err, errListenerClosed) + }) +} diff --git a/pkg/client/retry/stream.go b/pkg/client/retry/stream.go new file mode 100644 index 0000000000..3ce1fcbec7 --- /dev/null +++ b/pkg/client/retry/stream.go @@ -0,0 +1,82 @@ +package retry + +import ( + "context" + "errors" + "io" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// StreamDecision describes how a stream error should be handled. +type StreamDecision int + +const ( + // StreamDecisionRetry indicates the stream should reconnect after a transient failure. + StreamDecisionRetry StreamDecision = iota + // StreamDecisionStop indicates the stream should exit without reconnecting. + StreamDecisionStop + // StreamDecisionNoProgress indicates reconnect would not advance the stream. + StreamDecisionNoProgress +) + +const ( + // StreamBaseDelay is the base delay for stream reconnect backoff. + StreamBaseDelay = 1 * time.Second + // StreamMaxDelay caps stream reconnect backoff delay. + StreamMaxDelay = 30 * time.Second + // StreamBackoffFactor is the exponential multiplier for stream reconnect backoff. + StreamBackoffFactor = 2.0 + // StreamSyncMaxAttempts bounds synchronous subscribe/send reconnect attempts. + StreamSyncMaxAttempts = 5 +) + +// StreamBackoffDelay returns a full-jitter delay for the given stream reconnect attempt. +func StreamBackoffDelay(attempt int) time.Duration { + clk := defaultClock() + return fullJitterDelay(attempt, StreamBaseDelay, StreamBackoffFactor, StreamMaxDelay, clk.jitter) +} + +// Sleep waits for d or until ctx is cancelled. +func Sleep(ctx context.Context, d time.Duration) error { + return sleepContext(ctx, d) +} + +// SleepStreamBackoff waits for the stream reconnect backoff delay or until ctx is cancelled. +func SleepStreamBackoff(ctx context.Context, attempt int) error { + return Sleep(ctx, StreamBackoffDelay(attempt)) +} + +// ClassifyStreamError maps a stream error to a reconnect decision. +func ClassifyStreamError(ctx context.Context, err error) StreamDecision { + if err == nil { + return StreamDecisionStop + } + + if ctx.Err() != nil { + return StreamDecisionStop + } + + if errors.Is(err, io.EOF) { + return StreamDecisionNoProgress + } + + st, ok := status.FromError(err) + if !ok { + return StreamDecisionNoProgress + } + + switch st.Code() { + case codes.Canceled: + return StreamDecisionStop + case codes.Unauthenticated, codes.PermissionDenied, codes.InvalidArgument, + codes.FailedPrecondition, codes.NotFound, codes.Unimplemented: + return StreamDecisionStop + case codes.Unavailable, codes.Internal, codes.DeadlineExceeded, codes.ResourceExhausted: + return StreamDecisionRetry + default: + return StreamDecisionNoProgress + } +} diff --git a/pkg/client/retry/stream_test.go b/pkg/client/retry/stream_test.go new file mode 100644 index 0000000000..5721571fb8 --- /dev/null +++ b/pkg/client/retry/stream_test.go @@ -0,0 +1,107 @@ +package retry + +import ( + "context" + "fmt" + "io" + "math/rand/v2" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestStreamBackoffDelayFullJitter(t *testing.T) { + t.Parallel() + + r := rand.New(rand.NewPCG(1, 1)) + delay := fullJitterDelay(3, StreamBaseDelay, StreamBackoffFactor, StreamMaxDelay, r.Int64N) + assert.GreaterOrEqual(t, delay, time.Duration(0)) + assert.LessOrEqual(t, delay, StreamMaxDelay) +} + +func TestStreamBackoffDelayCap(t *testing.T) { + t.Parallel() + + r := rand.New(rand.NewPCG(2, 2)) + delay := fullJitterDelay(100, StreamBaseDelay, StreamBackoffFactor, StreamMaxDelay, r.Int64N) + assert.LessOrEqual(t, delay, StreamMaxDelay) +} + +func TestSleepCancellation(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := Sleep(ctx, time.Second) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestSleepStreamBackoffCancellation(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := SleepStreamBackoff(ctx, 0) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestClassifyStreamErrorTransientCodes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + codesToRetry := []codes.Code{ + codes.Unavailable, + codes.Internal, + codes.DeadlineExceeded, + codes.ResourceExhausted, + } + + for _, code := range codesToRetry { + err := status.Error(code, "transient") + assert.Equal(t, StreamDecisionRetry, ClassifyStreamError(ctx, err), "code %s", code) + } +} + +func TestClassifyStreamErrorPermanentCodes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + codesToStop := []codes.Code{ + codes.Canceled, + codes.Unauthenticated, + codes.PermissionDenied, + codes.InvalidArgument, + codes.FailedPrecondition, + codes.NotFound, + codes.Unimplemented, + } + + for _, code := range codesToStop { + err := status.Error(code, "permanent") + assert.Equal(t, StreamDecisionStop, ClassifyStreamError(ctx, err), "code %s", code) + } +} + +func TestClassifyStreamErrorEOFAndUnknown(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + assert.Equal(t, StreamDecisionNoProgress, ClassifyStreamError(ctx, io.EOF)) + assert.Equal(t, StreamDecisionNoProgress, ClassifyStreamError(ctx, fmt.Errorf("plain error"))) +} + +func TestClassifyStreamErrorContextCancelled(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := status.Error(codes.Unavailable, "transient") + assert.Equal(t, StreamDecisionStop, ClassifyStreamError(ctx, err)) +} diff --git a/pkg/client/stream_classify.go b/pkg/client/stream_classify.go new file mode 100644 index 0000000000..d8eaddc598 --- /dev/null +++ b/pkg/client/stream_classify.go @@ -0,0 +1,81 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "errors" + "io" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/hatchet-dev/hatchet/pkg/client/retry" +) + +// maxConsecutiveStreamNoProgress caps consecutive no-progress failures +// (recv or reconnect) before a listen loop surfaces the error to its owner. +const maxConsecutiveStreamNoProgress = 10 + +var errListenerClosed = errors.New("listener is closed") + +// streamVerdict is the single classification of an error observed on a +// long-lived stream, for both recv and reconnect failures. +type streamVerdict int + +const ( + // verdictRetry reconnects without counting toward the no-progress cap. + verdictRetry streamVerdict = iota + // verdictNoProgress reconnects and counts toward maxConsecutiveStreamNoProgress. + verdictNoProgress + // verdictStopClean exits the listen loop returning nil. + verdictStopClean + // verdictStopError exits the listen loop returning the error. + verdictStopError +) + +// streamClassifier maps a stream error to a verdict. The listen loop consults +// it once per error, so implementations may carry per-error side effects +// (e.g. the action listener's V2→V1 strategy fallback). +type streamClassifier func(ctx context.Context, err error) streamVerdict + +// newStreamClassifier builds the default classifier. reconnectOnEOF is +// consulted every time io.EOF is observed because handler registration can +// change between errors. +func newStreamClassifier(reconnectOnEOF func(ctx context.Context) bool) streamClassifier { + return func(ctx context.Context, err error) streamVerdict { + switch { + case ctx.Err() != nil, + errors.Is(err, errListenerClosed), + errors.Is(err, context.Canceled), + status.Code(err) == codes.Canceled: + return verdictStopClean + case errors.Is(err, io.EOF): + if reconnectOnEOF(ctx) { + return verdictRetry + } + return verdictStopClean + } + switch retry.ClassifyStreamError(ctx, err) { + case retry.StreamDecisionRetry: + return verdictRetry + case retry.StreamDecisionStop: + return verdictStopError + default: + return verdictNoProgress + } + } +} + +// shouldLogReconnectMilestone rate-limits reconnect warnings: first attempt +// and every fifth thereafter. +func shouldLogReconnectMilestone(attempt int) bool { + return attempt == 1 || attempt%5 == 0 +} + +func sendListenerError(ctx context.Context, errCh chan<- error, err error) { + select { + case errCh <- err: + case <-ctx.Done(): + } +} diff --git a/pkg/client/stream_classify_test.go b/pkg/client/stream_classify_test.go new file mode 100644 index 0000000000..6d215db01a --- /dev/null +++ b/pkg/client/stream_classify_test.go @@ -0,0 +1,133 @@ +package client + +import ( + "context" + "errors" + "fmt" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestNewStreamClassifierTransientCodes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + classify := newStreamClassifier(func(context.Context) bool { return false }) + + codesToRetry := []codes.Code{ + codes.Unavailable, + codes.Internal, + codes.DeadlineExceeded, + codes.ResourceExhausted, + } + + for _, code := range codesToRetry { + err := status.Error(code, "transient") + assert.Equal(t, verdictRetry, classify(ctx, err), "code %s", code) + } +} + +func TestNewStreamClassifierPermanentCodes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + classify := newStreamClassifier(func(context.Context) bool { return false }) + + codesToStop := []codes.Code{ + codes.Unauthenticated, + codes.PermissionDenied, + codes.InvalidArgument, + codes.FailedPrecondition, + codes.NotFound, + codes.Unimplemented, + } + + for _, code := range codesToStop { + err := status.Error(code, "permanent") + assert.Equal(t, verdictStopError, classify(ctx, err), "code %s", code) + } +} + +func TestNewStreamClassifierCleanStopConditions(t *testing.T) { + t.Parallel() + + classify := newStreamClassifier(func(context.Context) bool { return false }) + + t.Run("cancelled context", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + assert.Equal(t, verdictStopClean, classify(ctx, status.Error(codes.Unavailable, "transient"))) + }) + + t.Run("errListenerClosed", func(t *testing.T) { + assert.Equal(t, verdictStopClean, classify(context.Background(), errListenerClosed)) + }) + + t.Run("context.Canceled", func(t *testing.T) { + assert.Equal(t, verdictStopClean, classify(context.Background(), context.Canceled)) + }) + + t.Run("grpc Canceled", func(t *testing.T) { + assert.Equal(t, verdictStopClean, classify(context.Background(), status.Error(codes.Canceled, "cancelled"))) + }) +} + +func TestNewStreamClassifierEOFCallback(t *testing.T) { + t.Parallel() + + ctx := context.Background() + calls := 0 + classify := newStreamClassifier(func(context.Context) bool { + calls++ + return calls%2 == 1 + }) + + assert.Equal(t, verdictRetry, classify(ctx, io.EOF)) + assert.Equal(t, verdictStopClean, classify(ctx, io.EOF)) + assert.Equal(t, 2, calls) +} + +func TestNewStreamClassifierPlainError(t *testing.T) { + t.Parallel() + + classify := newStreamClassifier(func(context.Context) bool { return false }) + assert.Equal(t, verdictNoProgress, classify(context.Background(), fmt.Errorf("plain error"))) +} + +func TestNoProgressFatalClassifierWrapper(t *testing.T) { + t.Parallel() + + base := newStreamClassifier(func(context.Context) bool { return false }) + classify := func(ctx context.Context, err error) streamVerdict { + if v := base(ctx, err); v != verdictNoProgress { + return v + } + return verdictStopError + } + + assert.Equal(t, verdictStopError, classify(context.Background(), fmt.Errorf("plain error"))) + assert.Equal(t, verdictRetry, classify(context.Background(), status.Error(codes.Unavailable, "transient"))) +} + +func TestShouldLogReconnectMilestone(t *testing.T) { + t.Parallel() + + assert.True(t, shouldLogReconnectMilestone(1)) + assert.False(t, shouldLogReconnectMilestone(2)) + assert.False(t, shouldLogReconnectMilestone(4)) + assert.True(t, shouldLogReconnectMilestone(5)) + assert.True(t, shouldLogReconnectMilestone(10)) +} + +func TestSendListenerError(t *testing.T) { + t.Parallel() + + errCh := make(chan error, 1) + sendListenerError(context.Background(), errCh, errors.New("boom")) + require.Equal(t, "boom", (<-errCh).Error()) +} diff --git a/pkg/client/stream_listen.go b/pkg/client/stream_listen.go new file mode 100644 index 0000000000..8eac6e52d4 --- /dev/null +++ b/pkg/client/stream_listen.go @@ -0,0 +1,109 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "fmt" +) + +// listenStream runs the receive loop for a reconnecting stream. Reconnects +// use full-jitter backoff; after maxConsecutiveStreamNoProgress consecutive +// no-progress failures the loop returns an error; classify decides clean vs +// error exits. +// +// Context contract: ctx scopes the loop itself (recv classification and +// backoff sleeps; cancellation is a clean exit returning nil, because +// background loops shut down cleanly). Reconnect attempts always use +// stream.lifecycleContext(), which only Close() cancels, so one caller's ctx +// cannot destroy the shared stream's ability to reconnect. +// +// If no client is installed yet, the first iteration connects (this serves +// StreamByAdditionalMetadata's initial connect and makes Listen() usable on a +// never-connected listener). +func listenStream[C any, E any]( + ctx context.Context, + stream *reconnectingStream[C], + recv func(C) (E, error), + handle func(E) error, + classify streamClassifier, +) error { + noProgress := 0 + reconnects := 0 + + client, generation, connected := stream.snapshot() + defer func() { + if stream.closeSend != nil && connected { + if closeErr := stream.closeSend(client); closeErr != nil { + stream.l.Warn().Err(closeErr).Str("stream", stream.name).Msg("failed to close stream after listen exit") + } + } + }() + + for { + if !connected { + if reconnects > 0 { + if err := stream.sleep(ctx, reconnects-1); err != nil { + return nil + } + } + if err := stream.connectOnce(stream.lifecycleContext()); err != nil { + switch classify(ctx, err) { + case verdictStopClean: + return nil + case verdictStopError: + return fmt.Errorf("could not reconnect %s: %w", stream.name, err) + case verdictNoProgress: + noProgress++ + if noProgress >= maxConsecutiveStreamNoProgress { + return fmt.Errorf("%s made no progress after %d consecutive errors: %w", stream.name, noProgress, err) + } + } + reconnects++ + if shouldLogReconnectMilestone(reconnects) { + stream.l.Warn().Err(err).Str("stream", stream.name). + Int("reconnect_attempt", reconnects). + Int("consecutive_no_progress", noProgress). + Str("error_code", streamErrorCode(err)). + Msg("stream reconnect attempt continuing") + } + continue + } + client, generation, _ = stream.snapshot() + if reconnects > 0 { + stream.l.Info().Str("stream", stream.name).Int("attempts", reconnects). + Msg("stream reconnected") + } + connected, noProgress, reconnects = true, 0, 0 + } + + event, err := recv(client) + if err == nil { + noProgress, reconnects = 0, 0 + if herr := handle(event); herr != nil { + return herr + } + continue + } + + switch classify(ctx, err) { + case verdictStopClean: + return nil + case verdictStopError: + return err + case verdictNoProgress: + noProgress++ + if noProgress >= maxConsecutiveStreamNoProgress { + return fmt.Errorf("%s made no progress after %d consecutive errors: %w", stream.name, noProgress, err) + } + } + + if c, g, ok := stream.snapshot(); ok && g != generation { + client, generation = c, g + noProgress, reconnects = 0, 0 + continue + } + reconnects++ + connected = false + } +} diff --git a/pkg/client/stream_listen_test.go b/pkg/client/stream_listen_test.go new file mode 100644 index 0000000000..d79cec1ef8 --- /dev/null +++ b/pkg/client/stream_listen_test.go @@ -0,0 +1,443 @@ +package client + +import ( + "context" + "fmt" + "io" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/hatchet-dev/hatchet/pkg/client/retry" +) + +func TestListenReconnectingStreamHandlesEventsAndStopsOnEOF(t *testing.T) { + recvChan := make(chan testListenEvent, 1) + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + select { + case event, ok := <-recvChan: + if !ok { + return testListenEvent{}, io.EOF + } + return event, nil + default: + return testListenEvent{}, io.EOF + } + }, + } + + var handled atomic.Value + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + return client, nil + }) + + listenErr := make(chan error, 1) + go func() { + listenErr <- listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(event testListenEvent) error { + handled.Store(event.value) + return nil + }, + newStreamClassifier(func(context.Context) bool { return false }), + ) + }() + + recvChan <- testListenEvent{value: "event-1"} + require.Eventually(t, func() bool { + return handled.Load() == "event-1" + }, time.Second, 10*time.Millisecond) + + close(recvChan) + require.NoError(t, <-listenErr) + assert.True(t, client.closeCalled.Load()) +} + +func TestListenReconnectingStreamContextCancellationReturnsNil(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + cancel() + return testListenEvent{}, status.Error(codes.Canceled, "canceled") + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + return client, nil + }) + + err := listenStream(ctx, stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return true }), + ) + require.NoError(t, err) +} + +func TestListenReconnectingStreamPermanentRecvErrorReturnsError(t *testing.T) { + recvErr := status.Error(codes.PermissionDenied, "permission denied") + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, recvErr + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + t.Fatal("constructor should not run after permanent receive error") + return nil, nil + }) + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return true }), + ) + require.ErrorIs(t, err, recvErr) + assert.Equal(t, codes.PermissionDenied, status.Code(err)) +} + +func TestListenReconnectingStreamReconnectsOnEOFWhenPolicyAllows(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + initialClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, io.EOF + }, + } + replacementRecv := make(chan testListenEvent, 1) + replacementClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + event, ok := <-replacementRecv + if !ok { + return testListenEvent{}, io.EOF + } + return event, nil + }, + } + constructorCalls := atomic.Int32{} + + stream := newTestListenStream(t, initialClient, func(ctx context.Context) (*testListenClient, error) { + constructorCalls.Add(1) + return replacementClient, nil + }) + + var handled atomic.Value + listenErr := make(chan error, 1) + go func() { + listenErr <- listenStream(ctx, stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(event testListenEvent) error { + handled.Store(event.value) + cancel() + return nil + }, + newStreamClassifier(func(context.Context) bool { return ctx.Err() == nil }), + ) + }() + + require.Eventually(t, func() bool { + return constructorCalls.Load() == 1 + }, time.Second, 10*time.Millisecond) + + replacementRecv <- testListenEvent{value: "after-reconnect"} + require.Eventually(t, func() bool { + return handled.Load() == "after-reconnect" + }, time.Second, 10*time.Millisecond) + + close(replacementRecv) + require.NoError(t, <-listenErr) +} + +func TestListenReconnectingStreamChecksEOFPolicyEachRecvError(t *testing.T) { + policyCalls := atomic.Int32{} + recvCalls := atomic.Int32{} + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + recvCalls.Add(1) + return testListenEvent{}, io.EOF + }, + } + + constructorCalls := atomic.Int32{} + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + constructorCalls.Add(1) + return client, nil + }) + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { + return policyCalls.Add(1) == 1 + }), + ) + require.NoError(t, err) + assert.Equal(t, int32(2), policyCalls.Load()) + assert.Equal(t, int32(2), recvCalls.Load()) + assert.Equal(t, int32(1), constructorCalls.Load()) +} + +func TestListenReconnectingStreamNoProgressReconnectsBeforeCap(t *testing.T) { + recvCalls := atomic.Int32{} + constructorCalls := atomic.Int32{} + + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + if recvCalls.Add(1) == 1 { + return testListenEvent{}, fmt.Errorf("plain recv error") + } + return testListenEvent{}, io.EOF + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + constructorCalls.Add(1) + return client, nil + }) + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return false }), + ) + require.NoError(t, err) + assert.Equal(t, int32(1), constructorCalls.Load()) + assert.Equal(t, int32(2), recvCalls.Load()) +} + +func TestListenStreamNoProgressStopsAtCap(t *testing.T) { + recvCalls := atomic.Int32{} + constructorCalls := atomic.Int32{} + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + recvCalls.Add(1) + return testListenEvent{}, fmt.Errorf("plain recv error") + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + constructorCalls.Add(1) + return nil, fmt.Errorf("plain connect error") + }) + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return true }), + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "made no progress") + assert.Equal(t, int32(1), recvCalls.Load()) + assert.GreaterOrEqual(t, constructorCalls.Load(), int32(maxConsecutiveStreamNoProgress-2)) +} + +func TestListenStreamNoProgressFatalClassifierStopsImmediately(t *testing.T) { + recvErr := fmt.Errorf("plain recv error") + recvCalls := atomic.Int32{} + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + recvCalls.Add(1) + return testListenEvent{}, recvErr + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + t.Fatal("constructor should not run when no-progress policy stops immediately") + return nil, nil + }) + + base := newStreamClassifier(func(context.Context) bool { return false }) + classify := func(ctx context.Context, err error) streamVerdict { + if v := base(ctx, err); v != verdictNoProgress { + return v + } + return verdictStopError + } + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + classify, + ) + require.ErrorIs(t, err, recvErr) + assert.Equal(t, int32(1), recvCalls.Load()) +} + +func TestListenStreamConnectsUseLifecycleContext(t *testing.T) { + listenCtx := context.Background() + + var constructorCtx atomic.Value + recvCalls := atomic.Int32{} + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + if recvCalls.Add(1) == 1 { + return testListenEvent{}, status.Error(codes.Unavailable, "broken") + } + return testListenEvent{}, io.EOF + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + constructorCtx.Store(ctx) + return client, nil + }) + + err := listenStream(listenCtx, stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { + return recvCalls.Load() == 1 + }), + ) + require.NoError(t, err) + + storedCtx := constructorCtx.Load().(context.Context) + assert.Equal(t, stream.lifecycleContext(), storedCtx) +} + +func TestListenReconnectingStreamGenerationChangeFastPath(t *testing.T) { + releaseReconnect := make(chan struct{}) + initialClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, status.Error(codes.Unavailable, "broken") + }, + } + replacementClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, io.EOF + }, + } + + constructorCalls := atomic.Int32{} + stream := newTestListenStream(t, initialClient, func(ctx context.Context) (*testListenClient, error) { + if constructorCalls.Add(1) == 1 { + <-releaseReconnect + } + return replacementClient, nil + }) + + listenErr := make(chan error, 1) + go func() { + listenErr <- listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return false }), + ) + }() + + require.Eventually(t, func() bool { + return constructorCalls.Load() == 1 + }, time.Second, 10*time.Millisecond) + + require.NoError(t, stream.installClient(replacementClient)) + close(releaseReconnect) + + require.NoError(t, <-listenErr) + assert.Equal(t, int32(1), constructorCalls.Load()) +} + +func TestListenReconnectingStreamClosesListenedClientOnExit(t *testing.T) { + initialClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, io.EOF + }, + } + replacementClient := &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, io.EOF + }, + } + + stream := newTestListenStream(t, initialClient, func(ctx context.Context) (*testListenClient, error) { + return replacementClient, nil + }) + + require.NoError(t, listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return false }), + )) + assert.True(t, initialClient.closeCalled.Load()) + assert.False(t, replacementClient.closeCalled.Load()) +} + +func TestListenStreamReconnectsUnboundedlyOnTransientErrors(t *testing.T) { + constructorCalls := atomic.Int32{} + ctx, cancel := context.WithCancel(context.Background()) + + stream := newTestListenStream(t, nil, func(ctx context.Context) (*testListenClient, error) { + if constructorCalls.Add(1) >= retry.StreamSyncMaxAttempts+1 { + cancel() + } + return nil, status.Error(codes.Unavailable, "still down") + }) + + err := listenStream(ctx, stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return true }), + ) + require.NoError(t, err) + assert.Greater(t, constructorCalls.Load(), int32(retry.StreamSyncMaxAttempts)) +} + +func TestListenStreamRetryableRecvFailureAppliesReconnectBackoff(t *testing.T) { + var sleepAttempts []int + recvCalls := atomic.Int32{} + + client := &testListenClient{ + recvFn: func() (testListenEvent, error) { + if recvCalls.Add(1) == 1 { + return testListenEvent{}, status.Error(codes.Unavailable, "transient") + } + return testListenEvent{}, io.EOF + }, + } + + stream := newTestListenStream(t, client, func(ctx context.Context) (*testListenClient, error) { + return client, nil + }) + stream.sleep = func(_ context.Context, attempt int) error { + sleepAttempts = append(sleepAttempts, attempt) + return nil + } + + err := listenStream(context.Background(), stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return false }), + ) + require.NoError(t, err) + assert.Equal(t, []int{0}, sleepAttempts) +} + +func TestListenStreamSleepCancellationReturnsNil(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + stream := newTestListenStream(t, &testListenClient{ + recvFn: func() (testListenEvent, error) { + return testListenEvent{}, status.Error(codes.Unavailable, "stream broken") + }, + }, func(ctx context.Context) (*testListenClient, error) { + return nil, status.Error(codes.Unavailable, "still down") + }) + stream.sleep = func(sleepCtx context.Context, _ int) error { + cancel() + return sleepCtx.Err() + } + + err := listenStream(ctx, stream, + func(c *testListenClient) (testListenEvent, error) { return c.Recv() }, + func(testListenEvent) error { return nil }, + newStreamClassifier(func(context.Context) bool { return true }), + ) + require.NoError(t, err) +} diff --git a/pkg/client/workflow.go b/pkg/client/workflow.go index a2c602cced..54e184f475 100644 --- a/pkg/client/workflow.go +++ b/pkg/client/workflow.go @@ -5,7 +5,6 @@ package client import ( "encoding/json" "fmt" - "time" "github.com/google/uuid" ) @@ -90,54 +89,41 @@ func (r *WorkflowResult) Results() (interface{}, error) { // Result waits for the workflow run to complete and returns the results. // -// Retry strategy (best-effort): -// 1. This function retries AddWorkflowRun up to DefaultActionListenerRetryCount times with DefaultActionListenerRetryInterval intervals -// 2. AddWorkflowRun calls retrySend which retries up to DefaultActionListenerRetryCount times with DefaultActionListenerRetryInterval intervals -// 3. Each retrySend attempt calls retrySubscribe which itself retries up to DefaultActionListenerRetryCount times with DefaultActionListenerRetryInterval intervals -func (c *Workflow) Result() (*WorkflowResult, error) { +// AddWorkflowRun is attempted once; it uses bounded synchronous reconnect +// (StreamSyncMaxAttempts) for send and subscribe failures. The background +// listen loop reconnects unboundedly while the listener remains open. +func (r *Workflow) Result() (*WorkflowResult, error) { resChan := make(chan *WorkflowResult, 1) + failChan := make(chan error, 1) sessionId := uuid.NewString() - var err error - retries := 0 - - for retries < DefaultActionListenerRetryCount { - if retries > 0 { - time.Sleep(DefaultActionListenerRetryInterval) - } - - err = c.listener.AddWorkflowRun( - c.workflowRunId, - sessionId, - func(event WorkflowRunEvent) error { - resChan <- &WorkflowResult{ - workflowRun: event, - } - - return nil - }, - ) - - if err == nil { - defer c.listener.RemoveWorkflowRun(c.workflowRunId, sessionId) - - break - } - - retries++ - } - - if retries == DefaultActionListenerRetryCount && err != nil { + err := r.listener.addWorkflowRun(r.workflowRunId, sessionId, + func(event WorkflowRunEvent) error { + resChan <- &WorkflowResult{workflowRun: event} + return nil + }, + func(err error) { + select { + case failChan <- err: + default: + } + }, + ) + if err != nil { + r.listener.RemoveWorkflowRun(r.workflowRunId, sessionId) return nil, fmt.Errorf("failed to listen for workflow events: %w", err) } + defer r.listener.RemoveWorkflowRun(r.workflowRunId, sessionId) - res := <-resChan - - for _, stepRunResult := range res.workflowRun.Results { - if stepRunResult.Error != nil { - return nil, fmt.Errorf("%s", *stepRunResult.Error) + select { + case res := <-resChan: + for _, stepRunResult := range res.workflowRun.Results { + if stepRunResult.Error != nil { + return nil, fmt.Errorf("%s", *stepRunResult.Error) + } } + return res, nil + case err := <-failChan: + return nil, fmt.Errorf("workflow run listener terminated while waiting for %s: %w", r.workflowRunId, err) } - - return res, nil } diff --git a/pkg/client/workflow_test.go b/pkg/client/workflow_test.go new file mode 100644 index 0000000000..ba4e71c5f0 --- /dev/null +++ b/pkg/client/workflow_test.go @@ -0,0 +1,78 @@ +// Deprecated: This package is part of the legacy v0 workflow definition system. +// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go +package client + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/rs/zerolog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + dispatchercontracts "github.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts" + "github.com/hatchet-dev/hatchet/pkg/client/retry" +) + +func TestWorkflowResultSingleAddWorkflowRunAttempt(t *testing.T) { + logger := zerolog.Nop() + constructorCalls := atomic.Int32{} + + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + constructorCalls.Add(1) + return nil, status.Error(codes.Unavailable, "engine down") + }, nil) + + workflow := NewWorkflow("run-1", listener) + _, err := workflow.Result() + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to listen for workflow events") + assert.LessOrEqual(t, constructorCalls.Load(), int32(retry.StreamSyncMaxAttempts)) + assert.Greater(t, constructorCalls.Load(), int32(0)) + + assert.False(t, listener.reg.hasAny()) +} + +func TestWorkflowResultFailsWhenListenerDiesPermanently(t *testing.T) { + logger := zerolog.Nop() + runID := "run-permanent-fail" + + client := &mockSubscribeClient{ + recvFn: func() (*dispatchercontracts.WorkflowRunEvent, error) { + return nil, status.Error(codes.Unauthenticated, "auth failed") + }, + recvChan: make(chan *dispatchercontracts.WorkflowRunEvent), + } + + listener := newTestWorkflowRunsListener(t, &logger, func(ctx context.Context) (dispatchercontracts.Dispatcher_SubscribeToWorkflowRunsClient, error) { + return client, nil + }, client) + + workflow := NewWorkflow(runID, listener) + + resultDone := make(chan struct{}) + var resultErr error + go func() { + _, resultErr = workflow.Result() + close(resultDone) + }() + + require.Eventually(t, func() bool { + select { + case <-resultDone: + return true + default: + return false + } + }, 5*time.Second, 10*time.Millisecond) + + require.Error(t, resultErr) + assert.Contains(t, resultErr.Error(), runID) + assert.ErrorIs(t, resultErr, status.Error(codes.Unauthenticated, "auth failed")) + + require.NoError(t, listener.Close()) +}