From 72fddcf5641247a4a7cc94a38f3c88ad82214184 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:32:59 -0700 Subject: [PATCH 1/7] isolate worker completions on dedicated gRPC connections A DTS server allows 100 concurrent streams per HTTP/2 connection, and the long-lived GetWorkItems stream holds one. With one connection per worker, completions and abandons for a busy worker queue behind that limit. Workers now keep one intake connection (Hello, GetWorkItems) and route an exact allowlist of Complete*/Abandon* methods round-robin over a bounded group of owned completion connections (default 3, 1-8 via WithWorkerCompletionConnections). Work is completed on the connection generation that accepted it, and a retired generation closes only after its in-flight work drains. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 2 + README.md | 4 +- client/grpc_worker.go | 196 +++- client/grpc_worker_compat.go | 5 +- client/grpc_worker_connection.go | 186 ++++ client/grpc_worker_connection_test.go | 927 ++++++++++++++++++ client/grpc_worker_drain_test.go | 166 ++++ client/grpc_worker_processor.go | 45 +- client/grpc_worker_transport_test.go | 4 + durabletaskscheduler/README.md | 73 +- durabletaskscheduler/connection_test.go | 9 +- durabletaskscheduler/resource_id_test.go | 4 + .../worker_completion_test.go | 220 +++++ exporthistory/versioning_test.go | 12 + samples/largepayloads/largepayloads_test.go | 14 +- samples/worker/README.md | 5 +- samples/worker/worker.go | 32 +- 17 files changed, 1822 insertions(+), 82 deletions(-) create mode 100644 client/grpc_worker_connection.go create mode 100644 client/grpc_worker_connection_test.go create mode 100644 client/grpc_worker_drain_test.go create mode 100644 durabletaskscheduler/worker_completion_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 871dfa56..a3a53163 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Added required completion-transport isolation to gRPC workers: one intake and three dedicated completion/abandon connections by default, with a configurable 1–8 completion budget and bounded generation-safe retirement. `WithWorkerCompletionConnections(0)` is invalid. Borrowed workers and client listeners require separate caller-owned channels via `WithWorkerCompletionTransports`; owning factories must return distinct connections and non-nil closers for every invocation. - Added `ResourceId` connection-string support for every DTS authentication mode and optional `AuthorityHost` configuration for SDK-created DefaultAzure, WorkloadIdentity, Environment, and InteractiveBrowser credentials. Caller-supplied credentials own their authority; managed identity and developer-tool cloud configuration remain separate. - Added a flat catalogue of directly runnable DTS samples with outcome assertions, owned-resource cleanup, and process-level E2E validation selected by sample name using Go's `-run` flag. - Added the exported `task.CallActivityOption` type, per-activity tag options, completion-action tag propagation, and distinct activity and sub-orchestration action trace contexts. @@ -85,6 +86,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Graceful worker shutdown and run-context cancellation now preserve the intake stream until accepted work and acknowledgements drain, instead of releasing their server leases early. New dispatch is rejected after intake stop; deadline cancellation and genuine lease-loss handling remain unchanged. - Forced coroutine unloading no longer leaks durable actions or custom status from application defers: durable operations panic with `ErrTaskBlocked` before changing state. Normal-return defers remain awaitable across replay, and orchestration loggers suppress forced-unload output. - Long durable timers no longer schedule another chunk when a trailing timer event is processed after the orchestration has finalized. - DTS authentication now caches access tokens and immutable gRPC metadata, honors credential `RefreshOn` guidance, and coalesces concurrent refreshes and failures. Previously credentials such as `AzureCLICredential` were invoked for every RPC, serializing high-throughput workers behind external token acquisition. diff --git a/README.md b/README.md index 9a436781..61aadb69 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ You write the orchestrations, activities, and entities. [Azure Durable Task Sche DTS is the only supported runtime. This SDK does not include a storage backend. -The SDK gives you two connections: +The SDK provides two roles: - A **management client** starts orchestrations and reads their state. - A **worker** runs your registered orchestrators, activities, and entities. @@ -82,7 +82,7 @@ This project needs Go 1.25 or later. ## Connection to DTS -The [`durabletaskscheduler`](./durabletaskscheduler) package is the integration surface. The package validates your connection string and supplies Azure token credentials. It also owns the management connection and the worker connection separately. The worker stream recovers after a network failure. +The [`durabletaskscheduler`](./durabletaskscheduler) package is the integration surface. The package validates your connection string and supplies Azure token credentials. It owns the management connection separately from the worker's connection group. Each worker has one intake connection and three dedicated completion connections by default, while retaining one identity and executor. Use `durabletaskclient.WithWorkerCompletionConnections(n)` to tune the completion budget from 1 to 8; isolation cannot be disabled. The worker stream recovers after a network failure. ```go options, err := durabletaskscheduler.NewOptionsFromConnectionString( diff --git a/client/grpc_worker.go b/client/grpc_worker.go index dd64e9e3..868d9f08 100644 --- a/client/grpc_worker.go +++ b/client/grpc_worker.go @@ -40,7 +40,10 @@ type workItemsStream interface { type grpcWorkerClientFactory func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) // TaskHubGrpcWorkerConnectionFactory creates a gRPC connection for a worker -// stream generation. A non-nil closer transfers ownership to the worker. +// stream generation. Each generation requests one intake connection and the +// configured completion budget. Every call must return a distinct connection +// and a distinct non-nil closer transferring ownership. Factories must honor cancellation; +// closers must return promptly. type TaskHubGrpcWorkerConnectionFactory func(context.Context) (grpc.ClientConnInterface, io.Closer, error) type TaskHubGrpcWorkerOption func(*taskHubGrpcWorkerOptions) error @@ -65,6 +68,10 @@ const ( // orchestration-completion message limit. DefaultMaxOrchestratorCompletionBytes = 4_089_446 minOrchestratorCompletionBytes = 64 * 1024 + + // DefaultWorkerCompletionConnections is the owned worker's completion budget, + // excluding its intake connection. + DefaultWorkerCompletionConnections = 3 ) type WorkItemFilter struct { @@ -82,32 +89,35 @@ type WorkItemFilters struct { } type taskHubGrpcWorkerOptions struct { - maxConcurrentOrchestrations int - maxConcurrentActivities int - maxConcurrentEntities int - helloTimeout time.Duration - silentDisconnectTimeout time.Duration - rpcTimeout time.Duration - reconnectBaseDelay time.Duration - reconnectMaxDelay time.Duration - transientRetryMaxAttempts int - transientRetryBaseDelay time.Duration - transientRetryMaxDelay time.Duration - maxStreamedHistoryEvents int - maxStreamedHistoryBytes int64 - maxOrchestratorCompletionBytes int - maximumTimerInterval *time.Duration - taskExecutorOptions []task.TaskExecutorOption - versioning *task.VersioningOptions - capabilities []WorkerCapability - workItemFilters *WorkItemFilters - workItemFiltersConfigured bool - autoWorkItemFilters bool - largePayloads *api.LargePayloadOptions - converter api.DataConverter - unversionedOrchestrators map[string]struct{} - unversionedActivities map[string]struct{} - reconnectRandom randomInt64N + maxConcurrentOrchestrations int + maxConcurrentActivities int + maxConcurrentEntities int + completionConnections int + completionConnectionsConfigured bool + borrowedCompletionTransports []grpc.ClientConnInterface + helloTimeout time.Duration + silentDisconnectTimeout time.Duration + rpcTimeout time.Duration + reconnectBaseDelay time.Duration + reconnectMaxDelay time.Duration + transientRetryMaxAttempts int + transientRetryBaseDelay time.Duration + transientRetryMaxDelay time.Duration + maxStreamedHistoryEvents int + maxStreamedHistoryBytes int64 + maxOrchestratorCompletionBytes int + maximumTimerInterval *time.Duration + taskExecutorOptions []task.TaskExecutorOption + versioning *task.VersioningOptions + capabilities []WorkerCapability + workItemFilters *WorkItemFilters + workItemFiltersConfigured bool + autoWorkItemFilters bool + largePayloads *api.LargePayloadOptions + converter api.DataConverter + unversionedOrchestrators map[string]struct{} + unversionedActivities map[string]struct{} + reconnectRandom randomInt64N // waitFn overrides every delay the worker imposes on itself: reconnect // backoff, transient RPC retry backoff, and work-item abandon delays. It is // only set by tests so the deterministic delay schedule can be observed @@ -121,6 +131,7 @@ func defaultTaskHubGrpcWorkerOptions() taskHubGrpcWorkerOptions { maxConcurrentOrchestrations: defaultConcurrency, maxConcurrentActivities: defaultConcurrency, maxConcurrentEntities: defaultConcurrency, + completionConnections: DefaultWorkerCompletionConnections, helloTimeout: 30 * time.Second, silentDisconnectTimeout: 2 * time.Minute, rpcTimeout: 30 * time.Second, @@ -168,6 +179,40 @@ func WithMaxConcurrentEntityWorkItems(n int) TaskHubGrpcWorkerOption { } } +// WithWorkerCompletionConnections sets the dedicated completion/abandon budget +// per generation, excluding intake. The default is 3; valid values are 1 through +// 8. Isolation cannot be disabled. For a borrowed worker, an explicitly set +// budget must match the transports supplied by WithWorkerCompletionTransports. +func WithWorkerCompletionConnections(n int) TaskHubGrpcWorkerOption { + return func(options *taskHubGrpcWorkerOptions) error { + if err := validateWorkerCompletionBudget(n); err != nil { + return err + } + options.completionConnections = n + options.completionConnectionsConfigured = true + return nil + } +} + +// WithWorkerCompletionTransports supplies a borrowed worker's dedicated +// completion connections. They must be distinct from intake and one another, +// use equivalent authentication/worker metadata, and number between 1 and 8. +// The caller retains ownership and is responsible for readiness and recovery. +// This option is not valid with an owning connection factory. +func WithWorkerCompletionTransports(connections ...grpc.ClientConnInterface) TaskHubGrpcWorkerOption { + snapshot := slices.Clone(connections) + return func(options *taskHubGrpcWorkerOptions) error { + if err := validateWorkerCompletionBudget(len(snapshot)); err != nil { + return err + } + if err := validateWorkerCompletionTransports(nil, snapshot); err != nil { + return err + } + options.borrowedCompletionTransports = slices.Clone(snapshot) + return nil + } +} + // WithMaxStreamedHistoryEvents limits the number of history events retained // while receiving a streamed orchestration history. func WithMaxStreamedHistoryEvents(n int) TaskHubGrpcWorkerOption { @@ -492,6 +537,7 @@ type grpcWorkerRun struct { orchestrationSlots chan struct{} activitySlots chan struct{} entitySlots chan struct{} + dispatchMu sync.Mutex pending sync.WaitGroup retired sync.WaitGroup err error @@ -506,11 +552,12 @@ type grpcWorkerConnection struct { pending sync.WaitGroup } -// NewTaskHubGrpcWorker creates a worker that borrows a caller-owned connection. +// NewTaskHubGrpcWorker borrows intake and requires separate caller-owned +// completion connections supplied by WithWorkerCompletionTransports. // -// The worker never closes or replaces the supplied connection: a reconnect only -// recreates the work-item stream on that same connection. Recovery is therefore -// limited to what the caller's connection can do on its own. A connection that +// The worker never closes or replaces the supplied connections: reconnect only +// recreates the stream on the same intake. Recovery is limited to what the +// caller's connections can do on their own. A connection that // is permanently wedged (for example one whose credentials expired, or one whose // endpoint moved) keeps producing poisoned streams and the worker keeps retrying // on the escalating reconnect schedule without ever obtaining a fresh channel. @@ -522,17 +569,34 @@ func NewTaskHubGrpcWorker( logger api.Logger, opts ...TaskHubGrpcWorkerOption, ) (*TaskHubGrpcWorker, error) { - if cc == nil { + if isNilWorkerTransport(cc) { return nil, fmt.Errorf("gRPC connection is required") } - return newTaskHubGrpcWorker( + var completions []grpc.ClientConnInterface + worker, err := newTaskHubGrpcWorker( func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { - return protos.NewTaskHubSidecarServiceClient(cc), nil, nil + group := &workerCompletionTransport{intake: cc, completions: completions} + return protos.NewTaskHubSidecarServiceClient(group), nil, nil }, registry, logger, opts..., ) + if err != nil { + return nil, err + } + completions = worker.options.borrowedCompletionTransports + if len(completions) == 0 { + return nil, fmt.Errorf("borrowed workers require separate connections via WithWorkerCompletionTransports") + } + if err := validateWorkerCompletionTransports(cc, completions); err != nil { + return nil, err + } + if worker.options.completionConnectionsConfigured && worker.options.completionConnections != len(completions) { + return nil, fmt.Errorf("worker completion budget must match the number of borrowed completion transports") + } + worker.options.completionConnections = len(completions) + return worker, nil } // NewTaskHubGrpcWorkerWithConnectionFactory creates a worker whose connection @@ -547,24 +611,29 @@ func NewTaskHubGrpcWorkerWithConnectionFactory( if factory == nil { return nil, fmt.Errorf("gRPC worker connection factory is required") } - return newTaskHubGrpcWorker( + var worker *TaskHubGrpcWorker + var err error + worker, err = newTaskHubGrpcWorker( func(ctx context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { - cc, closer, err := factory(ctx) + ctx, cancel := context.WithTimeout(ctx, worker.options.helloTimeout) + defer cancel() + group, err := newWorkerCompletionTransport(ctx, factory, worker.options.completionConnections) if err != nil { return nil, nil, err } - if cc == nil { - if closer != nil { - _ = closer.Close() - } - return nil, nil, fmt.Errorf("gRPC worker connection factory returned a nil connection") - } - return protos.NewTaskHubSidecarServiceClient(cc), closer, nil + return protos.NewTaskHubSidecarServiceClient(group), group, nil }, registry, logger, opts..., ) + if err != nil { + return nil, err + } + if len(worker.options.borrowedCompletionTransports) != 0 { + return nil, fmt.Errorf("owning worker factories cannot use WithWorkerCompletionTransports") + } + return worker, nil } func newTaskHubGrpcWorker( @@ -881,7 +950,8 @@ func (w *TaskHubGrpcWorker) Run(ctx context.Context) error { return run.err } -// Shutdown stops intake and waits for in-flight work to finish. If ctx expires, +// Shutdown stops new dispatch and keeps the lease-owning intake stream open +// until accepted execution and acknowledgements finish. If ctx expires, // in-flight execution and completion RPCs are canceled. func (w *TaskHubGrpcWorker) Shutdown(ctx context.Context) error { w.mu.Lock() @@ -1003,9 +1073,10 @@ func (w *TaskHubGrpcWorker) runLoop(run *grpcWorkerRun, connection *grpcWorkerCo w.options.reconnectMaxDelay, w.options.reconnectRandom, ) + var previousRetirement <-chan struct{} for { observedMessage, err := w.consumeConnection(run, connection) - w.retireConnection(run, connection) + retirement := w.retireConnection(run, connection) if run.intakeCtx.Err() != nil { return nil @@ -1013,6 +1084,16 @@ func (w *TaskHubGrpcWorker) runLoop(run *grpcWorkerRun, connection *grpcWorkerCo if !isTransientWorkerError(err) { return fmt.Errorf("work item stream stopped with a non-retryable error: %w", err) } + // Allow one draining generation alongside the next intake, but do not + // accumulate connection groups during repeated reconnects. + if previousRetirement != nil { + select { + case <-previousRetirement: + case <-run.intakeCtx.Done(): + return nil + } + } + previousRetirement = retirement if code := status.Code(err); code == codes.Unauthenticated || code == codes.PermissionDenied { w.logger.Warnf("gRPC worker authentication or authorization failed; reconnecting so refreshed credentials or RBAC can recover: %v", err) } @@ -1059,21 +1140,28 @@ func (w *TaskHubGrpcWorker) connect(ctx context.Context) (*grpcWorkerConnection, if err != nil { return nil, fmt.Errorf("failed to create gRPC worker connection: %w", err) } - closeOnError := func() { + closeOnError := func(err error) error { if closer != nil { - _ = closer.Close() + if closeErr := closer.Close(); closeErr != nil { + return errors.Join(err, fmt.Errorf("failed to close gRPC worker connection: %w", closeErr)) + } } + return err } helloCtx, cancelHello := context.WithTimeout(ctx, w.options.helloTimeout) _, err = client.Hello(helloCtx, &emptypb.Empty{}) cancelHello() if err != nil { - closeOnError() - return nil, fmt.Errorf("gRPC worker Hello failed: %w", err) + return nil, closeOnError(fmt.Errorf("gRPC worker Hello failed: %w", err)) } - streamCtx, cancelStream := context.WithCancel(ctx) + if err := ctx.Err(); err != nil { + return nil, closeOnError(fmt.Errorf("worker intake stopped before opening the stream: %w", status.FromContextError(err).Err())) + } + // Ending the stream releases server-side leases, so logical intake stop + // must not cancel it until accepted work has acknowledged or abandoned. + streamCtx, cancelStream := context.WithCancel(context.WithoutCancel(ctx)) stream, err := client.GetWorkItems(streamCtx, &protos.GetWorkItemsRequest{ MaxConcurrentOrchestrationWorkItems: int32(w.options.maxConcurrentOrchestrations), MaxConcurrentActivityWorkItems: int32(w.options.maxConcurrentActivities), @@ -1083,8 +1171,7 @@ func (w *TaskHubGrpcWorker) connect(ctx context.Context) (*grpcWorkerConnection, }) if err != nil { cancelStream() - closeOnError() - return nil, fmt.Errorf("failed to open gRPC work item stream: %w", err) + return nil, closeOnError(fmt.Errorf("failed to open gRPC work item stream: %w", err)) } return &grpcWorkerConnection{ @@ -1247,11 +1334,13 @@ func workItemFiltersToProto(filters *WorkItemFilters) *protos.WorkItemFilters { return result } -func (w *TaskHubGrpcWorker) retireConnection(run *grpcWorkerRun, connection *grpcWorkerConnection) { +func (w *TaskHubGrpcWorker) retireConnection(run *grpcWorkerRun, connection *grpcWorkerConnection) <-chan struct{} { connection.cancelStream() + done := make(chan struct{}) run.retired.Add(1) go func() { defer run.retired.Done() + defer close(done) connection.pending.Wait() if connection.closer != nil { if err := connection.closer.Close(); err != nil { @@ -1259,6 +1348,7 @@ func (w *TaskHubGrpcWorker) retireConnection(run *grpcWorkerRun, connection *grp } } }() + return done } func isTransientWorkerError(err error) bool { diff --git a/client/grpc_worker_compat.go b/client/grpc_worker_compat.go index 600f51df..2f0572d0 100644 --- a/client/grpc_worker_compat.go +++ b/client/grpc_worker_compat.go @@ -7,8 +7,9 @@ import ( "github.com/microsoft/durabletask-go/task" ) -// StartWorkItemListener preserves the original client API. New applications -// should create a dedicated TaskHubGrpcWorker, ideally with its own connection. +// StartWorkItemListener starts a worker borrowing the client's intake +// connection. WithWorkerCompletionTransports must supply separate caller-owned +// completion channels. Prefer a dedicated owned worker for DTS applications. // // The listener borrows the client's connection, so it inherits the connection // ownership limitations documented on [NewTaskHubGrpcWorker]. diff --git a/client/grpc_worker_connection.go b/client/grpc_worker_connection.go new file mode 100644 index 00000000..f325a4f7 --- /dev/null +++ b/client/grpc_worker_connection.go @@ -0,0 +1,186 @@ +package client + +import ( + "context" + "errors" + "fmt" + "io" + "reflect" + "slices" + "sync" + "sync/atomic" + + "github.com/microsoft/durabletask-go/internal/protos" + "google.golang.org/grpc" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/status" +) + +const maxWorkerCompletionConnections = 8 + +func validateWorkerCompletionBudget(n int) error { + if n < 1 || n > maxWorkerCompletionConnections { + return fmt.Errorf("worker completion connections must be between 1 and %d", maxWorkerCompletionConnections) + } + return nil +} + +func isNilWorkerTransport(value any) bool { + if value == nil { + return true + } + v := reflect.ValueOf(value) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return v.IsNil() + default: + return false + } +} + +func sameWorkerTransport(a, b any) bool { + return a != nil && reflect.ValueOf(a).Comparable() && a == b +} + +func validateWorkerCompletionTransports(intake grpc.ClientConnInterface, completions []grpc.ClientConnInterface) error { + for i, connection := range completions { + if isNilWorkerTransport(connection) { + return fmt.Errorf("worker completion transport %d is nil", i) + } + if sameWorkerTransport(connection, intake) { + return fmt.Errorf("worker completion transports must be separate from intake") + } + for _, previous := range completions[:i] { + if sameWorkerTransport(connection, previous) { + return fmt.Errorf("worker completion transports must be distinct") + } + } + } + return nil +} + +// workerCompletionTransport belongs to one stream generation. Its caller must +// drain that generation's work before closing it. +type workerCompletionTransport struct { + intake grpc.ClientConnInterface + completions []grpc.ClientConnInterface + closers []io.Closer + next atomic.Uint64 + closeOnce sync.Once + closeErr error +} + +func newWorkerCompletionTransport( + ctx context.Context, + factory TaskHubGrpcWorkerConnectionFactory, + additional int, +) (*workerCompletionTransport, error) { + if err := validateWorkerCompletionBudget(additional); err != nil { + return nil, err + } + group := &workerCompletionTransport{} + for i := 0; i <= additional; i++ { + cc, closer, err := factory(ctx) + duplicateCloser := false + if !isNilWorkerTransport(closer) { + for _, previous := range group.closers { + duplicateCloser = duplicateCloser || sameWorkerTransport(closer, previous) + } + if !duplicateCloser { + group.closers = append(group.closers, closer) + } + } + switch { + case err != nil: + err = fmt.Errorf("worker connection %d: %w", i, err) + case isNilWorkerTransport(cc): + err = fmt.Errorf("gRPC worker connection factory returned a nil connection") + case isNilWorkerTransport(closer): + err = fmt.Errorf("completion connections require an owning factory with a closer for every connection") + case ctx.Err() != nil: + err = ctx.Err() + } + if err == nil { + err = validateWorkerCompletionTransports(group.intake, append(slices.Clone(group.completions), cc)) + } + if err == nil && duplicateCloser { + err = fmt.Errorf("owning worker factories must return distinct ownership closers") + } + if err == nil { + if connection, ok := cc.(*grpc.ClientConn); ok { + err = waitWorkerConnectionReady(ctx, connection) + } + } + if err != nil { + return nil, errors.Join(err, group.Close()) + } + if i == 0 { + group.intake = cc + } else { + group.completions = append(group.completions, cc) + } + } + return group, nil +} + +func waitWorkerConnectionReady(ctx context.Context, connection *grpc.ClientConn) error { + connection.Connect() + for { + if err := ctx.Err(); err != nil { + return fmt.Errorf("worker connection did not become ready: %w", status.FromContextError(err).Err()) + } + state := connection.GetState() + if state == connectivity.Ready { + return nil + } + connection.WaitForStateChange(ctx, state) + } +} + +func (c *workerCompletionTransport) Invoke( + ctx context.Context, + method string, + args, reply any, + opts ...grpc.CallOption, +) error { + connection := c.intake + if isWorkerCompletionMethod(method) { + index := (c.next.Add(1) - 1) % uint64(len(c.completions)) + connection = c.completions[index] + } + return connection.Invoke(ctx, method, args, reply, opts...) +} + +func isWorkerCompletionMethod(method string) bool { + switch method { + case protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName, + protos.TaskHubSidecarService_CompleteOrchestratorTask_FullMethodName, + protos.TaskHubSidecarService_CompleteEntityTask_FullMethodName, + protos.TaskHubSidecarService_AbandonTaskActivityWorkItem_FullMethodName, + protos.TaskHubSidecarService_AbandonTaskOrchestratorWorkItem_FullMethodName, + protos.TaskHubSidecarService_AbandonTaskEntityWorkItem_FullMethodName: + return true + default: + return false + } +} + +func (c *workerCompletionTransport) NewStream( + ctx context.Context, + desc *grpc.StreamDesc, + method string, + opts ...grpc.CallOption, +) (grpc.ClientStream, error) { + return c.intake.NewStream(ctx, desc, method, opts...) +} + +func (c *workerCompletionTransport) Close() error { + c.closeOnce.Do(func() { + for i := len(c.closers) - 1; i >= 0; i-- { + if err := c.closers[i].Close(); err != nil { + c.closeErr = errors.Join(c.closeErr, fmt.Errorf("close worker connection %d: %w", i, err)) + } + } + }) + return c.closeErr +} diff --git a/client/grpc_worker_connection_test.go b/client/grpc_worker_connection_test.go new file mode 100644 index 00000000..338c637a --- /dev/null +++ b/client/grpc_worker_connection_test.go @@ -0,0 +1,927 @@ +package client + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/microsoft/durabletask-go/api" + "github.com/microsoft/durabletask-go/internal/protos" + "github.com/microsoft/durabletask-go/task" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/stats" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/wrapperspb" +) + +type completionTestCall struct { + method string + transport int + token string +} + +type completionTestStream struct { + request *protos.GetWorkItemsRequest + transport int + commands chan fakeWorkItemResult +} + +type completionTestServer struct { + protos.UnimplementedTaskHubSidecarServiceServer + streams chan *completionTestStream + calls chan completionTestCall + complete func(context.Context, *protos.ActivityResponse) error + hellos atomic.Int32 +} + +func newCompletionTestServer() *completionTestServer { + return &completionTestServer{ + streams: make(chan *completionTestStream, 16), + calls: make(chan completionTestCall, 2048), + } +} + +func completionTransportID(ctx context.Context) int { + md, _ := metadata.FromIncomingContext(ctx) + id, _ := strconv.Atoi(md.Get("test-transport")[0]) + return id +} + +func (s *completionTestServer) Hello(context.Context, *emptypb.Empty) (*emptypb.Empty, error) { + s.hellos.Add(1) + return &emptypb.Empty{}, nil +} + +func (s *completionTestServer) GetWorkItems(request *protos.GetWorkItemsRequest, stream protos.TaskHubSidecarService_GetWorkItemsServer) error { + intake := &completionTestStream{ + request: request, + transport: completionTransportID(stream.Context()), + commands: make(chan fakeWorkItemResult, 1024), + } + select { + case s.streams <- intake: + case <-stream.Context().Done(): + return stream.Context().Err() + } + for { + select { + case <-stream.Context().Done(): + return stream.Context().Err() + case command := <-intake.commands: + if command.err != nil { + return command.err + } + if err := stream.Send(command.item); err != nil { + return err + } + } + } +} + +func (s *completionTestServer) CompleteActivityTask(ctx context.Context, request *protos.ActivityResponse) (*protos.CompleteTaskResponse, error) { + s.calls <- completionTestCall{ + method: protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName, + transport: completionTransportID(ctx), + token: request.CompletionToken, + } + if s.complete != nil { + if err := s.complete(ctx, request); err != nil { + return nil, err + } + } + return &protos.CompleteTaskResponse{}, nil +} + +func (s *completionTestServer) CompleteOrchestratorTask(context.Context, *protos.OrchestratorResponse) (*protos.CompleteTaskResponse, error) { + return &protos.CompleteTaskResponse{}, nil +} + +func (s *completionTestServer) CompleteEntityTask(context.Context, *protos.EntityBatchResult) (*protos.CompleteTaskResponse, error) { + return &protos.CompleteTaskResponse{}, nil +} + +func (s *completionTestServer) AbandonTaskActivityWorkItem(ctx context.Context, request *protos.AbandonActivityTaskRequest) (*protos.AbandonActivityTaskResponse, error) { + s.calls <- completionTestCall{ + method: protos.TaskHubSidecarService_AbandonTaskActivityWorkItem_FullMethodName, + transport: completionTransportID(ctx), + token: request.CompletionToken, + } + return &protos.AbandonActivityTaskResponse{}, nil +} + +func (s *completionTestServer) AbandonTaskOrchestratorWorkItem(context.Context, *protos.AbandonOrchestrationTaskRequest) (*protos.AbandonOrchestrationTaskResponse, error) { + return &protos.AbandonOrchestrationTaskResponse{}, nil +} + +func (s *completionTestServer) AbandonTaskEntityWorkItem(context.Context, *protos.AbandonEntityTaskRequest) (*protos.AbandonEntityTaskResponse, error) { + return &protos.AbandonEntityTaskResponse{}, nil +} + +func (s *completionTestServer) GetInstance(context.Context, *protos.GetInstanceRequest) (*protos.GetInstanceResponse, error) { + return &protos.GetInstanceResponse{}, nil +} + +func (s *completionTestServer) StreamInstanceHistory(_ *protos.StreamInstanceHistoryRequest, stream protos.TaskHubSidecarService_StreamInstanceHistoryServer) error { + s.calls <- completionTestCall{ + method: protos.TaskHubSidecarService_StreamInstanceHistory_FullMethodName, + transport: completionTransportID(stream.Context()), + } + return nil +} + +type completionMethodKey struct{} + +type completionTestStats struct { + begin chan struct{} + headers atomic.Int32 +} + +func (s *completionTestStats) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context { + return context.WithValue(ctx, completionMethodKey{}, info.FullMethodName) +} + +func (s *completionTestStats) HandleRPC(ctx context.Context, event stats.RPCStats) { + if ctx.Value(completionMethodKey{}) != protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName { + return + } + switch event.(type) { + case *stats.Begin: + s.begin <- struct{}{} + case *stats.OutHeader: + s.headers.Add(1) + } +} + +func (*completionTestStats) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context { + return ctx +} + +func (*completionTestStats) HandleConn(context.Context, stats.ConnStats) {} + +type completionTestCloser struct { + connection *grpc.ClientConn + closes atomic.Int32 + closed chan struct{} +} + +func (c *completionTestCloser) Close() error { + c.closes.Add(1) + err := c.connection.Close() + close(c.closed) + return err +} + +type completionTestHarness struct { + listener *bufconn.Listener + mu sync.Mutex + closers []*completionTestCloser + stats *completionTestStats +} + +func newCompletionTestHarness(t *testing.T, server *completionTestServer, options ...grpc.ServerOption) *completionTestHarness { + t.Helper() + harness := &completionTestHarness{ + listener: bufconn.Listen(1024 * 1024), + stats: &completionTestStats{begin: make(chan struct{}, 2048)}, + } + grpcServer := grpc.NewServer(append([]grpc.ServerOption{grpc.MaxConcurrentStreams(100)}, options...)...) + protos.RegisterTaskHubSidecarServiceServer(grpcServer, server) + done := make(chan struct{}) + go func() { + defer close(done) + _ = grpcServer.Serve(harness.listener) + }() + t.Cleanup(func() { + harness.mu.Lock() + defer harness.mu.Unlock() + for _, closer := range harness.closers { + if closer.closes.Load() == 0 { + require.NoError(t, closer.Close()) + } + } + grpcServer.Stop() + require.NoError(t, harness.listener.Close()) + awaitCompletionTest(t, done) + }) + return harness +} + +func (h *completionTestHarness) factory(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + h.mu.Lock() + defer h.mu.Unlock() + id := strconv.Itoa(len(h.closers)) + connection, err := grpc.NewClient("passthrough:///completion-test", + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return h.listener.Dial() }), + grpc.WithStatsHandler(h.stats), + grpc.WithUnaryInterceptor(func(ctx context.Context, method string, request, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + return invoker(metadata.AppendToOutgoingContext(ctx, "test-transport", id, "workerid", "one-worker"), method, request, reply, cc, opts...) + }), + grpc.WithStreamInterceptor(func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return streamer(metadata.AppendToOutgoingContext(ctx, "test-transport", id, "workerid", "one-worker"), desc, cc, method, opts...) + }), + ) + if err != nil { + return nil, nil, err + } + closer := &completionTestCloser{connection: connection, closed: make(chan struct{})} + h.closers = append(h.closers, closer) + return connection, closer, nil +} + +func (h *completionTestHarness) connections() []*completionTestCloser { + h.mu.Lock() + defer h.mu.Unlock() + return append([]*completionTestCloser(nil), h.closers...) +} + +func awaitCompletionTest[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case value := <-ch: + return value + case <-time.After(10 * time.Second): + t.Fatal("completion transport test did not make progress") + var zero T + return zero + } +} + +func completionTestActivity(token string) fakeWorkItemResult { + return fakeWorkItemResult{item: &protos.WorkItem{ + CompletionToken: token, + Request: &protos.WorkItem_ActivityRequest{ActivityRequest: &protos.ActivityRequest{ + Name: "echo", + OrchestrationInstance: &protos.OrchestrationInstance{InstanceId: token}, + }}, + }} +} + +func startCompletionTestWorker(t *testing.T, harness *completionTestHarness, additional, activityLimit int) (*TaskHubGrpcWorker, <-chan struct{}) { + t.Helper() + executed := make(chan struct{}, 1024) + registry := task.NewTaskRegistry() + require.NoError(t, registry.AddActivityN("echo", func(task.ActivityContext) (any, error) { + executed <- struct{}{} + return "done", nil + })) + options := []TaskHubGrpcWorkerOption{ + WithMaxConcurrentOrchestrationWorkItems(1000), + WithMaxConcurrentActivityWorkItems(activityLimit), + WithWorkerHelloTimeout(5 * time.Second), + WithWorkerSilentDisconnectTimeout(time.Minute), + WithWorkerRPCTimeout(time.Minute), + WithWorkerReconnectBackoff(time.Millisecond, time.Millisecond), + WithWorkerTransientRetryPolicy(3, time.Millisecond, time.Millisecond), + } + var worker *TaskHubGrpcWorker + var err error + if additional == 0 { + // Shared transport exists only in this private diagnostic baseline. + worker, err = newTaskHubGrpcWorker(func(ctx context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + cc, closer, err := harness.factory(ctx) + return protos.NewTaskHubSidecarServiceClient(cc), closer, err + }, registry, nil, options...) + } else { + if additional != DefaultWorkerCompletionConnections { + options = append(options, WithWorkerCompletionConnections(additional)) + } + worker, err = NewTaskHubGrpcWorkerWithConnectionFactory(harness.factory, registry, nil, options...) + } + require.NoError(t, err) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = worker.Shutdown(ctx) + require.NoError(t, worker.Wait(ctx)) + }) + require.NoError(t, worker.Start(context.Background())) + return worker, executed +} + +func TestWorkerCompletionConnectionsRelieveHTTP2StreamAdmission(t *testing.T) { + for _, additional := range []int{0, 3} { + t.Run(fmt.Sprintf("additional_%d", additional), func(t *testing.T) { + const tasks = 240 + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + server := newCompletionTestServer() + server.complete = func(ctx context.Context, _ *protos.ActivityResponse) error { + select { + case <-release: + return nil + case <-ctx.Done(): + return ctx.Err() + } + } + harness := newCompletionTestHarness(t, server) + worker, executed := startCompletionTestWorker(t, harness, additional, 1000) + intake := awaitCompletionTest(t, server.streams) + require.Zero(t, intake.transport) + require.EqualValues(t, 1000, intake.request.MaxConcurrentOrchestrationWorkItems) + require.EqualValues(t, 1000, intake.request.MaxConcurrentActivityWorkItems) + for i := range tasks { + intake.commands <- completionTestActivity(strconv.Itoa(i)) + } + for range tasks { + awaitCompletionTest(t, executed) + awaitCompletionTest(t, harness.stats.begin) + } + admitted := tasks + if additional == 0 { + admitted = 99 // The held intake stream uses the hundredth HTTP/2 slot. + } + byTransport := make(map[int]int) + for range admitted { + call := awaitCompletionTest(t, server.calls) + byTransport[call.transport]++ + } + require.EqualValues(t, admitted, harness.stats.headers.Load()) + require.EqualValues(t, 1, server.hellos.Load(), "additional channels must not register workers") + require.Empty(t, server.streams, "only one intake may be open") + require.Len(t, harness.connections(), additional+1) + if additional == 0 { + require.Equal(t, map[int]int{0: 99}, byTransport) + require.Equal(t, 141, tasks-int(harness.stats.headers.Load()), "calls must wait before OutHeader on stock transport") + } else { + require.Equal(t, map[int]int{1: 80, 2: 80, 3: 80}, byTransport) + } + releaseOnce.Do(func() { close(release) }) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + require.NoError(t, worker.Shutdown(ctx)) + for range tasks - admitted { + awaitCompletionTest(t, server.calls) + } + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } + }) + } +} + +func TestWorkerCompletionConnectionsKeepPendingAcknowledgementsOnOriginalGeneration(t *testing.T) { + for _, failCompletion := range []bool{false, true} { + t.Run(fmt.Sprintf("abandon_%t", failCompletion), func(t *testing.T) { + testWorkerCompletionGeneration(t, failCompletion) + }) + } +} + +func testWorkerCompletionGeneration(t *testing.T, failCompletion bool) { + t.Helper() + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + server := newCompletionTestServer() + server.complete = func(ctx context.Context, request *protos.ActivityResponse) error { + if request.CompletionToken == "old" { + select { + case <-release: + case <-ctx.Done(): + return ctx.Err() + } + if failCompletion { + return status.Error(codes.InvalidArgument, "test failed acknowledgement") + } + } + return nil + } + harness := newCompletionTestHarness(t, server) + worker, _ := startCompletionTestWorker(t, harness, 3, 1000) + first := awaitCompletionTest(t, server.streams) + first.commands <- completionTestActivity("old") + oldCall := awaitCompletionTest(t, server.calls) + require.Equal(t, 1, oldCall.transport) + first.commands <- fakeWorkItemResult{err: status.Error(codes.Unavailable, "reconnect")} + second := awaitCompletionTest(t, server.streams) + require.Equal(t, 4, second.transport) + for _, closer := range harness.connections()[:4] { + require.Zero(t, closer.closes.Load(), "a pending acknowledgement keeps the entire original group alive") + } + second.commands <- completionTestActivity("new") + newCall := awaitCompletionTest(t, server.calls) + require.Equal(t, "new", newCall.token) + require.Equal(t, 5, newCall.transport) + releaseOnce.Do(func() { close(release) }) + if failCompletion { + abandon := awaitCompletionTest(t, server.calls) + require.Equal(t, completionTestCall{ + method: protos.TaskHubSidecarService_AbandonTaskActivityWorkItem_FullMethodName, + transport: 2, + token: "old", + }, abandon, "old work must also abandon on its original group") + } + for _, closer := range harness.connections()[:4] { + awaitCompletionTest(t, closer.closed) + require.EqualValues(t, 1, closer.closes.Load()) + } + require.NoError(t, worker.Shutdown(context.Background())) + require.EqualValues(t, 2, server.hellos.Load()) + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } +} + +func TestWorkerCompletionConnectionsBoundRetiringGenerations(t *testing.T) { + for _, cancelDrain := range []bool{false, true} { + t.Run(fmt.Sprintf("cancel_%t", cancelDrain), func(t *testing.T) { + testWorkerCompletionRetirementBound(t, cancelDrain) + }) + } +} + +func testWorkerCompletionRetirementBound(t *testing.T, cancelDrain bool) { + t.Helper() + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + server := newCompletionTestServer() + server.complete = func(ctx context.Context, _ *protos.ActivityResponse) error { + select { + case <-release: + return nil + case <-ctx.Done(): + return ctx.Err() + } + } + harness := newCompletionTestHarness(t, server) + worker, _ := startCompletionTestWorker(t, harness, 3, 1000) + first := awaitCompletionTest(t, server.streams) + first.commands <- completionTestActivity("held") + awaitCompletionTest(t, server.calls) + first.commands <- fakeWorkItemResult{err: status.Error(codes.Unavailable, "first disconnect")} + second := awaitCompletionTest(t, server.streams) + second.commands <- fakeWorkItemResult{err: status.Error(codes.Unavailable, "second disconnect")} + for _, closer := range harness.connections()[4:8] { + awaitCompletionTest(t, closer.closed) + } + require.Len(t, harness.connections(), 8, "must not allocate a third group while the oldest is draining") + if cancelDrain { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + require.ErrorIs(t, worker.Shutdown(ctx), context.Canceled) + waitCtx, cancelWait := context.WithTimeout(context.Background(), 10*time.Second) + defer cancelWait() + require.NoError(t, worker.Wait(waitCtx)) + require.Len(t, harness.connections(), 8, "shutdown must unblock retirement without opening another group") + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } + return + } + releaseOnce.Do(func() { close(release) }) + third := awaitCompletionTest(t, server.streams) + require.Equal(t, 8, third.transport) + require.Len(t, harness.connections(), 12) + require.NoError(t, worker.Shutdown(context.Background())) + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } +} + +func TestWorkerCompletionConnectionsCancelSaturatedShutdown(t *testing.T) { + const tasks = 450 + server := newCompletionTestServer() + server.complete = func(ctx context.Context, _ *protos.ActivityResponse) error { + <-ctx.Done() + return ctx.Err() + } + harness := newCompletionTestHarness(t, server) + worker, executed := startCompletionTestWorker(t, harness, 3, 1000) + intake := awaitCompletionTest(t, server.streams) + for i := range tasks { + intake.commands <- completionTestActivity(strconv.Itoa(i)) + } + for range tasks { + awaitCompletionTest(t, executed) + awaitCompletionTest(t, harness.stats.begin) + } + for range 300 { + awaitCompletionTest(t, server.calls) + } + require.EqualValues(t, 300, harness.stats.headers.Load(), "each of three completion transports is saturated at 100 streams") + worker.mu.Lock() + run := worker.run + worker.mu.Unlock() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + require.ErrorIs(t, worker.Shutdown(ctx), context.Canceled) + waitCtx, cancelWait := context.WithTimeout(context.Background(), 10*time.Second) + defer cancelWait() + require.NoError(t, worker.Wait(waitCtx)) + require.Empty(t, run.activitySlots, "canceled RPCs must release execution slots") + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } +} + +func TestWorkerCompletionConnectionsNotFoundReleasesSlotWithoutAbandon(t *testing.T) { + server := newCompletionTestServer() + server.complete = func(_ context.Context, request *protos.ActivityResponse) error { + if request.CompletionToken == "gone" { + return status.Error(codes.NotFound, "lease expired") + } + return nil + } + harness := newCompletionTestHarness(t, server) + worker, executed := startCompletionTestWorker(t, harness, 3, 1) + intake := awaitCompletionTest(t, server.streams) + intake.commands <- completionTestActivity("gone") + intake.commands <- completionTestActivity("next") + for _, token := range []string{"gone", "next"} { + awaitCompletionTest(t, executed) + call := awaitCompletionTest(t, server.calls) + require.Equal(t, token, call.token) + require.Equal(t, protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName, call.method) + } + require.NoError(t, worker.Shutdown(context.Background())) + require.Empty(t, server.calls, "terminal NotFound must not cause retry or abandon") +} + +func TestWorkerCompletionTransportRoutesExactMethodsAndPreservesCallOptions(t *testing.T) { + observed := make(chan completionTestCall, 16) + server := newCompletionTestServer() + harness := newCompletionTestHarness(t, server, grpc.UnaryInterceptor( + func(ctx context.Context, request any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + md, _ := metadata.FromIncomingContext(ctx) + if md.Get("workerid")[0] != "one-worker" || md.Get("test-context")[0] != "preserved" { + return nil, status.Error(codes.InvalidArgument, "metadata lost") + } + observed <- completionTestCall{method: info.FullMethod, transport: completionTransportID(ctx)} + if err := grpc.SendHeader(ctx, metadata.Pairs("test-response", "header")); err != nil { + return nil, err + } + grpc.SetTrailer(ctx, metadata.Pairs("test-response", "trailer")) + return handler(ctx, request) + })) + ctx := metadata.AppendToOutgoingContext(context.Background(), "test-context", "preserved") + group, err := newWorkerCompletionTransport(ctx, harness.factory, 3) + require.NoError(t, err) + defer func() { require.NoError(t, group.Close()) }() + client := protos.NewTaskHubSidecarServiceClient(group) + calls := []struct { + method string + call func(...grpc.CallOption) error + want int + }{ + {protos.TaskHubSidecarService_Hello_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.Hello(ctx, &emptypb.Empty{}, opts...) + return err + }, 0}, + {protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.CompleteActivityTask(ctx, &protos.ActivityResponse{}, opts...) + return err + }, 1}, + {protos.TaskHubSidecarService_CompleteOrchestratorTask_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.CompleteOrchestratorTask(ctx, &protos.OrchestratorResponse{}, opts...) + return err + }, 2}, + {protos.TaskHubSidecarService_CompleteEntityTask_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.CompleteEntityTask(ctx, &protos.EntityBatchResult{}, opts...) + return err + }, 3}, + {protos.TaskHubSidecarService_AbandonTaskActivityWorkItem_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.AbandonTaskActivityWorkItem(ctx, &protos.AbandonActivityTaskRequest{}, opts...) + return err + }, 1}, + {protos.TaskHubSidecarService_AbandonTaskOrchestratorWorkItem_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.AbandonTaskOrchestratorWorkItem(ctx, &protos.AbandonOrchestrationTaskRequest{}, opts...) + return err + }, 2}, + {protos.TaskHubSidecarService_AbandonTaskEntityWorkItem_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.AbandonTaskEntityWorkItem(ctx, &protos.AbandonEntityTaskRequest{}, opts...) + return err + }, 3}, + {protos.TaskHubSidecarService_GetInstance_FullMethodName, func(opts ...grpc.CallOption) error { + _, err := client.GetInstance(ctx, &protos.GetInstanceRequest{}, opts...) + return err + }, 0}, + } + for _, test := range calls { + var header, trailer metadata.MD + require.NoError(t, test.call(grpc.Header(&header), grpc.Trailer(&trailer))) + require.Equal(t, completionTestCall{method: test.method, transport: test.want}, awaitCompletionTest(t, observed)) + require.Equal(t, []string{"header"}, header.Get("test-response")) + require.Equal(t, []string{"trailer"}, trailer.Get("test-response")) + if test.method == protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName || + test.method == protos.TaskHubSidecarService_AbandonTaskActivityWorkItem_FullMethodName { + require.Equal(t, test.method, awaitCompletionTest(t, server.calls).method) + } + } + require.False(t, isWorkerCompletionMethod("/TaskHubSidecarService/CompleteActivityTaskExtra")) + require.False(t, isWorkerCompletionMethod("/OtherService/CompleteActivityTask")) + require.False(t, isWorkerCompletionMethod("/TaskHubSidecarService/AbandonAnything")) + require.False(t, isWorkerCompletionMethod(protos.TaskHubSidecarService_GetWorkItems_FullMethodName)) + stream, err := client.GetWorkItems(ctx, &protos.GetWorkItemsRequest{}) + require.NoError(t, err) + require.Zero(t, awaitCompletionTest(t, server.streams).transport) + require.NoError(t, stream.CloseSend()) + history, err := client.StreamInstanceHistory(ctx, &protos.StreamInstanceHistoryRequest{}) + require.NoError(t, err) + _, err = history.Recv() + require.ErrorIs(t, err, io.EOF) + require.Equal(t, completionTestCall{ + method: protos.TaskHubSidecarService_StreamInstanceHistory_FullMethodName, + transport: 0, + }, awaitCompletionTest(t, server.calls)) + _, err = client.CompleteActivityTask(ctx, &protos.ActivityResponse{Result: wrapperspb.String("too large")}, grpc.MaxCallSendMsgSize(1)) + require.Equal(t, codes.ResourceExhausted, status.Code(err), "per-call size limits must not be discarded") +} + +func TestWorkerCompletionTransportRollbackAndOwnership(t *testing.T) { + for _, failure := range []string{"factory error", "nil connection", "borrowed connection", "canceled"} { + t.Run(failure, func(t *testing.T) { + var calls int + closers := []*countingCloser{{}, {}, {}} + borrowed := &fakeClientConn{} + factoryErr := errors.New("completion factory failed") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + group, err := newWorkerCompletionTransport(ctx, func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + index := calls + calls++ + if index == 2 { + switch failure { + case "factory error": + return borrowed, closers[index], factoryErr + case "nil connection": + return nil, closers[index], nil + case "borrowed connection": + return borrowed, nil, nil + case "canceled": + cancel() + } + } + return &fakeClientConn{}, closers[index], nil + }, 3) + require.Error(t, err) + require.Nil(t, group) + require.Equal(t, 3, calls) + require.Zero(t, borrowed.closes.Load(), "ownership must not be inferred from the connection type") + for i, closer := range closers { + if i == 2 && failure == "borrowed connection" { + require.Zero(t, closer.closes.Load()) + } else { + require.EqualValues(t, 1, closer.closes.Load()) + } + } + if failure == "factory error" { + require.ErrorIs(t, err, factoryErr) + } + if failure == "canceled" { + require.ErrorIs(t, err, context.Canceled) + } + }) + } +} + +func TestWorkerCompletionConnectionsReadinessCancellationRollsBack(t *testing.T) { + server := newCompletionTestServer() + harness := newCompletionTestHarness(t, server) + dialing := make(chan struct{}) + var dialOnce sync.Once + var calls int + var failedCloser *completionTestCloser + factory := func(ctx context.Context) (grpc.ClientConnInterface, io.Closer, error) { + calls++ + if calls == 1 { + return harness.factory(ctx) + } + connection, err := grpc.NewClient("passthrough:///unreachable", + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { + dialOnce.Do(func() { close(dialing) }) + <-ctx.Done() + return nil, ctx.Err() + })) + if err != nil { + return nil, nil, err + } + failedCloser = &completionTestCloser{connection: connection, closed: make(chan struct{})} + return connection, failedCloser, nil + } + worker, err := NewTaskHubGrpcWorkerWithConnectionFactory(factory, task.NewTaskRegistry(), nil, WithWorkerCompletionConnections(3)) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + go func() { result <- worker.Start(ctx) }() + awaitCompletionTest(t, dialing) + cancel() + err = awaitCompletionTest(t, result) + require.Equal(t, codes.Canceled, status.Code(err)) + require.ErrorContains(t, err, "did not become ready") + require.Equal(t, 2, calls, "a failing channel must not silently fall back or continue allocation") + require.False(t, worker.Running()) + require.EqualValues(t, 1, failedCloser.closes.Load()) + require.EqualValues(t, 1, harness.connections()[0].closes.Load()) + require.Zero(t, server.hellos.Load(), "setup failure must not register a partial group") +} + +func TestWorkerCompletionConnectionsHelloFailureClosesWholeGroup(t *testing.T) { + server := newCompletionTestServer() + harness := newCompletionTestHarness(t, server, grpc.UnaryInterceptor( + func(context.Context, any, *grpc.UnaryServerInfo, grpc.UnaryHandler) (any, error) { + return nil, status.Error(codes.PermissionDenied, "test Hello failure") + })) + worker, err := NewTaskHubGrpcWorkerWithConnectionFactory( + harness.factory, task.NewTaskRegistry(), nil, WithWorkerCompletionConnections(3)) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + err = worker.Start(ctx) + require.Equal(t, codes.PermissionDenied, status.Code(err)) + require.ErrorContains(t, err, "worker Hello failed") + require.False(t, worker.Running()) + require.Len(t, harness.connections(), 4) + require.Empty(t, server.streams) + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } +} + +type completionErrorCloser struct { + countingCloser + err error +} + +func (c *completionErrorCloser) Close() error { + c.closes.Add(1) + return c.err +} + +func TestWorkerCompletionTransportClosesEveryConnectionOnceDespiteCloseErrors(t *testing.T) { + closeErr := errors.New("test close failure") + closers := []*completionErrorCloser{{err: closeErr}, {}, {}, {}} + index := 0 + group, err := newWorkerCompletionTransport(context.Background(), func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + closer := closers[index] + index++ + return &fakeClientConn{}, closer, nil + }, 3) + require.NoError(t, err) + for range 2 { + require.ErrorIs(t, group.Close(), closeErr) + } + for _, closer := range closers { + require.EqualValues(t, 1, closer.closes.Load()) + } +} + +func TestWorkerCompletionConnectionsValidation(t *testing.T) { + for _, count := range []int{-1, 0, 9} { + _, err := NewTaskHubGrpcWorkerWithConnectionFactory(func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + t.Fatal("invalid options must not dial") + return nil, nil, nil + }, task.NewTaskRegistry(), api.DefaultLogger(), WithWorkerCompletionConnections(count)) + require.ErrorContains(t, err, "between 1 and 8") + } + borrowed := &fakeClientConn{} + _, err := NewTaskHubGrpcWorker(borrowed, task.NewTaskRegistry(), nil, WithWorkerCompletionConnections(3)) + require.ErrorContains(t, err, "require separate connections") + require.Zero(t, borrowed.closes.Load()) + require.Zero(t, borrowed.streamCount()) + worker, err := NewTaskHubGrpcWorkerWithConnectionFactory( + func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { return borrowed, nil, nil }, + task.NewTaskRegistry(), nil) + require.NoError(t, err) + require.ErrorContains(t, worker.Start(context.Background()), "require an owning factory") + require.Zero(t, borrowed.closes.Load(), "missing ownership must not close the caller's channel") +} + +func TestWorkerCompletionConnectionsAcceptBoundedBudgets(t *testing.T) { + for _, budget := range []int{1, 8} { + t.Run(strconv.Itoa(budget), func(t *testing.T) { + server := newCompletionTestServer() + harness := newCompletionTestHarness(t, server) + worker, _ := startCompletionTestWorker(t, harness, budget, 1000) + intake := awaitCompletionTest(t, server.streams) + intake.commands <- completionTestActivity("budget") + call := awaitCompletionTest(t, server.calls) + require.Equal(t, 1, call.transport) + require.Len(t, harness.connections(), budget+1) + require.EqualValues(t, 1, server.hellos.Load()) + require.NoError(t, worker.Shutdown(context.Background())) + for _, closer := range harness.connections() { + require.EqualValues(t, 1, closer.closes.Load()) + } + }) + } +} + +func TestWorkerCompletionTransportsValidateBorrowedIsolation(t *testing.T) { + intake := &fakeClientConn{} + completion := &fakeClientConn{} + var typedNil *fakeClientConn + for _, test := range []struct { + name string + transports []grpc.ClientConnInterface + options []TaskHubGrpcWorkerOption + message string + }{ + {"empty", nil, nil, "between 1 and 8"}, + {"nil", []grpc.ClientConnInterface{nil}, nil, "is nil"}, + {"typed nil", []grpc.ClientConnInterface{typedNil}, nil, "is nil"}, + {"shared intake", []grpc.ClientConnInterface{intake}, nil, "separate from intake"}, + {"duplicate", []grpc.ClientConnInterface{completion, completion}, nil, "must be distinct"}, + {"budget mismatch", []grpc.ClientConnInterface{completion}, []TaskHubGrpcWorkerOption{WithWorkerCompletionConnections(3)}, "budget must match"}, + } { + t.Run(test.name, func(t *testing.T) { + opts := append([]TaskHubGrpcWorkerOption{WithWorkerCompletionTransports(test.transports...)}, test.options...) + _, err := NewTaskHubGrpcWorker(intake, task.NewTaskRegistry(), nil, opts...) + require.ErrorContains(t, err, test.message) + }) + } + _, err := NewTaskHubGrpcWorkerWithConnectionFactory( + func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + t.Fatal("invalid ownership options must not dial") + return nil, nil, nil + }, task.NewTaskRegistry(), nil, WithWorkerCompletionTransports(completion)) + require.ErrorContains(t, err, "owning worker factories cannot use") + require.Zero(t, intake.closes.Load()) + require.Zero(t, completion.closes.Load()) +} + +func TestWorkerCompletionConnectionsRejectReusedFactoryTransport(t *testing.T) { + cc := &fakeClientConn{} + closer := &countingCloser{} + group, err := newWorkerCompletionTransport(context.Background(), + func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + return cc, closer, nil + }, 3) + require.ErrorContains(t, err, "separate from intake") + require.Nil(t, group) + require.EqualValues(t, 1, closer.closes.Load(), "aliased ownership must roll back only once") +} + +func TestWorkerCompletionConnectionsRejectReusedFactoryCloser(t *testing.T) { + closer := &countingCloser{} + group, err := newWorkerCompletionTransport(context.Background(), + func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { + return &fakeClientConn{}, closer, nil + }, 3) + require.ErrorContains(t, err, "distinct ownership closers") + require.Nil(t, group) + require.EqualValues(t, 1, closer.closes.Load()) +} + +func TestWorkerCompletionTransportsBorrowedIsolationAndSnapshot(t *testing.T) { + server := newCompletionTestServer() + harness := newCompletionTestHarness(t, server) + connections := make([]grpc.ClientConnInterface, 4) + for i := range connections { + cc, _, err := harness.factory(context.Background()) + require.NoError(t, err) + connections[i] = cc + } + transports := append([]grpc.ClientConnInterface(nil), connections[1:]...) + option := WithWorkerCompletionTransports(transports...) + transports[0] = connections[0] + registry := task.NewTaskRegistry() + require.NoError(t, registry.AddActivityN("echo", func(task.ActivityContext) (any, error) { return "done", nil })) + worker, err := NewTaskHubGrpcWorker(connections[0], registry, nil, + option, WithWorkerCompletionConnections(3), + WithWorkerReconnectBackoff(time.Millisecond, time.Millisecond)) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + defer func() { require.NoError(t, worker.Shutdown(ctx)) }() + require.NoError(t, worker.Start(ctx)) + for generation := range 2 { + intake := awaitCompletionTest(t, server.streams) + require.Zero(t, intake.transport) + for i := range 3 { + intake.commands <- completionTestActivity(fmt.Sprintf("%d/%d", generation, i)) + } + transports := make(map[int]int) + for range 3 { + call := awaitCompletionTest(t, server.calls) + transports[call.transport]++ + } + require.Equal(t, map[int]int{1: 1, 2: 1, 3: 1}, transports) + if generation == 0 { + intake.commands <- fakeWorkItemResult{err: status.Error(codes.Unavailable, "borrowed reconnect")} + } + } + require.NoError(t, worker.Shutdown(ctx)) + require.Len(t, harness.connections(), 4, "a borrowed worker must not synthesize channels") + for _, closer := range harness.connections() { + require.Zero(t, closer.closes.Load(), "worker retirement and shutdown must preserve caller ownership") + require.NoError(t, closer.Close()) + } +} diff --git a/client/grpc_worker_drain_test.go b/client/grpc_worker_drain_test.go new file mode 100644 index 00000000..30b94490 --- /dev/null +++ b/client/grpc_worker_drain_test.go @@ -0,0 +1,166 @@ +package client + +import ( + "context" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/microsoft/durabletask-go/api" + "github.com/microsoft/durabletask-go/internal/protos" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type leaseAwareSchedulerClient struct { + *fakeSchedulerClient + acknowledged atomic.Int32 + expired atomic.Int32 + ackStarted chan struct{} + releaseAck chan struct{} +} + +func (c *leaseAwareSchedulerClient) CompleteActivityTask( + ctx context.Context, + response *protos.ActivityResponse, + options ...grpc.CallOption, +) (*protos.CompleteTaskResponse, error) { + if c.ackStarted != nil { + close(c.ackStarted) + select { + case <-c.releaseAck: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if c.stream.ctx.Err() != nil { + c.expired.Add(1) + return nil, status.Error(codes.NotFound, "intake ended and released the work item lease") + } + result, err := c.fakeSchedulerClient.CompleteActivityTask(ctx, response, options...) + if err == nil { + c.acknowledged.Add(1) + } + return result, err +} + +func TestWorkerGracefulDrainKeepsPendingAcknowledgementAndSilentIntakeAlive(t *testing.T) { + client := &leaseAwareSchedulerClient{ + fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(1)}, + ackStarted: make(chan struct{}), + releaseAck: make(chan struct{}), + } + worker := newFakeWorker(t, client.fakeSchedulerClient, + WithMaxConcurrentActivityWorkItems(1), + WithWorkerSilentDisconnectTimeout(100*time.Millisecond)) + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, nil, nil + } + worker.executor = &recordingExecutor{ + executeActivity: func(context.Context, api.InstanceID, *protos.HistoryEvent) (*protos.HistoryEvent, error) { + return &protos.HistoryEvent{EventType: &protos.HistoryEvent_TaskCompleted{ + TaskCompleted: &protos.TaskCompletedEvent{}, + }}, nil + }, + } + var releaseOnce sync.Once + t.Cleanup(func() { + releaseOnce.Do(func() { close(client.releaseAck) }) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, worker.Shutdown(ctx)) + require.NoError(t, worker.Wait(ctx)) + }) + require.NoError(t, worker.Start(context.Background())) + client.stream.results <- completionTestActivity("pending-ack") + awaitCompletionTest(t, client.ackStarted) + worker.mu.Lock() + run := worker.run + worker.mu.Unlock() + done := make(chan error, 1) + go func() { done <- worker.Shutdown(context.Background()) }() + awaitCompletionTest(t, run.intakeCtx.Done()) + select { + case <-client.stream.ctx.Done(): + t.Fatal("silence timeout released the lease while its acknowledgement was pending") + case <-time.After(2 * worker.options.silentDisconnectTimeout): + } + releaseOnce.Do(func() { close(client.releaseAck) }) + require.NoError(t, awaitCompletionTest(t, done)) + require.EqualValues(t, 1, client.acknowledged.Load()) + require.Zero(t, client.expired.Load()) +} + +func TestWorkerStoppedDispatchDoesNotExecuteEvenWithFreeSlot(t *testing.T) { + worker := newFakeWorker(t, &fakeSchedulerClient{}) + for range 100 { + run := newRunLoopTestRun() + run.cancelIntake() + var abandoned bool + err := worker.dispatch(run, &grpcWorkerConnection{}, run.activitySlots, + func(context.Context) { abandoned = true }, + func(context.Context) { t.Error("work was executed after logical intake stop") }) + require.ErrorIs(t, err, context.Canceled) + require.True(t, abandoned) + require.Empty(t, run.activitySlots) + run.pending.Wait() + run.cancelProcessing() + } +} + +func TestWorkerGracefulDrainPreservesIntakeLeaseUntilAcknowledged(t *testing.T) { + for _, stop := range []string{"shutdown", "run context"} { + t.Run(stop, func(t *testing.T) { + client := &leaseAwareSchedulerClient{fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(1)}} + worker := newFakeWorker(t, client.fakeSchedulerClient, WithMaxConcurrentActivityWorkItems(1)) + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, nil, nil + } + started, release := make(chan struct{}), make(chan struct{}) + var releaseOnce sync.Once + worker.executor = &recordingExecutor{ + executeActivity: func(ctx context.Context, _ api.InstanceID, _ *protos.HistoryEvent) (*protos.HistoryEvent, error) { + close(started) + select { + case <-release: + return &protos.HistoryEvent{EventType: &protos.HistoryEvent_TaskCompleted{ + TaskCompleted: &protos.TaskCompletedEvent{}, + }}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(func() { + releaseOnce.Do(func() { close(release) }) + cancel() + shutdownCtx, stopShutdown := context.WithTimeout(context.Background(), 5*time.Second) + defer stopShutdown() + require.NoError(t, worker.Shutdown(shutdownCtx)) + require.NoError(t, worker.Wait(shutdownCtx)) + }) + require.NoError(t, worker.Start(ctx)) + client.stream.results <- completionTestActivity("drain") + awaitCompletionTest(t, started) + worker.mu.Lock() + run := worker.run + worker.mu.Unlock() + if stop == "run context" { + cancel() + } + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- worker.Shutdown(context.Background()) }() + awaitCompletionTest(t, run.intakeCtx.Done()) + require.NoError(t, client.stream.ctx.Err(), "logical intake stop must not revoke an accepted work item's lease") + releaseOnce.Do(func() { close(release) }) + require.NoError(t, awaitCompletionTest(t, shutdownDone)) + require.EqualValues(t, 1, client.acknowledged.Load()) + require.Zero(t, client.expired.Load()) + }) + } +} diff --git a/client/grpc_worker_processor.go b/client/grpc_worker_processor.go index 3e9ae4c8..03044713 100644 --- a/client/grpc_worker_processor.go +++ b/client/grpc_worker_processor.go @@ -24,9 +24,33 @@ import ( ) func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *grpcWorkerConnection) (bool, error) { + stopWatcher, watcherDone := make(chan struct{}), make(chan struct{}) + waitAccepted := func() { + // Synchronize with the last dispatch that raced logical intake stop. + run.dispatchMu.Lock() + run.dispatchMu.Unlock() + connection.pending.Wait() + } + go func() { + defer close(watcherDone) + select { + case <-run.intakeCtx.Done(): + waitAccepted() + connection.cancelStream() + case <-stopWatcher: + } + }() + defer func() { + if run.intakeCtx.Err() != nil { + waitAccepted() + } + close(stopWatcher) + <-watcherDone + }() + observedMessage := false for { - workItem, err := w.receiveWorkItem(connection) + workItem, err := w.receiveWorkItem(run, connection) if err != nil { return observedMessage, err } @@ -62,8 +86,15 @@ func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *gr } } -func (w *TaskHubGrpcWorker) receiveWorkItem(connection *grpcWorkerConnection) (*protos.WorkItem, error) { - workItem, err := recvBeforeSilenceTimeout(connection.stream.Recv, connection.cancelStream, w.options.silentDisconnectTimeout) +func (w *TaskHubGrpcWorker) receiveWorkItem(run *grpcWorkerRun, connection *grpcWorkerConnection) (*protos.WorkItem, error) { + cancelOnSilence := func() { + // No new work is expected during a graceful drain. Keep leases alive + // until pending acknowledgements finish rather than timing intake out. + if run.intakeCtx.Err() == nil { + connection.cancelStream() + } + } + workItem, err := recvBeforeSilenceTimeout(connection.stream.Recv, cancelOnSilence, w.options.silentDisconnectTimeout) if err != nil { return nil, err } @@ -169,8 +200,16 @@ func (w *TaskHubGrpcWorker) dispatch( return run.intakeCtx.Err() } + run.dispatchMu.Lock() + if err := run.intakeCtx.Err(); err != nil { + run.dispatchMu.Unlock() + <-slots + abandon(run.processingCtx) + return err + } run.pending.Add(1) connection.pending.Add(1) + run.dispatchMu.Unlock() go func() { defer func() { <-slots diff --git a/client/grpc_worker_transport_test.go b/client/grpc_worker_transport_test.go index eafcda4f..734fcd0f 100644 --- a/client/grpc_worker_transport_test.go +++ b/client/grpc_worker_transport_test.go @@ -901,10 +901,12 @@ func TestBorrowedConnectionIsReusedAcrossReconnectsAndNeverClosed(t *testing.T) connection := &fakeClientConn{items: []*protos.WorkItem{{ Request: &protos.WorkItem_HealthPing{HealthPing: &protos.HealthPing{}}, }}} + completion := &fakeClientConn{} worker, err := NewTaskHubGrpcWorker( connection, task.NewTaskRegistry(), api.DefaultLogger(), + WithWorkerCompletionTransports(completion), WithWorkerHelloTimeout(time.Second), WithWorkerSilentDisconnectTimeout(time.Second), WithWorkerReconnectBackoff(time.Millisecond, 5*time.Millisecond), @@ -925,5 +927,7 @@ func TestBorrowedConnectionIsReusedAcrossReconnectsAndNeverClosed(t *testing.T) // the connection-factory worker, which durabletaskscheduler.NewWorker uses, // can obtain a fresh channel and close the retired one. require.Zero(t, connection.closes.Load(), "the worker must never close a borrowed connection") + require.Zero(t, completion.closes.Load(), "completion ownership also remains with the caller") + require.Zero(t, completion.streamCount(), "completion channels must never receive intake") require.GreaterOrEqual(t, connection.streamCount(), 2) } diff --git a/durabletaskscheduler/README.md b/durabletaskscheduler/README.md index 0da4a562..923767cb 100644 --- a/durabletaskscheduler/README.md +++ b/durabletaskscheduler/README.md @@ -188,14 +188,81 @@ A hand-built `Options` value with a zero failure threshold disables recreation. caller's `grpc.ClientConnInterface`; they never replace or close it. `NewWorker` owns its channels, recreates them after transient disconnects, and closes retired channels after their in-flight completions drain. -`client.NewTaskHubGrpcWorkerWithConnectionFactory` provides the same lifecycle -when its factory returns a non-nil closer, which transfers ownership to the -worker. Use an owning configuration to recover from a permanently wedged +`client.NewTaskHubGrpcWorkerWithConnectionFactory` provides the same lifecycle. +Its factory must return distinct connections with non-nil ownership closers. +Use an owning configuration to recover from a permanently wedged channel. After a worker has started, `Unauthenticated` and `PermissionDenied` stream or reconnect-handshake responses are retried with backoff so token refresh and RBAC propagation can recover without restarting the worker. See `client.NewTaskHubGrpcWorker` for the full ownership contract. +#### Required completion isolation (unpublished local draft) + +Workers use one intake connection and a separate bounded group for completion +and abandonment. Owned workers default to three completion connections, plus +intake. `client.WithWorkerCompletionConnections(n)` tunes that budget from 1 +through 8; zero is invalid and there is no shared-transport mode. The shared +connection exists only in a private regression baseline. + +The required architecture is user-approved for this local draft. The exact +public surface, initial budget of three, and upper bound of eight remain subject +to final review before publication. Three is a measured starting point, not a +universal optimum or a guarantee of higher persisted throughput. + +```go +// One logical worker, one intake, three completion connections by default. +worker, err := durabletaskscheduler.NewWorker(options, registry, logger) +// A smaller transport budget does not change execution limits. +worker, err = durabletaskscheduler.NewWorker(options, registry, logger, + durabletaskclient.WithWorkerCompletionConnections(1)) +``` + +The worker still has one identity, executor, Hello handshake, and GetWorkItems +stream per generation. Only the three completion RPCs and three abandon RPCs +are round-robin routed to the additional connections; all other calls, including +streamed history, stay on intake. Execution limits, retry policies, terminal +NotFound handling, and large-payload processing do not change. + +Every DTS connection uses the same prepared options, credential, worker ID, +TLS, token audience, metadata, interceptors, and message limits. DTS connections +must become transport-ready within the group setup +timeout (`WithWorkerHelloTimeout`); only intake performs Hello, with its existing +Hello timeout. This readiness stage now applies to every owned worker. +Completion-channel authentication failures are reported by the RPC, not +hidden by falling back to intake. A custom factory must honor cancellation and +return a distinct connection with a distinct non-nil ownership closer for each call. +Custom connection wrappers are responsible for their own transport readiness. + +`client.NewTaskHubGrpcWorker` and `TaskHubGrpcClient.StartWorkItemListener` +require `client.WithWorkerCompletionTransports(completionConnections...)`. +The supplied 1–8 transports must be separate from intake and each other, with +equivalent authentication and worker metadata. The caller owns readiness, +recovery, and closing; the worker neither constructs nor closes these channels. +Their explicit count determines the borrowed worker's budget. If +`WithWorkerCompletionConnections(n)` is also supplied, it must match that count. +Borrowed transport slices are snapshotted. Owning factories reject this borrowed +option rather than silently ignoring it. + +```go +// All connections are created/configured/closed by the caller. +worker, err := durabletaskclient.NewTaskHubGrpcWorker(intake, registry, logger, + durabletaskclient.WithWorkerCompletionTransports(completion1, completion2)) +``` + +Accepted work retains its original generation's connections until completion +or abandonment finishes. At most two groups (active and draining) are live: +repeated disconnects wait for the older group to drain before creating another. +This can delay recovery when old work is slow. Shutdown stops new dispatch but +keeps the lease-owning intake stream alive until accepted work and +acknowledgements drain. Silence does not terminate intake after a graceful stop. +Only then is the stream canceled and its group retired; a real stream failure +can still invalidate leases. The shutdown deadline cancels processing and +completion RPCs. As before, +entity abandonment uses its independent, bounded retry context. `Wait` joins +processing and retirement after a timed-out shutdown. Factories and application +handlers must cooperate with cancellation, and ownership closers must return +promptly. + Reconnect and RPC retry delays are deterministic and always stay within `[baseDelay, maxDelay]`. A stream that delivers at least one message before it ends is treated as a drain and restarts the schedule at the base delay; a stream diff --git a/durabletaskscheduler/connection_test.go b/durabletaskscheduler/connection_test.go index 56d58e75..0a4820df 100644 --- a/durabletaskscheduler/connection_test.go +++ b/durabletaskscheduler/connection_test.go @@ -21,6 +21,7 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" "google.golang.org/grpc/test/bufconn" @@ -540,10 +541,16 @@ func TestClientCloseStopsCompatibilityListener(t *testing.T) { managementClient, err := NewClient(context.Background(), options, api.DefaultLogger()) require.NoError(t, err) <-server.metadata - require.NoError(t, managementClient.StartWorkItemListener(context.Background(), task.NewTaskRegistry())) + completion, err := connect(options, clientRole, "") + require.NoError(t, err) + defer func() { require.NoError(t, completion.Close()) }() + require.ErrorContains(t, managementClient.StartWorkItemListener(context.Background(), task.NewTaskRegistry()), "require separate connections") + require.NoError(t, managementClient.StartWorkItemListener(context.Background(), task.NewTaskRegistry(), + durabletaskclient.WithWorkerCompletionTransports(completion))) <-server.metadata <-server.metadata require.NoError(t, managementClient.Close()) + require.NotEqual(t, connectivity.Shutdown, completion.GetState(), "the listener must not close borrowed completion ownership") shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() diff --git a/durabletaskscheduler/resource_id_test.go b/durabletaskscheduler/resource_id_test.go index d8f76133..f48e4cbc 100644 --- a/durabletaskscheduler/resource_id_test.go +++ b/durabletaskscheduler/resource_id_test.go @@ -394,7 +394,11 @@ func TestPublicAuthenticationPathsPreserveResourceIDAcrossReconnects(t *testing. awaitAudienceMetadata(t, server) require.Len(t, recordedTokenOptions(credential), 1, "client acquires a token for its eager Hello") if path == "compatibility listener" { + completion, err := connect(options, clientRole, "") + require.NoError(t, err) + defer func() { require.NoError(t, completion.Close()) }() require.NoError(t, client.StartWorkItemListener(ctx, task.NewTaskRegistry(), + durabletaskclient.WithWorkerCompletionTransports(completion), durabletaskclient.WithWorkerReconnectBackoff(time.Millisecond, time.Millisecond))) awaitAudienceMetadata(t, server) awaitAudienceMetadata(t, server) diff --git a/durabletaskscheduler/worker_completion_test.go b/durabletaskscheduler/worker_completion_test.go new file mode 100644 index 00000000..a2861b5f --- /dev/null +++ b/durabletaskscheduler/worker_completion_test.go @@ -0,0 +1,220 @@ +package durabletaskscheduler + +import ( + "context" + "os" + "os/exec" + "strings" + "sync/atomic" + "testing" + "time" + + durabletaskclient "github.com/microsoft/durabletask-go/client" + "github.com/microsoft/durabletask-go/internal/protos" + "github.com/microsoft/durabletask-go/task" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/wrapperspb" +) + +type workerCompletionObservation struct { + method string + metadata metadata.MD + response *protos.ActivityResponse +} + +type workerCompletionOptionsServer struct { + protos.UnimplementedTaskHubSidecarServiceServer + observed chan workerCompletionObservation + disconnect chan struct{} + intakes atomic.Int32 +} + +func (s *workerCompletionOptionsServer) record(ctx context.Context, method string, response *protos.ActivityResponse) { + md, _ := metadata.FromIncomingContext(ctx) + s.observed <- workerCompletionObservation{method: method, metadata: md, response: response} +} + +func (s *workerCompletionOptionsServer) Hello(ctx context.Context, _ *emptypb.Empty) (*emptypb.Empty, error) { + s.record(ctx, protos.TaskHubSidecarService_Hello_FullMethodName, nil) + return &emptypb.Empty{}, nil +} + +func (s *workerCompletionOptionsServer) GetWorkItems(request *protos.GetWorkItemsRequest, stream protos.TaskHubSidecarService_GetWorkItemsServer) error { + s.intakes.Add(1) + s.record(stream.Context(), protos.TaskHubSidecarService_GetWorkItems_FullMethodName, nil) + if request.MaxConcurrentActivityWorkItems != 1000 || request.MaxConcurrentOrchestrationWorkItems != 1000 { + return status.Error(codes.InvalidArgument, "execution limits changed") + } + for i := range 6 { + if err := stream.Send(&protos.WorkItem{ + CompletionToken: "token", + Request: &protos.WorkItem_ActivityRequest{ActivityRequest: &protos.ActivityRequest{ + Name: "echo", + TaskId: int32(i), + OrchestrationInstance: &protos.OrchestrationInstance{InstanceId: "instance"}, + }}, + }); err != nil { + return err + } + } + select { + case <-s.disconnect: + return status.Error(codes.Unavailable, "test reconnect") + case <-stream.Context().Done(): + return stream.Context().Err() + } +} + +func (s *workerCompletionOptionsServer) CompleteActivityTask(ctx context.Context, response *protos.ActivityResponse) (*protos.CompleteTaskResponse, error) { + s.record(ctx, protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName, response) + return &protos.CompleteTaskResponse{}, nil +} + +func (*workerCompletionOptionsServer) GetInstance(context.Context, *protos.GetInstanceRequest) (*protos.GetInstanceResponse, error) { + return &protos.GetInstanceResponse{ + Exists: true, + OrchestrationState: &protos.OrchestrationState{Output: wrapperspb.String(strings.Repeat("x", 70*1024))}, + }, nil +} + +func TestNewWorkerCompletionConnectionsPreservePreparedOptions(t *testing.T) { + // Like the audience tests, isolate fallback roots in a subprocess instead + // of modifying trust for other tests or installing a system certificate. + if os.Getenv("DTS_COMPLETION_TLS_TEST") != "1" { + command := exec.Command(os.Args[0], "-test.run=^TestNewWorkerCompletionConnectionsPreservePreparedOptions$", "-test.timeout=2m") + command.Env = append(os.Environ(), "DTS_COMPLETION_TLS_TEST=1") + output, err := command.CombinedOutput() + require.NoError(t, err, "%s", output) + return + } + t.Setenv("GODEBUG", os.Getenv("GODEBUG")+",x509usefallbackroots=1") + server := &workerCompletionOptionsServer{ + observed: make(chan workerCompletionObservation, 32), + disconnect: make(chan struct{}, 1), + } + listener := bufconn.Listen(1024 * 1024) + grpcServer := grpc.NewServer(grpc.Creds(credentials.NewTLS(audienceTestTLS()))) + protos.RegisterTaskHubSidecarServiceServer(grpcServer, server) + done := make(chan struct{}) + go func() { + defer close(done) + _ = grpcServer.Serve(listener) + }() + defer func() { + grpcServer.Stop() + require.NoError(t, listener.Close()) + <-done + }() + + credential := &recordingCredential{} + options := NewOptionsWithCredential("https://127.0.0.1", "test-hub", credential) + options.ResourceID = "api://completion-test/.DEFAULT/" + options.WorkerID = "one-worker" + options.UserAgent = "completion-test-agent" + options.DataConverter = recreationDataConverter{} + options.MaxReceiveMessageSize = 64 * 1024 + options.MaxSendMessageSize = 64 * 1024 + options.dialer = bufconnDialer(listener) + connections := make(chan *grpc.ClientConn, 32) + options.UnaryInterceptors = []grpc.UnaryClientInterceptor{ + func(ctx context.Context, method string, request, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, callOptions ...grpc.CallOption) error { + connections <- cc + return invoker(metadata.AppendToOutgoingContext(ctx, "test-unary", "preserved"), method, request, reply, cc, callOptions...) + }, + } + options.StreamInterceptors = []grpc.StreamClientInterceptor{ + func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, callOptions ...grpc.CallOption) (grpc.ClientStream, error) { + return streamer(metadata.AppendToOutgoingContext(ctx, "test-stream", "preserved"), desc, cc, method, callOptions...) + }, + } + registry := task.NewTaskRegistry() + require.NoError(t, registry.AddActivityN("echo", func(task.ActivityContext) (any, error) { return "done", nil })) + worker, err := NewWorker(options, registry, nil, + durabletaskclient.WithMaxConcurrentActivityWorkItems(1000), + durabletaskclient.WithMaxConcurrentOrchestrationWorkItems(1000), + durabletaskclient.WithWorkerReconnectBackoff(time.Millisecond, time.Millisecond)) + require.NoError(t, err) + // A worker snapshots options once, including the credential and identity. + options.ResourceID = "api://changed" + options.WorkerID = "changed" + options.UserAgent = "changed" + options.MaxSendMessageSize = 128 * 1024 + options.UnaryInterceptors = nil + require.Empty(t, recordedTokenOptions(credential), "authentication stays lazy before Start") + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, worker.Start(ctx)) + defer func() { require.NoError(t, worker.Shutdown(ctx)) }() + + allConnections := make(map[*grpc.ClientConn]struct{}) + for generation := range 2 { + counts := make(map[string]int) + for range 8 { + var observation workerCompletionObservation + select { + case observation = <-server.observed: + case <-ctx.Done(): + t.Fatal("worker did not complete all test work") + } + counts[observation.method]++ + require.Equal(t, []string{"test-hub"}, observation.metadata.Get("taskhub")) + require.Equal(t, []string{"one-worker"}, observation.metadata.Get("workerid")) + require.Equal(t, []string{"completion-test-agent"}, observation.metadata.Get("x-user-agent")) + require.Equal(t, []string{"Bearer token"}, observation.metadata.Get("authorization")) + if observation.method == protos.TaskHubSidecarService_GetWorkItems_FullMethodName { + require.Equal(t, []string{"preserved"}, observation.metadata.Get("test-stream")) + } else { + require.Equal(t, []string{"preserved"}, observation.metadata.Get("test-unary")) + } + if observation.response != nil { + require.Equal(t, "recreated:done", observation.response.Result.GetValue()) + } + } + require.Equal(t, map[string]int{ + protos.TaskHubSidecarService_Hello_FullMethodName: 1, + protos.TaskHubSidecarService_GetWorkItems_FullMethodName: 1, + protos.TaskHubSidecarService_CompleteActivityTask_FullMethodName: 6, + }, counts) + generationConnections := make(map[*grpc.ClientConn]struct{}) + for range 7 { + select { + case connection := <-connections: + generationConnections[connection] = struct{}{} + allConnections[connection] = struct{}{} + case <-ctx.Done(): + t.Fatal("missing intercepted RPC") + } + } + require.Len(t, generationConnections, 4, "one intake and three completion channels") + require.Len(t, recordedTokenOptions(credential), 4*(generation+1)) + if generation == 0 { + for connection := range generationConnections { + client := protos.NewTaskHubSidecarServiceClient(connection) + _, err = client.CompleteActivityTask(ctx, &protos.ActivityResponse{Result: wrapperspb.String(strings.Repeat("x", 70*1024))}) + require.Equal(t, codes.ResourceExhausted, status.Code(err), "send limits must match intake on every channel") + _, err = client.GetInstance(ctx, &protos.GetInstanceRequest{}) + require.Equal(t, codes.ResourceExhausted, status.Code(err), "receive limits must match intake on every channel") + <-connections + <-connections + } + server.disconnect <- struct{}{} + } + } + for _, recorded := range recordedTokenOptions(credential) { + require.Equal(t, []string{"api://completion-test/.default"}, recorded.Scopes) + } + require.NoError(t, worker.Shutdown(ctx)) + require.EqualValues(t, 2, server.intakes.Load()) + require.Len(t, allConnections, 8) + for connection := range allConnections { + require.Equal(t, connectivity.Shutdown, connection.GetState(), "every owned channel must close after drain") + } +} diff --git a/exporthistory/versioning_test.go b/exporthistory/versioning_test.go index 8a4c56c5..854d2f84 100644 --- a/exporthistory/versioning_test.go +++ b/exporthistory/versioning_test.go @@ -96,6 +96,17 @@ func startFilterCapturingWorker( ) require.NoError(t, err) + completion, err := grpc.NewClient( + "passthrough:///export-history-filters", + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return listener.Dial() + }), + ) + require.NoError(t, err) + workerOptions = append([]durabletaskclient.TaskHubGrpcWorkerOption{ + durabletaskclient.WithWorkerCompletionTransports(completion), + }, workerOptions...) worker, err := durabletaskclient.NewTaskHubGrpcWorker( connection, registry, @@ -109,6 +120,7 @@ func startFilterCapturingWorker( defer cancel() require.NoError(t, worker.Shutdown(shutdownCtx)) require.NoError(t, connection.Close()) + require.NoError(t, completion.Close()) server.Stop() require.NoError(t, listener.Close()) }) diff --git a/samples/largepayloads/largepayloads_test.go b/samples/largepayloads/largepayloads_test.go index 871a5066..88c3ce89 100644 --- a/samples/largepayloads/largepayloads_test.go +++ b/samples/largepayloads/largepayloads_test.go @@ -78,6 +78,7 @@ func TestStorageCleanupWaitsForAcceptedUploads(t *testing.T) { require.NoError(t, err) items := make(chan *protos.WorkItem, 1) intakeStopped := make(chan struct{}) + cleanupStarted := make(chan struct{}) server := grpc.NewServer( grpc.UnaryInterceptor(func(ctx context.Context, request any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { switch request := request.(type) { @@ -116,6 +117,7 @@ func TestStorageCleanupWaitsForAcceptedUploads(t *testing.T) { return &protos.TerminateResponse{}, nil case *protos.PurgeInstancesRequest: purged.Store(true) + close(cleanupStarted) return &protos.PurgeInstancesResponse{DeletedInstanceCount: 1}, nil case *protos.ActivityResponse: completed.Store(true) @@ -136,6 +138,9 @@ func TestStorageCleanupWaitsForAcceptedUploads(t *testing.T) { return nil } <-stream.Context().Done() + if !shutdownTimesOut && !completed.Load() { + t.Error("intake released its lease before the accepted upload was acknowledged") + } close(intakeStopped) return nil }), @@ -161,12 +166,12 @@ func TestStorageCleanupWaitsForAcceptedUploads(t *testing.T) { } }) select { - case <-intakeStopped: + case <-cleanupStarted: case <-time.After(5 * time.Second): t.Fatal("sample did not reach worker shutdown with an accepted upload") } require.True(t, purged.Load()) - require.False(t, deleted.Load(), "intake cancellation alone does not drain accepted uploads") + require.False(t, deleted.Load(), "orchestration cleanup alone does not drain accepted uploads") if !shutdownTimesOut { release() } @@ -186,6 +191,11 @@ func TestStorageCleanupWaitsForAcceptedUploads(t *testing.T) { case <-time.After(35 * time.Second): t.Fatal("sample did not bound worker shutdown") } + select { + case <-intakeStopped: + case <-time.After(5 * time.Second): + t.Fatal("worker intake did not close after drain or deadline cancellation") + } release() select { case <-uploadFinished: diff --git a/samples/worker/README.md b/samples/worker/README.md index 1b2f8c87..7861bf31 100644 --- a/samples/worker/README.md +++ b/samples/worker/README.md @@ -22,8 +22,9 @@ go run ./samples/worker the orchestration continuation after the drained activity completes after restart. - `Run` blocks until its context is canceled and exits cleanly. -- The compatibility listener on a management client demonstrates the borrowed - connection contract; new applications should prefer owned workers. +- An owned worker can use a smaller completion budget (one dedicated completion + connection) independently of its activity-execution limit. Completion isolation + is required; the normal worker default is three completion connections. Activities are at-least-once, so the drain scenario uses an idempotent output and does not assert an exact delivery count. The recovery scenario is limited to diff --git a/samples/worker/worker.go b/samples/worker/worker.go index 18c1b895..c5cb57b8 100644 --- a/samples/worker/worker.go +++ b/samples/worker/worker.go @@ -65,8 +65,8 @@ func run() (err error) { drainID := dtssample.NewInstanceID("worker-drain") restartID := dtssample.NewInstanceID("worker-restart") runID := dtssample.NewInstanceID("worker-run") - borrowedID := dtssample.NewInstanceID("worker-borrowed") - ids = append(ids, drainID, restartID, runID, borrowedID) + budgetID := dtssample.NewInstanceID("worker-budget") + ids = append(ids, drainID, restartID, runID, budgetID) if err := verifyStartConcurrencyAndDrain(ctx, options, client, registry, concurrencyIDs, drainID); err != nil { return err @@ -77,7 +77,7 @@ func run() (err error) { if err := verifyRun(ctx, options, client, registry, runID); err != nil { return err } - if err := verifyBorrowedConnectionListener(ctx, client, registry, borrowedID); err != nil { + if err := verifyCompletionBudget(ctx, options, client, registry, budgetID); err != nil { return err } return nil @@ -253,35 +253,39 @@ func verifyRun( return nil } -func verifyBorrowedConnectionListener( +func verifyCompletionBudget( ctx context.Context, + options *durabletaskscheduler.Options, client *durabletaskscheduler.Client, registry *task.TaskRegistry, id api.InstanceID, -) error { - if err := client.StartWorkItemListener(ctx, registry, +) (err error) { + worker, err := durabletaskscheduler.NewWorker(options, registry, api.DefaultLogger(), durabletaskclient.WithAutoWorkItemFilters(), - durabletaskclient.WithMaxConcurrentActivityWorkItems(1)); err != nil { + durabletaskclient.WithMaxConcurrentActivityWorkItems(1), + durabletaskclient.WithWorkerCompletionConnections(1)) + if err != nil { + return err + } + if err := worker.Start(ctx); err != nil { return err } defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - _ = client.StopWorkItemListener(shutdownCtx) + err = errors.Join(err, shutdownWorker(worker)) }() if _, err := client.ScheduleNewOrchestration(ctx, "SampleWorkerEcho", - api.WithInstanceID(id), api.WithInput("borrowed")); err != nil { + api.WithInstanceID(id), api.WithInput("budget")); err != nil { return err } - if err := waitForWorkerOutput(ctx, client, id, "echo:borrowed"); err != nil { + if err := waitForWorkerOutput(ctx, client, id, "echo:budget"); err != nil { return err } shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := client.StopWorkItemListener(shutdownCtx); err != nil { + if err := worker.Shutdown(shutdownCtx); err != nil { return err } - fmt.Println("verified borrowed connection compatibility listener") + fmt.Println("verified worker with one dedicated completion connection") return nil } From 4729a07a284eb18370dccc07673fccf5f046e16b Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:38:49 -0700 Subject: [PATCH 2/7] fix lint findings in worker completion transport Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker_connection_test.go | 4 +++- client/grpc_worker_processor.go | 11 ++++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/client/grpc_worker_connection_test.go b/client/grpc_worker_connection_test.go index 338c637a..5c3f97b0 100644 --- a/client/grpc_worker_connection_test.go +++ b/client/grpc_worker_connection_test.go @@ -566,7 +566,9 @@ func TestWorkerCompletionTransportRoutesExactMethodsAndPreservesCallOptions(t *t if err := grpc.SendHeader(ctx, metadata.Pairs("test-response", "header")); err != nil { return nil, err } - grpc.SetTrailer(ctx, metadata.Pairs("test-response", "trailer")) + if err := grpc.SetTrailer(ctx, metadata.Pairs("test-response", "trailer")); err != nil { + return nil, err + } return handler(ctx, request) })) ctx := metadata.AppendToOutgoingContext(context.Background(), "test-context", "preserved") diff --git a/client/grpc_worker_processor.go b/client/grpc_worker_processor.go index 03044713..505b676e 100644 --- a/client/grpc_worker_processor.go +++ b/client/grpc_worker_processor.go @@ -23,12 +23,17 @@ import ( "google.golang.org/protobuf/types/known/wrapperspb" ) +// awaitDispatchBarrier synchronizes with any dispatch that raced logical intake +// stop, so its pending count is visible before callers wait on it. +func (run *grpcWorkerRun) awaitDispatchBarrier() { + run.dispatchMu.Lock() + defer run.dispatchMu.Unlock() +} + func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *grpcWorkerConnection) (bool, error) { stopWatcher, watcherDone := make(chan struct{}), make(chan struct{}) waitAccepted := func() { - // Synchronize with the last dispatch that raced logical intake stop. - run.dispatchMu.Lock() - run.dispatchMu.Unlock() + run.awaitDispatchBarrier() connection.pending.Wait() } go func() { From 6747da8f2c34e9ecc7ae4fb3fb235feac585f326 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:07:32 -0700 Subject: [PATCH 3/7] address review feedback on completion isolation docs and errors Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker_connection.go | 2 +- client/grpc_worker_connection_test.go | 2 +- durabletaskscheduler/README.md | 12 ++++-------- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/client/grpc_worker_connection.go b/client/grpc_worker_connection.go index f325a4f7..acde82a8 100644 --- a/client/grpc_worker_connection.go +++ b/client/grpc_worker_connection.go @@ -96,7 +96,7 @@ func newWorkerCompletionTransport( case isNilWorkerTransport(cc): err = fmt.Errorf("gRPC worker connection factory returned a nil connection") case isNilWorkerTransport(closer): - err = fmt.Errorf("completion connections require an owning factory with a closer for every connection") + err = fmt.Errorf("owned worker connections require a factory that returns a closer for every connection, including intake") case ctx.Err() != nil: err = ctx.Err() } diff --git a/client/grpc_worker_connection_test.go b/client/grpc_worker_connection_test.go index 5c3f97b0..73c588d1 100644 --- a/client/grpc_worker_connection_test.go +++ b/client/grpc_worker_connection_test.go @@ -802,7 +802,7 @@ func TestWorkerCompletionConnectionsValidation(t *testing.T) { func(context.Context) (grpc.ClientConnInterface, io.Closer, error) { return borrowed, nil, nil }, task.NewTaskRegistry(), nil) require.NoError(t, err) - require.ErrorContains(t, worker.Start(context.Background()), "require an owning factory") + require.ErrorContains(t, worker.Start(context.Background()), "require a factory that returns a closer") require.Zero(t, borrowed.closes.Load(), "missing ownership must not close the caller's channel") } diff --git a/durabletaskscheduler/README.md b/durabletaskscheduler/README.md index 923767cb..cfdf2fa5 100644 --- a/durabletaskscheduler/README.md +++ b/durabletaskscheduler/README.md @@ -196,18 +196,14 @@ stream or reconnect-handshake responses are retried with backoff so token refresh and RBAC propagation can recover without restarting the worker. See `client.NewTaskHubGrpcWorker` for the full ownership contract. -#### Required completion isolation (unpublished local draft) +#### Completion isolation Workers use one intake connection and a separate bounded group for completion and abandonment. Owned workers default to three completion connections, plus intake. `client.WithWorkerCompletionConnections(n)` tunes that budget from 1 -through 8; zero is invalid and there is no shared-transport mode. The shared -connection exists only in a private regression baseline. - -The required architecture is user-approved for this local draft. The exact -public surface, initial budget of three, and upper bound of eight remain subject -to final review before publication. Three is a measured starting point, not a -universal optimum or a guarantee of higher persisted throughput. +through 8; zero is invalid and there is no shared-transport mode. Three is a +starting point, not a universal optimum or a guarantee of higher persisted +throughput. ```go // One logical worker, one intake, three completion connections by default. From 615abb67b86d11f1b754b9324fc48b937f79cf23 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:23:30 -0700 Subject: [PATCH 4/7] keep intake stop able to interrupt a blocked stream open Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker.go | 15 ++++++- client/grpc_worker_drain_test.go | 75 ++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/client/grpc_worker.go b/client/grpc_worker.go index 868d9f08..1ff2c3eb 100644 --- a/client/grpc_worker.go +++ b/client/grpc_worker.go @@ -1159,9 +1159,12 @@ func (w *TaskHubGrpcWorker) connect(ctx context.Context) (*grpcWorkerConnection, if err := ctx.Err(); err != nil { return nil, closeOnError(fmt.Errorf("worker intake stopped before opening the stream: %w", status.FromContextError(err).Err())) } - // Ending the stream releases server-side leases, so logical intake stop - // must not cancel it until accepted work has acknowledged or abandoned. + // Ending an established stream releases its server-side leases, so logical + // intake stop must not cancel it until accepted work has acknowledged or + // abandoned. While the stream is still opening nothing is leased, so intake + // stop must be able to interrupt a blocked open. streamCtx, cancelStream := context.WithCancel(context.WithoutCancel(ctx)) + stopOpenCancel := context.AfterFunc(ctx, cancelStream) stream, err := client.GetWorkItems(streamCtx, &protos.GetWorkItemsRequest{ MaxConcurrentOrchestrationWorkItems: int32(w.options.maxConcurrentOrchestrations), MaxConcurrentActivityWorkItems: int32(w.options.maxConcurrentActivities), @@ -1169,6 +1172,14 @@ func (w *TaskHubGrpcWorker) connect(ctx context.Context) (*grpcWorkerConnection, Capabilities: slices.Clone(w.options.capabilities), WorkItemFilters: workItemFiltersToProto(w.options.workItemFilters), }) + if !stopOpenCancel() { + // Intake stopped during the open; the stream was or is being cancelled. + cancelStream() + if err == nil { + err = status.FromContextError(ctx.Err()).Err() + } + return nil, closeOnError(fmt.Errorf("worker intake stopped while opening the stream: %w", err)) + } if err != nil { cancelStream() return nil, closeOnError(fmt.Errorf("failed to open gRPC work item stream: %w", err)) diff --git a/client/grpc_worker_drain_test.go b/client/grpc_worker_drain_test.go index 30b94490..0915b951 100644 --- a/client/grpc_worker_drain_test.go +++ b/client/grpc_worker_drain_test.go @@ -164,3 +164,78 @@ func TestWorkerGracefulDrainPreservesIntakeLeaseUntilAcknowledged(t *testing.T) }) } } + +// openingSchedulerClient lets a test control how GetWorkItems establishes its +// stream, such as blocking at a saturated HTTP/2 stream limit. +type openingSchedulerClient struct { + *fakeSchedulerClient + open func(context.Context) (protos.TaskHubSidecarService_GetWorkItemsClient, error) +} + +func (c *openingSchedulerClient) GetWorkItems( + ctx context.Context, + _ *protos.GetWorkItemsRequest, + _ ...grpc.CallOption, +) (protos.TaskHubSidecarService_GetWorkItemsClient, error) { + return c.open(ctx) +} + +func TestWorkerIntakeStopInterruptsStreamOpen(t *testing.T) { + for _, stop := range []string{"shutdown", "run context"} { + t.Run(stop, func(t *testing.T) { + opening := make(chan struct{}) + client := &openingSchedulerClient{ + fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(0)}, + open: func(ctx context.Context) (protos.TaskHubSidecarService_GetWorkItemsClient, error) { + close(opening) + <-ctx.Done() + return nil, status.FromContextError(ctx.Err()).Err() + }, + } + worker := newFakeWorker(t, client.fakeSchedulerClient) + closer := &countingCloser{} + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, closer, nil + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + startDone := make(chan error, 1) + go func() { startDone <- worker.Start(ctx) }() + awaitCompletionTest(t, opening) + + if stop == "run context" { + cancel() + } else { + shutdownCtx, stopShutdown := context.WithTimeout(context.Background(), 5*time.Second) + defer stopShutdown() + _ = worker.Shutdown(shutdownCtx) + } + require.Error(t, awaitCompletionTest(t, startDone)) + require.False(t, worker.Running()) + require.EqualValues(t, 1, closer.closes.Load()) + }) + } +} + +func TestWorkerIntakeStopRacingStreamOpenDoesNotEstablishCancelledStream(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var streamCtx context.Context + client := &openingSchedulerClient{fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(0)}} + client.open = func(ctx context.Context) (protos.TaskHubSidecarService_GetWorkItemsClient, error) { + streamCtx = ctx + // The stop lands after the open succeeded but before connect returns. + cancel() + return client.stream, nil + } + worker := newFakeWorker(t, client.fakeSchedulerClient) + closer := &countingCloser{} + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, closer, nil + } + + require.Error(t, worker.Start(ctx)) + require.False(t, worker.Running()) + require.EqualValues(t, 1, closer.closes.Load()) + awaitCompletionTest(t, streamCtx.Done()) +} From 89f6263a80c9dc9b155aaadf79a6494d9d2b5a91 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:36:44 -0700 Subject: [PATCH 5/7] drain canceled dispatch abandonments before releasing intake leases Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker.go | 2 +- client/grpc_worker_drain_test.go | 183 ++++++++++++++++++++++++++++++- client/grpc_worker_processor.go | 17 ++- durabletaskscheduler/README.md | 4 +- 4 files changed, 196 insertions(+), 10 deletions(-) diff --git a/client/grpc_worker.go b/client/grpc_worker.go index 1ff2c3eb..34582b8e 100644 --- a/client/grpc_worker.go +++ b/client/grpc_worker.go @@ -951,7 +951,7 @@ func (w *TaskHubGrpcWorker) Run(ctx context.Context) error { } // Shutdown stops new dispatch and keeps the lease-owning intake stream open -// until accepted execution and acknowledgements finish. If ctx expires, +// until received work completes or is abandoned. If ctx expires, // in-flight execution and completion RPCs are canceled. func (w *TaskHubGrpcWorker) Shutdown(ctx context.Context) error { w.mu.Lock() diff --git a/client/grpc_worker_drain_test.go b/client/grpc_worker_drain_test.go index 0915b951..3019760d 100644 --- a/client/grpc_worker_drain_test.go +++ b/client/grpc_worker_drain_test.go @@ -6,6 +6,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/microsoft/durabletask-go/api" @@ -18,10 +19,13 @@ import ( type leaseAwareSchedulerClient struct { *fakeSchedulerClient - acknowledged atomic.Int32 - expired atomic.Int32 - ackStarted chan struct{} - releaseAck chan struct{} + acknowledged atomic.Int32 + expired atomic.Int32 + ackStarted chan struct{} + releaseAck chan struct{} + abandoned atomic.Int32 + abandonStarted chan struct{} + releaseAbandon chan struct{} } func (c *leaseAwareSchedulerClient) CompleteActivityTask( @@ -48,6 +52,177 @@ func (c *leaseAwareSchedulerClient) CompleteActivityTask( return result, err } +func (c *leaseAwareSchedulerClient) AbandonTaskActivityWorkItem( + ctx context.Context, + request *protos.AbandonActivityTaskRequest, + options ...grpc.CallOption, +) (*protos.AbandonActivityTaskResponse, error) { + if c.abandonStarted != nil { + close(c.abandonStarted) + select { + case <-c.releaseAbandon: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if c.stream.ctx.Err() != nil { + c.expired.Add(1) + return nil, status.Error(codes.NotFound, "intake ended and released the work item lease") + } + result, err := c.fakeSchedulerClient.AbandonTaskActivityWorkItem(ctx, request, options...) + if err == nil { + c.abandoned.Add(1) + } + return result, err +} + +func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { + for _, phase := range []string{"waiting for slot", "slot reserved"} { + t.Run(phase, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + worker := newFakeWorker(t, &fakeSchedulerClient{}) + run := newRunLoopTestRun() + defer run.cancelIntake() + defer run.cancelProcessing() + connection := &grpcWorkerConnection{} + slots := make(chan struct{}) + if phase == "waiting for slot" { + slots = make(chan struct{}, 1) + slots <- struct{}{} + } + started, release := make(chan struct{}), make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) }) + dispatchDone := make(chan error, 1) + go func() { + dispatchDone <- worker.dispatch(run, connection, slots, + func(context.Context) { + close(started) + <-release + }, + func(context.Context) { t.Error("work was executed after logical intake stop") }) + }() + synctest.Wait() + if phase == "slot reserved" { + run.dispatchMu.Lock() + <-slots + run.cancelIntake() + run.dispatchMu.Unlock() + // Return the admission token so the rejected dispatch can + // release its slot after observing cancellation. + slots <- struct{}{} + } else { + run.cancelIntake() + } + awaitCompletionTest(t, started) + + runDrained, connectionDrained := make(chan struct{}), make(chan struct{}) + go func() { + run.pending.Wait() + close(runDrained) + }() + go func() { + connection.pending.Wait() + close(connectionDrained) + }() + synctest.Wait() + for _, drained := range []<-chan struct{}{runDrained, connectionDrained} { + select { + case <-drained: + t.Error("drain finished before the abandonment returned") + default: + } + } + + releaseOnce.Do(func() { close(release) }) + require.ErrorIs(t, awaitCompletionTest(t, dispatchDone), context.Canceled) + awaitCompletionTest(t, runDrained) + awaitCompletionTest(t, connectionDrained) + if phase == "waiting for slot" { + <-slots + } + require.Empty(t, slots) + }) + }) + } +} + +func TestWorkerGracefulStopKeepsRejectedItemLeaseUntilAbandoned(t *testing.T) { + for _, stop := range []string{"shutdown", "run context", "shutdown deadline"} { + t.Run(stop, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + client := &leaseAwareSchedulerClient{ + fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(0)}, + abandonStarted: make(chan struct{}), + releaseAbandon: make(chan struct{}), + } + worker := newFakeWorker(t, client.fakeSchedulerClient, WithMaxConcurrentActivityWorkItems(1)) + closer := &countingCloser{} + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, closer, nil + } + executionStarted, releaseExecution := make(chan struct{}), make(chan struct{}) + worker.executor = &recordingExecutor{ + executeActivity: func(context.Context, api.InstanceID, *protos.HistoryEvent) (*protos.HistoryEvent, error) { + close(executionStarted) + <-releaseExecution + return &protos.HistoryEvent{EventType: &protos.HistoryEvent_TaskCompleted{ + TaskCompleted: &protos.TaskCompletedEvent{}, + }}, nil + }, + } + ctx, cancel := context.WithCancel(context.Background()) + var executionOnce, abandonOnce sync.Once + t.Cleanup(func() { + executionOnce.Do(func() { close(releaseExecution) }) + abandonOnce.Do(func() { close(client.releaseAbandon) }) + cancel() + shutdownCtx, stopShutdown := context.WithTimeout(context.Background(), 5*time.Second) + defer stopShutdown() + require.NoError(t, worker.Shutdown(shutdownCtx)) + require.NoError(t, worker.Wait(shutdownCtx)) + }) + require.NoError(t, worker.Start(ctx)) + client.stream.results <- completionTestActivity("accepted") + awaitCompletionTest(t, executionStarted) + client.stream.results <- completionTestActivity("rejected") + synctest.Wait() + + stopped := make(chan error, 1) + if stop == "run context" { + cancel() + go func() { stopped <- worker.Wait(context.Background()) }() + } else { + go func() { stopped <- worker.Shutdown(context.Background()) }() + } + awaitCompletionTest(t, client.abandonStarted) + executionOnce.Do(func() { close(releaseExecution) }) + synctest.Wait() + require.EqualValues(t, 1, client.acknowledged.Load()) + require.NoError(t, client.stream.ctx.Err(), "the rejected item's lease must remain alive during abandonment") + require.Zero(t, closer.closes.Load()) + require.True(t, worker.Running()) + + if stop == "shutdown deadline" { + shutdownCtx, stopShutdown := context.WithTimeout(context.Background(), time.Second) + defer stopShutdown() + require.ErrorIs(t, worker.Shutdown(shutdownCtx), context.DeadlineExceeded) + } else { + abandonOnce.Do(func() { close(client.releaseAbandon) }) + } + require.NoError(t, awaitCompletionTest(t, stopped)) + require.EqualValues(t, 1, closer.closes.Load()) + require.Zero(t, client.expired.Load()) + if stop == "shutdown deadline" { + require.Zero(t, client.abandoned.Load()) + } else { + require.EqualValues(t, 1, client.abandoned.Load()) + } + }) + }) + } +} + func TestWorkerGracefulDrainKeepsPendingAcknowledgementAndSilentIntakeAlive(t *testing.T) { client := &leaseAwareSchedulerClient{ fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(1)}, diff --git a/client/grpc_worker_processor.go b/client/grpc_worker_processor.go index 505b676e..f582d29c 100644 --- a/client/grpc_worker_processor.go +++ b/client/grpc_worker_processor.go @@ -191,6 +191,7 @@ func (w *TaskHubGrpcWorker) dispatchEntity( // dispatch reserves a concurrency slot and runs process in the background under // the processing context, so a graceful drain can still complete in-flight work. // If intake is canceled before a slot is free, the work item is abandoned instead. +// Received work remains pending until execution or abandonment returns. func (w *TaskHubGrpcWorker) dispatch( run *grpcWorkerRun, connection *grpcWorkerConnection, @@ -198,9 +199,19 @@ func (w *TaskHubGrpcWorker) dispatch( abandon func(context.Context), process func(context.Context), ) error { + run.dispatchMu.Lock() + run.pending.Add(1) + connection.pending.Add(1) + run.dispatchMu.Unlock() + finish := func() { + connection.pending.Done() + run.pending.Done() + } + select { case slots <- struct{}{}: case <-run.intakeCtx.Done(): + defer finish() abandon(run.processingCtx) return run.intakeCtx.Err() } @@ -209,17 +220,15 @@ func (w *TaskHubGrpcWorker) dispatch( if err := run.intakeCtx.Err(); err != nil { run.dispatchMu.Unlock() <-slots + defer finish() abandon(run.processingCtx) return err } - run.pending.Add(1) - connection.pending.Add(1) run.dispatchMu.Unlock() go func() { defer func() { <-slots - connection.pending.Done() - run.pending.Done() + finish() }() process(run.processingCtx) }() diff --git a/durabletaskscheduler/README.md b/durabletaskscheduler/README.md index cfdf2fa5..ca301823 100644 --- a/durabletaskscheduler/README.md +++ b/durabletaskscheduler/README.md @@ -250,7 +250,9 @@ or abandonment finishes. At most two groups (active and draining) are live: repeated disconnects wait for the older group to drain before creating another. This can delay recovery when old work is slow. Shutdown stops new dispatch but keeps the lease-owning intake stream alive until accepted work and -acknowledgements drain. Silence does not terminate intake after a graceful stop. +acknowledgements drain. Received items still waiting for an execution slot are +abandoned before their intake lease is released. Silence does not terminate +intake after a graceful stop. Only then is the stream canceled and its group retired; a real stream failure can still invalidate leases. The shutdown deadline cancels processing and completion RPCs. As before, From 45c7c41f83d4d739aa4f8bcde5d5b6f663f0c0e0 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:51:10 -0700 Subject: [PATCH 6/7] fence intake cancellation against late work registration Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker_drain_test.go | 64 ++++++++++++++++++++++++++------ client/grpc_worker_processor.go | 22 ++++------- 2 files changed, 60 insertions(+), 26 deletions(-) diff --git a/client/grpc_worker_drain_test.go b/client/grpc_worker_drain_test.go index 3019760d..1b02d575 100644 --- a/client/grpc_worker_drain_test.go +++ b/client/grpc_worker_drain_test.go @@ -76,6 +76,16 @@ func (c *leaseAwareSchedulerClient) AbandonTaskActivityWorkItem( return result, err } +type cancelOnIntakeCheckContext struct { + context.Context + cancel context.CancelFunc +} + +func (c cancelOnIntakeCheckContext) Err() error { + c.cancel() + return c.Context.Err() +} + func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { for _, phase := range []string{"waiting for slot", "slot reserved"} { t.Run(phase, func(t *testing.T) { @@ -85,10 +95,13 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { defer run.cancelIntake() defer run.cancelProcessing() connection := &grpcWorkerConnection{} - slots := make(chan struct{}) + slots := make(chan struct{}, 1) if phase == "waiting for slot" { - slots = make(chan struct{}, 1) slots <- struct{}{} + } else { + // Cancel at the post-admission check, after the only + // ready select case has reserved the free slot. + run.intakeCtx = cancelOnIntakeCheckContext{run.intakeCtx, run.cancelIntake} } started, release := make(chan struct{}), make(chan struct{}) var releaseOnce sync.Once @@ -103,15 +116,7 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { func(context.Context) { t.Error("work was executed after logical intake stop") }) }() synctest.Wait() - if phase == "slot reserved" { - run.dispatchMu.Lock() - <-slots - run.cancelIntake() - run.dispatchMu.Unlock() - // Return the admission token so the rejected dispatch can - // release its slot after observing cancellation. - slots <- struct{}{} - } else { + if phase == "waiting for slot" { run.cancelIntake() } awaitCompletionTest(t, started) @@ -147,6 +152,43 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { } } +func TestWorkerGracefulStopRacingBufferedReceive(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + for range 300 { + client := &fakeSchedulerClient{stream: newFakeWorkItemStream(1)} + worker := newFakeWorker(t, client) + started, release := make(chan struct{}), make(chan struct{}) + worker.executor = &recordingExecutor{ + executeActivity: func(context.Context, api.InstanceID, *protos.HistoryEvent) (*protos.HistoryEvent, error) { + close(started) + <-release + return &protos.HistoryEvent{EventType: &protos.HistoryEvent_TaskCompleted{ + TaskCompleted: &protos.TaskCompletedEvent{}, + }}, nil + }, + } + ctx, cancel := context.WithCancel(context.Background()) + require.NoError(t, worker.Start(ctx)) + client.stream.results <- completionTestActivity("accepted") + awaitCompletionTest(t, started) + cancel() + synctest.Wait() + // Deliver a final buffered item as the previously accepted work + // drains the count to zero and wakes the stop watcher. + close(release) + client.stream.results <- completionTestActivity("raced-intake-stop") + synctest.Wait() + require.NoError(t, worker.Shutdown(context.Background())) + require.False(t, worker.Running()) + client.mu.Lock() + completed, abandoned := len(client.activityCompletions), client.activityAbandons + client.mu.Unlock() + require.Equal(t, 1, completed) + require.LessOrEqual(t, abandoned, 1) + } + }) +} + func TestWorkerGracefulStopKeepsRejectedItemLeaseUntilAbandoned(t *testing.T) { for _, stop := range []string{"shutdown", "run context", "shutdown deadline"} { t.Run(stop, func(t *testing.T) { diff --git a/client/grpc_worker_processor.go b/client/grpc_worker_processor.go index f582d29c..165a0fd5 100644 --- a/client/grpc_worker_processor.go +++ b/client/grpc_worker_processor.go @@ -23,31 +23,24 @@ import ( "google.golang.org/protobuf/types/known/wrapperspb" ) -// awaitDispatchBarrier synchronizes with any dispatch that raced logical intake -// stop, so its pending count is visible before callers wait on it. -func (run *grpcWorkerRun) awaitDispatchBarrier() { - run.dispatchMu.Lock() - defer run.dispatchMu.Unlock() -} - func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *grpcWorkerConnection) (bool, error) { stopWatcher, watcherDone := make(chan struct{}), make(chan struct{}) - waitAccepted := func() { - run.awaitDispatchBarrier() - connection.pending.Wait() - } go func() { defer close(watcherDone) select { case <-run.intakeCtx.Done(): - waitAccepted() + // Fence registration through cancellation: a final receive may + // race the drain, including when the pending count was zero. + run.dispatchMu.Lock() + connection.pending.Wait() connection.cancelStream() + run.dispatchMu.Unlock() case <-stopWatcher: } }() defer func() { if run.intakeCtx.Err() != nil { - waitAccepted() + connection.pending.Wait() } close(stopWatcher) <-watcherDone @@ -202,7 +195,6 @@ func (w *TaskHubGrpcWorker) dispatch( run.dispatchMu.Lock() run.pending.Add(1) connection.pending.Add(1) - run.dispatchMu.Unlock() finish := func() { connection.pending.Done() run.pending.Done() @@ -211,12 +203,12 @@ func (w *TaskHubGrpcWorker) dispatch( select { case slots <- struct{}{}: case <-run.intakeCtx.Done(): + run.dispatchMu.Unlock() defer finish() abandon(run.processingCtx) return run.intakeCtx.Err() } - run.dispatchMu.Lock() if err := run.intakeCtx.Err(); err != nil { run.dispatchMu.Unlock() <-slots From 1d83a1f3170ad0f2552f4b3b68f73d877c6188c9 Mon Sep 17 00:00:00 2001 From: Tomer Rosenthal <17064840+torosent@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:16:54 -0700 Subject: [PATCH 7/7] register received work before handing it to dispatch Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- client/grpc_worker.go | 52 ++++++++++++++++- client/grpc_worker_drain_test.go | 97 ++++++++++++++++++++++++++++++-- client/grpc_worker_processor.go | 39 ++++++------- durabletaskscheduler/README.md | 4 +- 4 files changed, 162 insertions(+), 30 deletions(-) diff --git a/client/grpc_worker.go b/client/grpc_worker.go index 34582b8e..ea367103 100644 --- a/client/grpc_worker.go +++ b/client/grpc_worker.go @@ -537,7 +537,6 @@ type grpcWorkerRun struct { orchestrationSlots chan struct{} activitySlots chan struct{} entitySlots chan struct{} - dispatchMu sync.Mutex pending sync.WaitGroup retired sync.WaitGroup err error @@ -549,7 +548,54 @@ type grpcWorkerConnection struct { cancelStream context.CancelFunc closer io.Closer - pending sync.WaitGroup + pendingMu sync.Mutex + pendingChanged *sync.Cond + pendingWorkItems int + intakeClosed bool +} + +func (c *grpcWorkerConnection) registerWorkItem(run *grpcWorkerRun) bool { + c.pendingMu.Lock() + defer c.pendingMu.Unlock() + if c.intakeClosed { + return false + } + c.pendingWorkItems++ + run.pending.Add(1) + return true +} + +func (c *grpcWorkerConnection) finishWorkItem(run *grpcWorkerRun) { + c.pendingMu.Lock() + c.pendingWorkItems-- + if c.pendingWorkItems == 0 && c.pendingChanged != nil { + c.pendingChanged.Broadcast() + } + c.pendingMu.Unlock() + run.pending.Done() +} + +func (c *grpcWorkerConnection) waitForWorkItemsLocked() { + if c.pendingWorkItems > 0 && c.pendingChanged == nil { + c.pendingChanged = sync.NewCond(&c.pendingMu) + } + for c.pendingWorkItems > 0 { + c.pendingChanged.Wait() + } +} + +func (c *grpcWorkerConnection) waitForWorkItems() { + c.pendingMu.Lock() + defer c.pendingMu.Unlock() + c.waitForWorkItemsLocked() +} + +func (c *grpcWorkerConnection) stopIntake() { + c.pendingMu.Lock() + defer c.pendingMu.Unlock() + c.waitForWorkItemsLocked() + c.intakeClosed = true + c.cancelStream() } // NewTaskHubGrpcWorker borrows intake and requires separate caller-owned @@ -1352,7 +1398,7 @@ func (w *TaskHubGrpcWorker) retireConnection(run *grpcWorkerRun, connection *grp go func() { defer run.retired.Done() defer close(done) - connection.pending.Wait() + connection.waitForWorkItems() if connection.closer != nil { if err := connection.closer.Close(); err != nil { w.logger.Warnf("failed to close retired gRPC worker connection: %v", err) diff --git a/client/grpc_worker_drain_test.go b/client/grpc_worker_drain_test.go index 1b02d575..c59fa49c 100644 --- a/client/grpc_worker_drain_test.go +++ b/client/grpc_worker_drain_test.go @@ -95,6 +95,7 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { defer run.cancelIntake() defer run.cancelProcessing() connection := &grpcWorkerConnection{} + require.True(t, connection.registerWorkItem(run)) slots := make(chan struct{}, 1) if phase == "waiting for slot" { slots <- struct{}{} @@ -127,7 +128,7 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { close(runDrained) }() go func() { - connection.pending.Wait() + connection.waitForWorkItems() close(connectionDrained) }() synctest.Wait() @@ -155,8 +156,11 @@ func TestWorkerCanceledDispatchWaitsForAbandonment(t *testing.T) { func TestWorkerGracefulStopRacingBufferedReceive(t *testing.T) { synctest.Test(t, func(t *testing.T) { for range 300 { - client := &fakeSchedulerClient{stream: newFakeWorkItemStream(1)} - worker := newFakeWorker(t, client) + client := &leaseAwareSchedulerClient{fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(1)}} + worker := newFakeWorker(t, client.fakeSchedulerClient) + worker.clientFactory = func(context.Context) (protos.TaskHubSidecarServiceClient, io.Closer, error) { + return client, nil, nil + } started, release := make(chan struct{}), make(chan struct{}) worker.executor = &recordingExecutor{ executeActivity: func(context.Context, api.InstanceID, *protos.HistoryEvent) (*protos.HistoryEvent, error) { @@ -185,10 +189,93 @@ func TestWorkerGracefulStopRacingBufferedReceive(t *testing.T) { client.mu.Unlock() require.Equal(t, 1, completed) require.LessOrEqual(t, abandoned, 1) + require.Zero(t, client.expired.Load(), "no acknowledgement may use a lease released by shutdown") } }) } +func TestWorkerReceiveHandoffPreservesLeaseUntilAbandoned(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + client := &leaseAwareSchedulerClient{ + fakeSchedulerClient: &fakeSchedulerClient{stream: newFakeWorkItemStream(1)}, + abandonStarted: make(chan struct{}), + releaseAbandon: make(chan struct{}), + } + run := newRunLoopTestRun() + defer run.cancelIntake() + defer run.cancelProcessing() + streamCtx, cancelStream := context.WithCancel(context.Background()) + defer cancelStream() + client.stream.ctx = streamCtx + connection := &grpcWorkerConnection{client: client, stream: client.stream, cancelStream: cancelStream} + worker := newFakeWorker(t, client.fakeSchedulerClient) + var releaseOnce sync.Once + dispatched := false + t.Cleanup(func() { + releaseOnce.Do(func() { close(client.releaseAbandon) }) + if !dispatched { + connection.pendingMu.Lock() + registered := connection.pendingWorkItems > 0 + connection.pendingMu.Unlock() + if registered { + connection.finishWorkItem(run) + } + } + }) + client.stream.results <- completionTestActivity("receive-handoff") + item, err := worker.receiveWorkItem(run, connection) + require.NoError(t, err) + + // Pause after successful receipt, before entering dispatch at all. + run.cancelIntake() + stopped := make(chan struct{}) + go func() { + connection.stopIntake() + close(stopped) + }() + synctest.Wait() + require.NoError(t, streamCtx.Err(), "receipt must already be registered before dispatch") + + dispatchDone := make(chan error, 1) + dispatched = true + go func() { + dispatchDone <- worker.dispatchActivity(run, connection, item.GetCompletionToken(), item.GetActivityRequest()) + }() + awaitCompletionTest(t, client.abandonStarted) + synctest.Wait() + require.NoError(t, streamCtx.Err(), "the abandonment must retain the received item's lease") + releaseOnce.Do(func() { close(client.releaseAbandon) }) + require.ErrorIs(t, awaitCompletionTest(t, dispatchDone), context.Canceled) + awaitCompletionTest(t, stopped) + run.pending.Wait() + require.EqualValues(t, 1, client.abandoned.Load()) + require.Zero(t, client.expired.Load()) + }) +} + +func TestWorkerStoppedIntakeRejectsUnregisteredDelivery(t *testing.T) { + run := newRunLoopTestRun() + defer run.cancelIntake() + defer run.cancelProcessing() + streamCtx, cancelStream := context.WithCancel(context.Background()) + defer cancelStream() + stream := &racingStream{ctx: streamCtx, item: completionTestActivity("late-delivery").item} + client := &fakeSchedulerClient{} + connection := &grpcWorkerConnection{client: client, stream: stream, cancelStream: cancelStream} + worker := newFakeWorker(t, client) + run.cancelIntake() + connection.stopIntake() + + // Recv returns buffered work even though cancellation already released its lease. + item, err := worker.receiveWorkItem(run, connection) + require.Equal(t, codes.Canceled, status.Code(err)) + require.Nil(t, item) + connection.waitForWorkItems() + run.pending.Wait() + require.Empty(t, client.activityCompletions) + require.Zero(t, client.activityAbandons) +} + func TestWorkerGracefulStopKeepsRejectedItemLeaseUntilAbandoned(t *testing.T) { for _, stop := range []string{"shutdown", "run context", "shutdown deadline"} { t.Run(stop, func(t *testing.T) { @@ -317,8 +404,10 @@ func TestWorkerStoppedDispatchDoesNotExecuteEvenWithFreeSlot(t *testing.T) { for range 100 { run := newRunLoopTestRun() run.cancelIntake() + connection := &grpcWorkerConnection{} + require.True(t, connection.registerWorkItem(run)) var abandoned bool - err := worker.dispatch(run, &grpcWorkerConnection{}, run.activitySlots, + err := worker.dispatch(run, connection, run.activitySlots, func(context.Context) { abandoned = true }, func(context.Context) { t.Error("work was executed after logical intake stop") }) require.ErrorIs(t, err, context.Canceled) diff --git a/client/grpc_worker_processor.go b/client/grpc_worker_processor.go index 165a0fd5..92675571 100644 --- a/client/grpc_worker_processor.go +++ b/client/grpc_worker_processor.go @@ -29,18 +29,13 @@ func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *gr defer close(watcherDone) select { case <-run.intakeCtx.Done(): - // Fence registration through cancellation: a final receive may - // race the drain, including when the pending count was zero. - run.dispatchMu.Lock() - connection.pending.Wait() - connection.cancelStream() - run.dispatchMu.Unlock() + connection.stopIntake() case <-stopWatcher: } }() defer func() { if run.intakeCtx.Err() != nil { - connection.pending.Wait() + connection.waitForWorkItems() } close(stopWatcher) <-watcherDone @@ -48,6 +43,9 @@ func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *gr observedMessage := false for { + if err := run.intakeCtx.Err(); err != nil { + return observedMessage, err + } workItem, err := w.receiveWorkItem(run, connection) if err != nil { return observedMessage, err @@ -56,6 +54,7 @@ func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *gr switch request := workItem.Request.(type) { case *protos.WorkItem_HealthPing: + connection.finishWorkItem(run) continue case *protos.WorkItem_OrchestratorRequest: if err := w.dispatchOrchestration(run, connection, workItem.GetCompletionToken(), request.OrchestratorRequest); err != nil { @@ -80,6 +79,7 @@ func (w *TaskHubGrpcWorker) consumeConnection(run *grpcWorkerRun, connection *gr } default: w.logger.Warnf("received unknown work item type with completion token present=%t", workItem.GetCompletionToken() != "") + connection.finishWorkItem(run) } } } @@ -99,6 +99,12 @@ func (w *TaskHubGrpcWorker) receiveWorkItem(run *grpcWorkerRun, connection *grpc if workItem == nil { return nil, status.Error(codes.Internal, "received a nil work item") } + if !connection.registerWorkItem(run) { + // Cancellation won the receipt boundary; closing intake already + // releases this unregistered delivery's lease on the server. + w.logger.Debug("work item delivery raced closed intake; its lease was released with the stream") + return nil, status.Error(codes.Canceled, "worker intake closed before work item registration") + } return workItem, nil } @@ -184,7 +190,7 @@ func (w *TaskHubGrpcWorker) dispatchEntity( // dispatch reserves a concurrency slot and runs process in the background under // the processing context, so a graceful drain can still complete in-flight work. // If intake is canceled before a slot is free, the work item is abandoned instead. -// Received work remains pending until execution or abandonment returns. +// The receive path registers work before exposing it to dispatch. func (w *TaskHubGrpcWorker) dispatch( run *grpcWorkerRun, connection *grpcWorkerConnection, @@ -192,35 +198,24 @@ func (w *TaskHubGrpcWorker) dispatch( abandon func(context.Context), process func(context.Context), ) error { - run.dispatchMu.Lock() - run.pending.Add(1) - connection.pending.Add(1) - finish := func() { - connection.pending.Done() - run.pending.Done() - } - select { case slots <- struct{}{}: case <-run.intakeCtx.Done(): - run.dispatchMu.Unlock() - defer finish() + defer connection.finishWorkItem(run) abandon(run.processingCtx) return run.intakeCtx.Err() } if err := run.intakeCtx.Err(); err != nil { - run.dispatchMu.Unlock() <-slots - defer finish() + defer connection.finishWorkItem(run) abandon(run.processingCtx) return err } - run.dispatchMu.Unlock() go func() { defer func() { <-slots - finish() + connection.finishWorkItem(run) }() process(run.processingCtx) }() diff --git a/durabletaskscheduler/README.md b/durabletaskscheduler/README.md index ca301823..82fa765e 100644 --- a/durabletaskscheduler/README.md +++ b/durabletaskscheduler/README.md @@ -252,7 +252,9 @@ This can delay recovery when old work is slow. Shutdown stops new dispatch but keeps the lease-owning intake stream alive until accepted work and acknowledgements drain. Received items still waiting for an execution slot are abandoned before their intake lease is released. Silence does not terminate -intake after a graceful stop. +intake after a graceful stop. Receipt is registered before work is passed to +dispatch. A delivery racing closed intake is not dispatched or abandoned again: +closing the stream releases its lease on the server. Only then is the stream canceled and its group retired; a real stream failure can still invalidate leases. The shutdown deadline cancels processing and completion RPCs. As before,