From 47b48c8503658cef1d22efb066b00b89fe86309c Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Thu, 4 Jun 2026 15:38:36 +0100 Subject: [PATCH 01/35] Add sharded force replication workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a parallel ShardedForceReplicationWorkflow alongside the existing ForceReplicationWorkflow variants. The new variant routes each execution by destination history shard, packs batches across shards, and gates on per-shard exclusivity so a single hot shard can't dominate the apply pipeline. Defaults are unchanged — the legacy workflow stays the default; callers opt into the sharded variant by starting a workflow of type "force-replication-sharded" on the new MigrationShardedActivityTQ. Wired as a second WorkerComponent in migration.Module: dedicated workflow + activity workers polling primitives.MigrationShardedActivityTQ. The shared *activities struct picks up adminClient (from a local admin client) and sdkClientFactory fields so the new ReplicateBatch activity can drive both the via-frontend inject path and the mid-flight ReleaseShards signal without relying on per-workflow plumbing. Side effect: the legacy *activities builder on main never populated adminClient, so the existing generateMigrationTaskViaFrontend code path (activities.go GenerateLastHistoryReplicationTasks call) would have nil-pointer'd if that dynamic-config flag were enabled. The new ClientBean-based wiring in newActivitiesFromParams populates adminClient for both the legacy and sharded *activities instances, fixing that latent NPE for the legacy via-frontend path as well. --- common/metrics/metric_defs.go | 21 + common/primitives/task_queues.go | 1 + service/worker/migration/activities.go | 7 + service/worker/migration/fx.go | 140 ++- service/worker/migration/sharded_activity.go | 640 ++++++++++++ service/worker/migration/sharded_types.go | 374 +++++++ .../worker/migration/sharded_types_test.go | 91 ++ service/worker/migration/sharded_workflow.go | 915 ++++++++++++++++++ .../worker/migration/sharded_workflow_test.go | 490 ++++++++++ 9 files changed, 2656 insertions(+), 23 deletions(-) create mode 100644 service/worker/migration/sharded_activity.go create mode 100644 service/worker/migration/sharded_types.go create mode 100644 service/worker/migration/sharded_types_test.go create mode 100644 service/worker/migration/sharded_workflow.go create mode 100644 service/worker/migration/sharded_workflow_test.go diff --git a/common/metrics/metric_defs.go b/common/metrics/metric_defs.go index 1de307e9cff..3cff13bc40d 100644 --- a/common/metrics/metric_defs.go +++ b/common/metrics/metric_defs.go @@ -1487,6 +1487,27 @@ var ( VerifyReplicationTasksLatency = NewTimerDef("verify_replication_tasks_latency") VerifyDescribeMutableStateLatency = NewTimerDef("verify_describe_mutable_state_latency") + // Sharded force replication. The sharded ReplicateBatch activity runs + // many executions per invocation, so per-exec timing is the meaningful + // granularity — the batch-level *_tasks_latency timers above scale with + // BatchSize and aren't comparable across configurations. + GenerateReplicationTaskLatency = NewTimerDef("generate_replication_task_latency") + VerifyReplicationTaskLatency = NewTimerDef("verify_replication_task_latency") + // VerifyReplicationTaskBusy counts verify attempts where the passive + // cluster returned RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW — the cache + // lock is held while history is being applied. A sign of progress that + // doesn't reset the per-shard no-progress timer. + VerifyReplicationTaskBusy = NewCounterDef("verify_replication_task_busy") + // VerifyReplicationTaskPending counts verify attempts where + // DescribeMutableState succeeded but the workflowVerifier saw the target + // lagging the source. A high pending vs. success ratio means verify is + // polling faster than apply can catch up. + VerifyReplicationTaskPending = NewCounterDef("verify_replication_task_pending") + // ReplicatedWorkflowCount accumulates verified-exec counts across each + // ReplicateBatch activity return. Emitted from the workflow so the + // counter is monotonic across activity retries. + ReplicatedWorkflowCount = NewCounterDef("replicated_workflow_count") + // Replication NamespaceReplicationTaskAckLevelGauge = NewGaugeDef("namespace_replication_task_ack_level") NamespaceReplicationDLQAckLevelGauge = NewGaugeDef("namespace_dlq_ack_level") diff --git a/common/primitives/task_queues.go b/common/primitives/task_queues.go index 0c6c049f11a..0c6dd2dab59 100644 --- a/common/primitives/task_queues.go +++ b/common/primitives/task_queues.go @@ -17,6 +17,7 @@ const ( internalTaskQueuePerNSPrefix = "temporal-sys-per-ns-" MigrationActivityTQ = "temporal-sys-migration-activity-tq" + MigrationShardedActivityTQ = "temporal-sys-migration-sharded-activity-tq" AddSearchAttributesActivityTQ = "temporal-sys-add-search-attributes-activity-tq" DeleteNamespaceActivityTQ = "temporal-sys-delete-namespace-activity-tq" DLQActivityTQ = "temporal-sys-dlq-activity-tq" diff --git a/service/worker/migration/activities.go b/service/worker/migration/activities.go index 22173be166e..0818670daaa 100644 --- a/service/worker/migration/activities.go +++ b/service/worker/migration/activities.go @@ -31,6 +31,7 @@ import ( "go.temporal.io/server/common/persistence" "go.temporal.io/server/common/quotas" "go.temporal.io/server/common/rpc/interceptor" + "go.temporal.io/server/common/sdk" workercommon "go.temporal.io/server/service/worker/common" "google.golang.org/grpc/metadata" ) @@ -129,6 +130,12 @@ type ( enableHistoryRateLimiter dynamicconfig.BoolPropertyFn workflowVerifier WorkflowVerifier chasmRegistry *chasm.Registry + // sdkClientFactory resolves the system SDK client lazily for the + // sharded ReplicateBatch activity's mid-flight ReleaseShards signal. + // Eager resolution at fx-wire time tries to dial the frontend before + // it's listening; the factory's internal sync.Once guarantees a + // single dial on first use. + sdkClientFactory sdk.ClientFactory } shardStatus struct { diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index c1eda546ef3..dc3021ce08c 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -2,6 +2,7 @@ package migration import ( "context" + "fmt" "go.temporal.io/api/workflowservice/v1" sdkworker "go.temporal.io/sdk/worker" @@ -9,6 +10,7 @@ import ( "go.temporal.io/server/api/adminservice/v1" "go.temporal.io/server/chasm" serverClient "go.temporal.io/server/client" + "go.temporal.io/server/common/cluster" "go.temporal.io/server/common/config" "go.temporal.io/server/common/dynamicconfig" "go.temporal.io/server/common/headers" @@ -18,6 +20,7 @@ import ( "go.temporal.io/server/common/persistence" "go.temporal.io/server/common/primitives" "go.temporal.io/server/common/resource" + "go.temporal.io/server/common/sdk" workercommon "go.temporal.io/server/service/worker/common" "go.uber.org/fx" ) @@ -32,6 +35,7 @@ type ( FrontendClient workflowservice.WorkflowServiceClient ClientFactory serverClient.Factory ClientBean serverClient.Bean + ClusterMetadata cluster.Metadata NamespaceReplicationQueue persistence.NamespaceReplicationQueue TaskManager persistence.TaskManager Logger log.Logger @@ -39,6 +43,7 @@ type ( DynamicCollection *dynamicconfig.Collection WorkflowVerifier WorkflowVerifier ChasmRegistry *chasm.Registry + SDKClientFactory sdk.ClientFactory } fxResult struct { @@ -48,21 +53,49 @@ type ( replicationWorkerComponent struct { initParams + activities *activities + } + + // shardedWorkerComponent registers the sharded force-replication + // workflow + ReplicateBatch activity on their dedicated TQ. Holds an + // *activities-sized clone so its activity registration is isolated + // from the default-TQ worker — sharded inject paths don't accidentally + // land on the legacy MigrationActivityTQ. + shardedWorkerComponent struct { + activities *activities } ) var Module = fx.Options( fx.Provide(NewResult), + fx.Provide(NewShardedResult), fx.Provide(workflowVerifierProvider), ) -func NewResult(params initParams) fxResult { - component := &replicationWorkerComponent{ - initParams: params, +func NewResult(params initParams) (fxResult, error) { + a, err := newActivitiesFromParams(params, forceReplicationWorkflowName) + if err != nil { + return fxResult{}, err } return fxResult{ - Component: component, + Component: &replicationWorkerComponent{ + initParams: params, + activities: a, + }, + }, nil +} + +// NewShardedResult constructs the sharded WorkerComponent. The component +// owns its own *activities clone so registration against the sharded TQ +// doesn't bleed into the legacy worker. +func NewShardedResult(params initParams) (fxResult, error) { + a, err := newActivitiesFromParams(params, shardedForceReplicationWorkflowName) + if err != nil { + return fxResult{}, err } + return fxResult{ + Component: &shardedWorkerComponent{activities: a}, + }, nil } func (wc *replicationWorkerComponent) RegisterWorkflow(registry sdkworker.Registry) { @@ -80,7 +113,7 @@ func (wc *replicationWorkerComponent) DedicatedWorkflowWorkerOptions() *workerco } func (wc *replicationWorkerComponent) RegisterActivities(registry sdkworker.Registry) { - registry.RegisterActivity(wc.activities()) + registry.RegisterActivity(wc.activities) } func (wc *replicationWorkerComponent) DedicatedActivityWorkerOptions() *workercommon.DedicatedWorkerOptions { @@ -92,6 +125,48 @@ func (wc *replicationWorkerComponent) DedicatedActivityWorkerOptions() *workerco } } +func (sc *shardedWorkerComponent) RegisterWorkflow(registry sdkworker.Registry) { + registry.RegisterWorkflowWithOptions(ShardedForceReplicationWorkflow, workflow.RegisterOptions{ + Name: shardedForceReplicationWorkflowName, + }) +} + +func (sc *shardedWorkerComponent) DedicatedWorkflowWorkerOptions() *workercommon.DedicatedWorkerOptions { + // Workflow + activity share the same TQ so the workflow's default + // ExecuteActivity (no explicit TaskQueue) routes to our dedicated + // activity worker rather than the default-TQ worker. Without a + // dedicated workflow worker here the workflow would land on + // default-worker-tq and its activities would pile up on the (separate, + // our-TQ) dedicated activity worker, unscheduled. + // + // LocalActivityWorkerOnly is essential: by default a worker polls for + // both workflow and activity tasks on its TQ. Since the activity + // worker (a separate sdkworker.Worker) also polls this TQ and is the + // one that owns the registered activities, leaving activity polling + // enabled here means this worker races for activity tasks and + // dispatches them with no registrations — ActivityNotRegisteredError, + // "Supported types: []". + return &workercommon.DedicatedWorkerOptions{ + TaskQueue: primitives.MigrationShardedActivityTQ, + Options: sdkworker.Options{ + LocalActivityWorkerOnly: true, + }, + } +} + +func (sc *shardedWorkerComponent) RegisterActivities(registry sdkworker.Registry) { + registry.RegisterActivity(sc.activities) +} + +func (sc *shardedWorkerComponent) DedicatedActivityWorkerOptions() *workercommon.DedicatedWorkerOptions { + return &workercommon.DedicatedWorkerOptions{ + TaskQueue: primitives.MigrationShardedActivityTQ, + Options: sdkworker.Options{ + BackgroundActivityContext: headers.SetCallerType(context.Background(), headers.CallerTypePreemptable), + }, + } +} + func workflowVerifierProvider() WorkflowVerifier { return func( _ context.Context, @@ -108,23 +183,42 @@ func workflowVerifierProvider() WorkflowVerifier { } } -func (wc *replicationWorkerComponent) activities() *activities { - return &activities{ - HistoryShardCount: wc.PersistenceConfig.NumHistoryShards, - executionManager: wc.ExecutionManager, - NamespaceRegistry: wc.NamespaceRegistry, - HistoryClient: wc.HistoryClient, - frontendClient: wc.FrontendClient, - clientFactory: wc.ClientFactory, - clientBean: wc.ClientBean, - namespaceReplicationQueue: wc.NamespaceReplicationQueue, - taskManager: wc.TaskManager, - Logger: wc.Logger, - MetricsHandler: wc.MetricsHandler, - forceReplicationMetricsHandler: wc.MetricsHandler.WithTags(metrics.WorkflowTypeTag(forceReplicationWorkflowName)), - generateMigrationTaskViaFrontend: dynamicconfig.WorkerGenerateMigrationTaskViaFrontend.Get(wc.DynamicCollection), - enableHistoryRateLimiter: dynamicconfig.WorkerEnableHistoryRateLimiter.Get(wc.DynamicCollection), - workflowVerifier: wc.WorkflowVerifier, - chasmRegistry: wc.ChasmRegistry, +// newActivitiesFromParams builds the shared *activities struct from the +// fx params. workflowTypeName tags the forceReplicationMetricsHandler so +// the legacy and sharded variants emit force-replication metrics under +// distinct workflow_type tags. +// +// adminClient is the local admin client cached by ClientBean at startup. +// Routing through the bean (rather than constructing a fresh wrapper via +// NewLocalAdminClientWithTimeout) reuses the same retry+metric wrapper +// every other consumer in the process sees, and guarantees adminClient +// is non-nil so the inject and verify paths can use it without nil +// guarding. A lookup failure indicates ClusterMetadata is misconfigured; +// surfacing it as an fx error fails app start cleanly rather than mid-run. +func newActivitiesFromParams(params initParams, workflowTypeName string) (*activities, error) { + localCluster := params.ClusterMetadata.GetCurrentClusterName() + localAdmin, err := params.ClientBean.GetRemoteAdminClient(localCluster) + if err != nil { + return nil, fmt.Errorf("migration: local admin client missing from ClientBean for cluster %q: %w", localCluster, err) } + return &activities{ + HistoryShardCount: params.PersistenceConfig.NumHistoryShards, + executionManager: params.ExecutionManager, + NamespaceRegistry: params.NamespaceRegistry, + HistoryClient: params.HistoryClient, + frontendClient: params.FrontendClient, + adminClient: localAdmin, + clientFactory: params.ClientFactory, + clientBean: params.ClientBean, + namespaceReplicationQueue: params.NamespaceReplicationQueue, + taskManager: params.TaskManager, + Logger: params.Logger, + MetricsHandler: params.MetricsHandler, + forceReplicationMetricsHandler: params.MetricsHandler.WithTags(metrics.WorkflowTypeTag(workflowTypeName)), + generateMigrationTaskViaFrontend: dynamicconfig.WorkerGenerateMigrationTaskViaFrontend.Get(params.DynamicCollection), + enableHistoryRateLimiter: dynamicconfig.WorkerEnableHistoryRateLimiter.Get(params.DynamicCollection), + workflowVerifier: params.WorkflowVerifier, + chasmRegistry: params.ChasmRegistry, + sdkClientFactory: params.SDKClientFactory, + }, nil } diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go new file mode 100644 index 00000000000..06b901dacec --- /dev/null +++ b/service/worker/migration/sharded_activity.go @@ -0,0 +1,640 @@ +package migration + +import ( + "context" + "errors" + "fmt" + "math" + "slices" + "time" + + commonpb "go.temporal.io/api/common/v1" + enumspb "go.temporal.io/api/enums/v1" + "go.temporal.io/api/serviceerror" + "go.temporal.io/sdk/activity" + "go.temporal.io/sdk/temporal" + "go.temporal.io/server/api/adminservice/v1" + "go.temporal.io/server/client/admin" + "go.temporal.io/server/common/metrics" + "go.temporal.io/server/common/namespace" + "go.temporal.io/server/common/quotas" +) + +// ReplicateBatch is the per-batch activity body for the sharded force +// replication workflow. It runs inject (skip on Resume) followed by +// verify, signal-releasing completed shards mid-flight once their +// cumulative idle cost crosses IdleShardCost. On workflow-initiated +// cancellation it enters drain mode: continues verifying for up to +// DrainGrace, then returns a replicateBatchResult carrying any +// still-unverified execs grouped by shard. Drain-mode signal traffic +// is intentionally suppressed — once we know we're about to return, +// there's no point racing a signal against the return value. +func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + // Flatten the nested wire shape once on entry so the per-exec + // bookkeeping (verified[], attempts[], nextRetryAt[]) can stay + // index-based. flatten() walks shards ascending then BIDs + // alphabetical, so the order is deterministic — replays of the same + // payload produce the same slice. + execs := req.Executions.flatten() + of := len(execs) + if of == 0 { + return replicateBatchResult{}, nil + } + + remoteAdminClient := a.clientFactory.NewRemoteAdminClientWithTimeout( + req.TargetClusterEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout) + + // Namespace lookup feeds the verify phase's retention/zombie skip + // check (checkSkipWorkflowExecution needs ns.Retention()). Looked up + // once per activity, since the namespace registry is local and + // immutable across the run. + ns, err := a.NamespaceRegistry.GetNamespaceByID(namespace.ID(req.NamespaceID)) + if err != nil { + return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) + } + + // ---- Inject phase (skipped on Resume: the execs have already been + // injected by an earlier activity that was drained for CAN). + if !req.Resume { + // One limiter per activity invocation, sized for RPS — the + // post-call ReserveN reservation pulls extra tokens proportional + // to history size, so a single limiter shared across this + // batch's execs is what enforces the per-batch RPS budget + // end-to-end. + rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) + for _, ex := range execs { + if err := ctx.Err(); err != nil { + // Worker shutdown or workflow drain mid-inject. Return a + // recognizable CanceledError so spawnBatch's + // IsCanceledError check fires and the batch's batchExecs + // entry is preserved for RecoveredBuckets re-injection + // next cycle. Any execs already injected here are + // re-injected harmlessly — replication dedupes per + // (namespace, wf, run). + return replicateBatchResult{}, temporal.NewCanceledError("inject phase cancelled") + } + if err := a.generateReplicationTaskForExec(ctx, rateLimiter, req, ex); err != nil { + if ctx.Err() != nil { + return replicateBatchResult{}, temporal.NewCanceledError("inject phase cancelled") + } + return replicateBatchResult{}, err + } + } + } + + // ---- Verify phase ---- + verified := make([]bool, of) + attempts := make([]int, of) + nextRetryAt := make([]time.Time, of) + doneCount := 0 + + shards := newShardVerifyTracker(execs, req.Resume, req.NoProgressByShard) + + var draining bool + var drainStartAt time.Time + + // callCtx is what attemptVerifyExec uses for DescribeMutableState. In + // normal mode it's the activity's parent ctx; when drain mode kicks + // in, it swaps to a fresh detached context with a DrainGrace timeout + // so DMS calls keep working long enough for nearly-verified execs to + // land. The parent ctx is already dead by the time we transition + // (that's what triggers the transition), so using it for drain-mode + // RPCs would defeat the grace window — every call would fail + // instantly. + callCtx := ctx + // Pre-create the drain context up front and defer cancel right away + // so go vet's lostcancel pass sees the canonical pattern. The drain + // timer (started below on the cancel transition) plus the unconditional + // defer cancel make the lifetime obvious; the context only matters + // once draining is set, so creating it early costs nothing. + drainCtx, drainCancel := context.WithCancel(context.Background()) + defer drainCancel() + + for { + // Worker shutdown short-circuits drain mode entirely. The SDK + // closes activity.GetWorkerStopChannel a fixed WorkerStopTimeout + // (10s default) before forcibly returning; burning that window + // on DescribeMutableState calls that won't get to drive their + // results back is worse than returning current state and letting + // the next cycle's ResumeShards / RecoveredBuckets paths recover. + // + // Checked at the top of each outer iteration (not just on initial + // drain transition) because worker shutdown can fire after we've + // already entered workflow-initiated drain — e.g. CAN cancel + // arrives, drain starts under detached ctx, then a deploy hits + // mid-window. The detached ctx wouldn't notice on its own. + select { + case <-activity.GetWorkerStopChannel(ctx): + return replicateBatchResult{ + CompletedShards: shards.allCompleted(), + InFlight: a.buildInFlight(execs, verified, shards, time.Now()), + VerifiedCount: int64(doneCount), + }, nil + default: + } + + // Workflow-initiated activity cancellation (drainForCAN) + // transitions us into drain mode with the full DrainGrace window. + // WaitForCancellation=true on the activity options guarantees the + // workflow blocks for us, so the grace window is genuinely + // available — swap callCtx onto a detached deadline so + // DescribeMutableState keeps working after the parent ctx died. + if !draining { + if err := ctx.Err(); err != nil { + draining = true + drainStartAt = time.Now() + // Start the drain budget timer here rather than at activity + // entry so the grace window measures from drain transition, + // not from activity start. + time.AfterFunc(req.DrainGrace, drainCancel) + callCtx = drainCtx + } + } + + now := time.Now() + var minNextRetry time.Time + ctxAborted := false + for i := range of { + if verified[i] { + continue + } + if !nextRetryAt[i].IsZero() && nextRetryAt[i].After(now) { + if minNextRetry.IsZero() || nextRetryAt[i].Before(minNextRetry) { + minNextRetry = nextRetryAt[i] + } + continue + } + + ex := execs[i] + shard := ex.Shard + ok, err := a.attemptVerifyExec(callCtx, remoteAdminClient, ns, req, ex) + if err != nil { + if callCtx.Err() != nil { + // callCtx is dead. Two cases: (1) normal-mode parent + // ctx was just cancelled mid-call — the outer-loop top + // will promote to drain on the next iteration; + // (2) drain-mode detached ctx expired — the drain + // exit check below will fire. Either way, drop out of + // the inner loop now. + ctxAborted = true + break + } + return replicateBatchResult{}, err + } + + if ok { + verified[i] = true + doneCount++ + shards.recordVerified(shard, time.Now()) + continue + } + + attempts[i]++ + nextRetryAt[i] = time.Now().Add(backoffDelay(attempts[i])) + if minNextRetry.IsZero() || nextRetryAt[i].Before(minNextRetry) { + minNextRetry = nextRetryAt[i] + } + } + + activity.RecordHeartbeat(ctx, doneCount) + + // Clean completion — every exec verified. + if doneCount >= of { + return replicateBatchResult{ + CompletedShards: shards.allCompleted(), + VerifiedCount: int64(doneCount), + }, nil + } + + // Per-shard cumulative no-progress backstop. Trips on any shard + // still holding pending execs whose last verified outcome + // (carried across CAN via tracker seeding) is older than + // ShardNoProgress. + if stuckShard, stuckDur, ok := shards.pickStuck(time.Now(), req.ShardNoProgress); ok { + stuckIdx := a.firstUnverifiedOnShard(execs, verified, stuckShard) + stuck := execs[stuckIdx] + return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( + fmt.Sprintf("shard %d no progress for %v on %s/%s (%d/%d done)", + stuckShard, stuckDur, stuck.BusinessID, stuck.RunID, doneCount, of), + "ShardNoProgress", nil) + } + + // Drain-mode exit checks. Drain mode is entered when the workflow + // has cancelled the activity for CAN. No signals here — the + // return value carries everything the workflow needs (completed + // shards + unverified execs grouped by shard with their cumulative + // no-progress duration). + if draining { + elapsedDrain := time.Since(drainStartAt) + if elapsedDrain >= req.DrainGrace || shards.totalIdleCost(time.Now()) >= req.IdleShardCost { + return replicateBatchResult{ + CompletedShards: shards.allCompleted(), + InFlight: a.buildInFlight(execs, verified, shards, time.Now()), + VerifiedCount: int64(doneCount), + }, nil + } + } else { + // Normal mode: if the cumulative idle cost across + // completed-but-not-yet-signaled shards crosses the threshold, + // signal-release them so the workflow can dispatch new batches + // against those shards while this activity keeps draining its + // still-pending ones. + if shards.totalIdleCost(time.Now()) >= req.IdleShardCost { + releaseList := shards.awaitingRelease() + if len(releaseList) > 0 { + if err := a.signalReleaseShards(ctx, req, releaseList); err != nil { + return replicateBatchResult{}, err + } + shards.markReleased(releaseList) + } + } + } + + // If the inner loop aborted because callCtx died, skip the sleep + // entirely so the outer-loop top sees the new state promptly + // (normal → drain transition, or drain → exit). + if ctxAborted { + continue + } + + // Sleep until the next exec is due for retry. Drain mode honours + // the same scheduling so we don't burn the grace window + // busy-spinning when every remaining exec is in backoff. + sleepDur := 50 * time.Millisecond + if !minNextRetry.IsZero() { + if delta := time.Until(minNextRetry); delta > sleepDur { + sleepDur = delta + } + } + if draining { + remaining := req.DrainGrace - time.Since(drainStartAt) + if remaining > 0 && remaining < sleepDur { + sleepDur = remaining + } + // Parent ctx is already dead in drain mode, so the + // select-on-ctx.Done() in normal mode would tight-loop here. + // Use the detached drain ctx (and a pure wall-clock fallback) + // instead. + select { + case <-time.After(sleepDur): + case <-callCtx.Done(): + } + } else { + select { + case <-time.After(sleepDur): + case <-ctx.Done(): + // ctx cancel just sets draining on the next iteration; + // don't unwind here. + } + } + } +} + +// generateReplicationTaskForExec injects one execution into the +// replication queue. Delegates to generateWorkflowReplicationTask so the +// rateLimiter wait, frontend-vs-history RPC choice, archetype lookup, +// and history-size-proportional token reservation stay in one place — +// sharded only wraps it to supply a single-element TargetClusters slice +// and to thread the dynamic-config-driven generateViaFrontend flag +// through. +func (a *activities) generateReplicationTaskForExec( + ctx context.Context, + rateLimiter quotas.RateLimiter, + req *shardedBatchReq, + ex *shardedExecutionInfo, +) error { + start := time.Now() + defer func() { + a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(req.Namespace)). + Timer(metrics.GenerateReplicationTaskLatency.Name()).Record(time.Since(start)) + }() + return a.generateWorkflowReplicationTask( + ctx, + rateLimiter, + req.Namespace, + req.NamespaceID, + ex.ExecutionInfo, + []string{req.TargetClusterName}, + a.generateMigrationTaskViaFrontend(), + ) +} + +// attemptVerifyExec runs the source-describe + target-applied check for +// a single execution. Mirrors verifySingleReplicationTask's +// classification — DescribeMutableState on the remote followed by +// workflowVerifier content comparison on success or +// checkSkipWorkflowExecution on NotFound — and returns whether the exec +// is now verified. +// +// Per-classification metric counters are emitted inline (success, +// pending, not-found, busy-workflow, failed). The caller only needs the +// verified bit; not-verified outcomes (missing/busy) drive per-exec +// backoff identically but stay distinct in metrics so a "passive +// cluster apply is in progress" signal stays visible. +// +// We do the DescribeMutableState call inline rather than delegating +// straight to verifySingleReplicationTask so the busy-workflow branch +// keeps its distinct counter. The existing helper folds BUSY_WORKFLOW +// into the generic notVerified result, losing the signal that the +// passive cluster is making progress. +func (a *activities) attemptVerifyExec( + ctx context.Context, + remoteAdminClient adminservice.AdminServiceClient, + ns *namespace.Namespace, + req *shardedBatchReq, + ex *shardedExecutionInfo, +) (bool, error) { + attemptStart := time.Now() + defer func() { + a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(req.Namespace)). + Timer(metrics.VerifyReplicationTaskLatency.Name()).Record(time.Since(attemptStart)) + }() + + archetype, err := a.archetypeIDToName(ctx, ex.ArchetypeID) + if err != nil { + return false, err + } + + vreq := &verifyReplicationTasksRequest{ + Namespace: req.Namespace, + NamespaceID: req.NamespaceID, + TargetClusterEndpoint: req.TargetClusterEndpoint, + TargetClusterName: req.TargetClusterName, + } + + describeStart := time.Now() + mu, err := remoteAdminClient.DescribeMutableState(ctx, &adminservice.DescribeMutableStateRequest{ + Namespace: req.Namespace, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: ex.BusinessID, + RunId: ex.RunID, + }, + Archetype: archetype, + ArchetypeId: ex.ArchetypeID, + SkipForceReload: true, + }) + a.forceReplicationMetricsHandler.Timer(metrics.VerifyDescribeMutableStateLatency.Name()).Record(time.Since(describeStart)) + + nsTag := metrics.NamespaceTag(req.Namespace) + + if err == nil { + result, vErr := a.workflowVerifier(ctx, vreq, remoteAdminClient, a.adminClient, ns, ex.ExecutionInfo, mu) + if vErr != nil { + return false, vErr + } + if result.isVerified() { + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskSuccess.Name()).Record(1) + return true, nil + } + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskPending.Name()).Record(1) + return false, nil + } + + if _, ok := errors.AsType[*serviceerror.NotFound](err); ok { + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskNotFound.Name()).Record(1) + // Retention/zombie path: a not-found execution may already be + // deleted on source (zombie or past retention), in which case it + // never needs to replicate — treat that as verified so the + // shard's completion accounting moves forward. + result, sErr := a.checkSkipWorkflowExecution(ctx, vreq, ex.ExecutionInfo, ns) + if sErr != nil { + return false, sErr + } + return result.isVerified(), nil + } + + if _, ok := errors.AsType[*serviceerror.NamespaceNotFound](err); ok { + return false, temporal.NewNonRetryableApplicationError( + "failed to describe workflow from the remote cluster", "NamespaceNotFound", err) + } + + if resExhausted, ok := errors.AsType[*serviceerror.ResourceExhausted](err); ok && resExhausted.Cause == enumspb.RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW { + // Passive cluster holds the workflow cache lock while applying + // history during SyncWorkflowStateTask. Counted separately from + // pending so the "apply is in progress" signal stays visible, + // but the workflow-side treatment matches pending — per-exec + // backoff applies and the per-shard last-progress timer does + // not move (it only updates on verified outcomes). + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskBusy.Name()).Record(1) + return false, nil + } + + a.forceReplicationMetricsHandler.WithTags(nsTag, metrics.ServiceErrorTypeTag(err)). + Counter(metrics.VerifyReplicationTaskFailed.Name()).Record(1) + return false, fmt.Errorf("describe workflow on remote cluster: %w", err) +} + +// signalReleaseShards sends the mid-flight ReleaseShards signal to the +// parent workflow, freeing the listed shards in the workflow's +// shardInFlight set so the packer can dispatch fresh batches against them +// while this activity stays running on its still-pending shards. +// +// No retry wrapping: a transient failure propagates up so the activity +// fails, the workflow records it via lastErr, and the in-flight batch is +// recovered into the next CAN's RecoveredBuckets — preferable to silently +// swallowing the error here and stranding completed shards. +func (a *activities) signalReleaseShards(ctx context.Context, req *shardedBatchReq, shards []int32) error { + info := activity.GetInfo(ctx) + return a.sdkClientFactory.GetSystemClient().SignalWorkflow(ctx, info.WorkflowExecution.ID, info.WorkflowExecution.RunID, releaseShardsSignalName, releaseShardsPayload{ + BatchID: req.BatchID, + Shards: shards, + }) +} + +// shardVerify holds per-shard verify-phase state for one batch. +type shardVerify struct { + pending int + doneAt time.Time // set when pending first hits zero; cleared on signal release + released bool // ReleaseShards signal already sent + lastProgress time.Time // wall time of the most recent verified outcome +} + +// shardVerifyTracker is keyed by history shard ID. +type shardVerifyTracker map[int32]shardVerify + +func newShardVerifyTracker( + execs []*shardedExecutionInfo, + resume bool, + noProgressByShard map[int32]time.Duration, +) shardVerifyTracker { + t := shardVerifyTracker{} + for _, ex := range execs { + sv := t[ex.Shard] + sv.pending++ + t[ex.Shard] = sv + } + nowSeed := time.Now() + for sh, sv := range t { + if resume { + sv.lastProgress = nowSeed.Add(-noProgressByShard[sh]) + } else { + sv.lastProgress = nowSeed + } + t[sh] = sv + } + return t +} + +func (t shardVerifyTracker) recordVerified(sh int32, now time.Time) { + sv := t[sh] + sv.pending-- + sv.lastProgress = now + if sv.pending == 0 { + sv.doneAt = now + } + t[sh] = sv +} + +func (t shardVerifyTracker) markReleased(shards []int32) { + for _, sh := range shards { + sv := t[sh] + sv.released = true + sv.doneAt = time.Time{} + t[sh] = sv + } +} + +// totalIdleCost sums idle time across shards that are completed but not +// yet released — the "shard-seconds" unit the IdleShardCost threshold is +// denominated in. +func (t shardVerifyTracker) totalIdleCost(now time.Time) time.Duration { + var total time.Duration + for _, sv := range t { + if !sv.doneAt.IsZero() { + total += now.Sub(sv.doneAt) + } + } + return total +} + +// awaitingRelease returns completed-but-not-yet-signaled shard IDs in +// ascending order so the signal payload is deterministic across replays. +func (t shardVerifyTracker) awaitingRelease() []int32 { + var out []int32 + for sh, sv := range t { + if !sv.doneAt.IsZero() { + out = append(out, sh) + } + } + slices.Sort(out) + return out +} + +// allCompleted returns every shard that finished during this activity +// run — both signal-released and still awaiting release at return. +func (t shardVerifyTracker) allCompleted() []int32 { + var out []int32 + for sh, sv := range t { + if sv.released || !sv.doneAt.IsZero() { + out = append(out, sh) + } + } + slices.Sort(out) + return out +} + +// pickStuck returns (shard, age, true) for the lowest-numbered shard +// whose cumulative no-progress duration meets or exceeds threshold. +func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (int32, time.Duration, bool) { + var ( + minShard int32 + minAge time.Duration + found bool + ) + for sh, sv := range t { + if sv.pending <= 0 { + continue + } + age := now.Sub(sv.lastProgress) + if age < threshold { + continue + } + if !found || sh < minShard { + minShard = sh + minAge = age + found = true + } + } + return minShard, minAge, found +} + +// buildInFlight groups unverified execs by shard then businessID and +// attaches the cumulative no-progress duration per shard, for the +// drain-mode activity return. Shards with zero unverified execs are not +// included — fully verified shards are reported via CompletedShards (and +// any that completed mid-flight have already been signal-released so the +// workflow's packer could reuse them). +func (a *activities) buildInFlight( + execs []*shardedExecutionInfo, + verified []bool, + shards shardVerifyTracker, + now time.Time, +) []ResumeShard { + byShard := map[int32]map[string][]RunEntry{} + for i, ex := range execs { + if verified[i] { + continue + } + if byShard[ex.Shard] == nil { + byShard[ex.Shard] = map[string][]RunEntry{} + } + byShard[ex.Shard][ex.BusinessID] = append(byShard[ex.Shard][ex.BusinessID], RunEntry{ + RunID: ex.RunID, + ArchetypeID: ex.ArchetypeID, + }) + } + if len(byShard) == 0 { + return nil + } + shardIDs := make([]int32, 0, len(byShard)) + for sh := range byShard { + shardIDs = append(shardIDs, sh) + } + slices.Sort(shardIDs) + out := make([]ResumeShard, 0, len(shardIDs)) + for _, sh := range shardIDs { + out = append(out, ResumeShard{ + Shard: sh, + Execs: byShard[sh], + NoProgressDuration: now.Sub(shards[sh].lastProgress), + }) + } + return out +} + +// firstUnverifiedOnShard returns the index of the first execution in the +// flattened execs slice that targets the given shard and hasn't verified +// yet. Used to name a concrete (BusinessID, RunID) in the ShardNoProgress +// failure event. +// +// Panics if no such execution exists — the only caller is the stuck-shard +// backstop, which by construction only fires for shards with at least one +// pending exec. A silent fallback to index 0 would hide a future +// invariant violation behind a misleading error message. +func (a *activities) firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shard int32) int { + for i, ex := range execs { + if verified[i] { + continue + } + if ex.Shard == shard { + return i + } + } + panic(fmt.Sprintf("firstUnverifiedOnShard: no unverified exec on shard %d", shard)) +} + +// backoffDelay returns the per-exec retry delay after `attempt` +// consecutive failed verify attempts: 100ms × 2^(attempt-1), capped at +// 5s. The cap bounds how long after the apply pipeline recovers we'd +// take to notice; the per-shard no-progress timer fires after enough of +// these capped retries to distinguish "actively checking" from "lazy +// polling gave up". +func backoffDelay(attempt int) time.Duration { + if attempt < 1 { + attempt = 1 + } + if attempt > 6 { + return 5 * time.Second + } + return 100 * time.Millisecond * (1 << (attempt - 1)) +} diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go new file mode 100644 index 00000000000..6bfd3fd2bfb --- /dev/null +++ b/service/worker/migration/sharded_types.go @@ -0,0 +1,374 @@ +package migration + +import ( + "encoding/json" + "fmt" + "slices" + "time" +) + +const ( + // shardedForceReplicationWorkflowName is the registered workflow name. + // Distinct from the legacy ForceReplicationWorkflow so both variants + // coexist in the same worker — pick which to use at workflow start + // time by setting the start request's workflow type. + shardedForceReplicationWorkflowName = "force-replication-sharded" + + // shardedBatchActivityName is the registered activity name; the + // workflow dispatches by name so the same activity body can be + // referenced across CAN cycles without re-registering. + shardedBatchActivityName = "ReplicateBatch" + + // releaseShardsSignalName carries mid-flight ReleaseShards signals + // from active replicate-batch activities back to their parent + // workflow. Drain-mode shard completions ride the activity return + // value instead, so this signal only fires while the activity is + // still running normally. + releaseShardsSignalName = "ReleaseShards" + + // defaultShardedListPageSize is the ListWorkflows page size when + // the sharded workflow's params.ListWorkflowsPageSize is unset. + // Named to avoid clashing with the legacy + // defaultListWorkflowsPageSize already declared in + // force_replication_workflow.go. + defaultShardedListPageSize = 1000 + + // defaultBatchSize bounds the total executions in any single + // ReplicateBatch activity. Activity-payload sizing knob, not a + // shard-fill threshold. + defaultBatchSize = 100 + + // defaultMaxExecsPerShard bounds the executions any single shard + // can contribute to a batch — i.e. the per-shard inject blast + // radius before that shard's apply queue has to absorb a burst. + // 50 keeps a hot shard's contribution under half a default-sized + // batch (BatchSize=100), so a batch still spans ≥2 shards. + defaultMaxExecsPerShard = 50 + + // defaultShardNoProgress is the per-shard cumulative no-progress + // backstop. While a shard's pending exec count is non-zero and + // no exec on that shard has produced a verified outcome for this + // many seconds (carried across CAN via the resume payload), the + // activity fails non-retryably naming the stuck shard. + defaultShardNoProgress = 5 * time.Minute + + // defaultDrainGrace is the wall-budget the activity gets after + // the workflow cancels it for CAN. Continues verifying until + // either the grace expires, the idle-cost trigger fires, or + // every exec verifies. + defaultDrainGrace = 15 * time.Second + + // defaultIdleShardCost is the cumulative idle-time threshold + // (the "shard-seconds" unit: 30 s with 1 idle shard equals + // 3.3 s with 9 idle) at which the activity signal-releases its + // completed-but-not-yet-released shards mid-flight. + defaultIdleShardCost = 30 * time.Second + + // defaultPerBatchGenerateRPS is the per-batch inject-phase target. + // Sharded dispatches many concurrent batches and each builds its own + // limiter, so this caps the per-batch generate-replication-task rate; + // the workflow does not normalise against a global cap the way the + // existing migration's OverallRps does. Sits near the mid-range of + // what existing force-replication deployments configure per-activity + // (OverallRps 10–100 divided across ConcurrentActivityCount 2–10). + defaultPerBatchGenerateRPS = 30.0 + + // defaultConcurrentBatchCap is the ceiling applied to the derived + // default of TargetClusterShardCount/4. Keeps the in-flight batch + // count safely inside per-worker concurrent-activity suggestions + // even on the largest cells (4k+ shards), and absolutely bounds the + // cluster blast radius of a single force-rep run. + defaultConcurrentBatchCap = 500 + +) + +// RunEntry is the per-run leaf in the nested batch payload. Carries the +// RunID plus an optional ArchetypeID, serialised as a JSON tuple: +// `["runID"]` when ArchetypeID is zero, `["runID", N]` when set. +// +// Why a tuple, not an object: a tuple omits the JSON field names +// (`"r":`, `"a":`) that would otherwise repeat on every run, which is +// the main lever behind the nested payload's byte savings. With many +// runs per BusinessID (the heavy-reuse case), this collapses the +// per-run encoding overhead from ~47 bytes (flat ExecutionInfo) to +// ~10 bytes per run. +type RunEntry struct { + RunID string + ArchetypeID uint32 +} + +// MarshalJSON / UnmarshalJSON: custom because Go's default JSON has no +// way to express a heterogeneous tuple, and we want the archetype-omission +// to happen by changing the tuple length rather than emitting an explicit +// zero. +func (r RunEntry) MarshalJSON() ([]byte, error) { + if r.ArchetypeID == 0 { + return fmt.Appendf(nil, `[%q]`, r.RunID), nil + } + return fmt.Appendf(nil, `[%q,%d]`, r.RunID, r.ArchetypeID), nil +} + +func (r *RunEntry) UnmarshalJSON(data []byte) error { + var raw []json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return fmt.Errorf("RunEntry: %w", err) + } + if len(raw) < 1 || len(raw) > 2 { + return fmt.Errorf("RunEntry: expected [runID] or [runID,archetypeID], got %d-element array", len(raw)) + } + if err := json.Unmarshal(raw[0], &r.RunID); err != nil { + return fmt.Errorf("RunEntry runID: %w", err) + } + r.ArchetypeID = 0 + if len(raw) == 2 { + if err := json.Unmarshal(raw[1], &r.ArchetypeID); err != nil { + return fmt.Errorf("RunEntry archetypeID: %w", err) + } + } + return nil +} + +// BatchPayload groups runs by (shard, businessID) so a single businessID +// with many runs costs one BID-string-worth of bytes instead of one per +// run. The wire shape behind shardedBatchReq.Executions, ResumeShard.Execs +// (the per-shard inner map), and ShardedForceReplicationParams.RecoveredBuckets. +// +// On-wire form: +// +// {"shardID": {"businessID": [["runID"], ["runID", archetypeID], ...], ...}, ...} +// +// Top-level keys are shard IDs; inner-map keys are businessIDs; values +// are RunEntry tuples. +type BatchPayload map[int32]map[string][]RunEntry + +// totalRuns counts runs across all (shard, BID) groups. +func (p BatchPayload) totalRuns() int { + n := 0 + for _, byBID := range p { + for _, runs := range byBID { + n += len(runs) + } + } + return n +} + +// sortedShards returns shard IDs in ascending order. Used to give the +// activity-side flatten a deterministic iteration order for replays. +func (p BatchPayload) sortedShards() []int32 { + out := make([]int32, 0, len(p)) + for sh := range p { + out = append(out, sh) + } + slices.Sort(out) + return out +} + +// shardedExecutionInfo pairs an upstream ExecutionInfo with the destination +// history shard the sharded design routes by. Shard is kept out of the +// upstream ExecutionInfo struct because no other workflow needs it; the +// wire format (BatchPayload) carries shard as the outer map key. +type shardedExecutionInfo struct { + *ExecutionInfo + Shard int32 +} + +// flatten produces a deterministically-ordered slice of execs paired with +// their destination shard: shards ascending, BIDs alphabetical within +// shard, runs in input order. Lets the inner verify loop stay index-based +// even though the wire shape is nested. +func (p BatchPayload) flatten() []*shardedExecutionInfo { + n := p.totalRuns() + if n == 0 { + return nil + } + out := make([]*shardedExecutionInfo, 0, n) + for _, sh := range p.sortedShards() { + byBID := p[sh] + bids := make([]string, 0, len(byBID)) + for bid := range byBID { + bids = append(bids, bid) + } + slices.Sort(bids) + for _, bid := range bids { + for _, r := range byBID[bid] { + out = append(out, &shardedExecutionInfo{ + ExecutionInfo: &ExecutionInfo{ + BusinessID: bid, + RunID: r.RunID, + ArchetypeID: r.ArchetypeID, + }, + Shard: sh, + }) + } + } + } + return out +} + +// merge folds src into p. Used by collectRecoveredBuckets to combine +// per-batch payloads back into a single carry-over map. +func (p BatchPayload) merge(src BatchPayload) { + for sh, byBID := range src { + if p[sh] == nil { + p[sh] = map[string][]RunEntry{} + } + for bid, runs := range byBID { + p[sh][bid] = append(p[sh][bid], runs...) + } + } +} + +// ShardedForceReplicationParams is the workflow input. Configuration +// fields are read-only across CAN cycles; the carry-over block at the +// bottom is mutated each cycle. +type ShardedForceReplicationParams struct { + // ---- Configuration ---- + Namespace string + BatchSize int + MaxExecsPerShard int + ListWorkflowsPageSize int + TargetClusterEndpoint string + TargetClusterName string + TargetClusterShardCount int32 + + ShardNoProgress time.Duration + DrainGrace time.Duration + IdleShardCost time.Duration + + // PerBatchGenerateRPS is the inject-phase rate-limiter target inside + // each ReplicateBatch activity. Each batch builds its own + // quotas.RateLimiter at this rate, so N concurrent batches inject at + // N× this rate in aggregate — unlike the existing migration package's + // OverallRps, sharded does not normalise against a workflow-global + // budget. Defaults to defaultPerBatchGenerateRPS. + PerBatchGenerateRPS float64 + + // ConcurrentBatchCount is the absolute ceiling on in-flight + // ReplicateBatch activities. Per-shard exclusivity already bounds + // concurrency to TargetClusterShardCount, but at production cell + // sizes (1k–4k shards) that's well past the worker's + // concurrent-activity budget. This cap keeps the workflow inside + // that budget and limits the cluster blast radius of a single + // force-rep run. Defaults to + // min(TargetClusterShardCount/4, defaultConcurrentBatchCap). + ConcurrentBatchCount int + + // EstimationMultiplier sizes the QPSQueue's initial slice capacity + // (multiplier × ConcurrentBatchCount + 1). Pure allocation hint — + // the sliding window's logical max stays ConcurrentBatchCount + 1. + // Defaults to 2. + EstimationMultiplier int + + // ---- Continue-as-new carry-over ---- + NextPageToken []byte + ContinuedAsNewCount int + TotalForceReplicateWorkflowCount int64 + ReplicatedWorkflowCount int64 + ReplicatedWorkflowCountPerSecond float64 + + // QPSQueue carries the sliding-window samples across CAN so the + // per-second rate doesn't drop to zero on every cycle boundary. + QPSQueue QPSQueue + + // ResumeShards carries unverified execs from drained activities in + // the prior CAN cycle. The new run dispatches resume activities for + // these before the page loop runs, so their shards are claimed in + // shardInFlight from the start and the packer treats them as busy. + ResumeShards []ResumeShard + + // RecoveredBuckets carries execs whose dispatching activity returned + // a cancellation without returning a result — i.e., the activity + // body never ran (cancel-before-start race). They were dispatched + // but never injected, so the new cycle restores them into the + // streaming buckets to be dispatched as fresh inject+verify batches. + RecoveredBuckets BatchPayload +} + +// ResumeShard carries one shard's worth of unverified execs from a drained +// activity across a CAN boundary to the resume activity that picks them +// up. NoProgressDuration is the cumulative time the shard went without a +// verified outcome at drain time; the resume activity initialises its own +// per-shard last-progress clock to (now - NoProgressDuration) so the +// backstop check sees the full elapsed no-progress window, not just the +// current activity's slice. +// +// Execs is keyed by businessID: each entry is a list of RunEntry tuples +// for that BID. Grouping by BID at the wire level lets a hot BID (with +// many runs) collapse to one BID-string + N tuples rather than N copies +// of the BID; see BatchPayload's docstring. +type ResumeShard struct { + Shard int32 + Execs map[string][]RunEntry + NoProgressDuration time.Duration +} + +// shardedBatchReq is the per-batch activity input. Executions is the +// per-shard, per-BID nested payload — the workflow has marked every +// shard appearing as a top-level key in shardInFlight before dispatch, +// and the activity is responsible for either signal-releasing each shard +// mid-flight or listing it in the return value's CompletedShards / InFlight +// set. +// +// Resume=true skips the inject phase: the execs were already injected by +// some earlier activity that was cancelled at drain time and returned its +// unverified execs in its result. NoProgressByShard carries the cumulative +// pre-resume no-progress duration so the per-shard backstop stays +// meaningful across resume cycles. +// +// PerBatchGenerateRPS is the inject-phase rate-limiter target. Mirrors +// generateReplicationTasksRequest.RPS in the existing migration package — +// each batch builds its own quotas.RateLimiter at that rate, so two +// concurrent batches inject at 2× this rate in aggregate. +type shardedBatchReq struct { + BatchID int64 + Namespace string + NamespaceID string + Executions BatchPayload + + TargetClusterEndpoint string + TargetClusterName string + + Resume bool + NoProgressByShard map[int32]time.Duration + + PerBatchGenerateRPS float64 + + ShardNoProgress time.Duration + DrainGrace time.Duration + IdleShardCost time.Duration +} + +// replicateBatchResult is the activity's return payload. The activity is +// the source of truth for which execs verified vs. are still outstanding +// when it returns — only it has the per-exec verify state — so the drain +// payload rides the return value rather than a signal. The workflow's +// dispatch coroutine reads InFlight into drainPayload on nil-error return. +// +// CompletedShards is informational (the dispatch coroutine's defer clears +// heldByBatch + shardInFlight regardless), but keeping it in the result +// gives metrics and future bookkeeping a clean handle on "which shards +// this batch finished". +type replicateBatchResult struct { + CompletedShards []int32 + InFlight []ResumeShard + + // VerifiedCount is the number of executions this activity invocation + // finished verifying (including retention/zombie skips that resolve + // as verified). The workflow accumulates this into its running + // ReplicatedWorkflowCount and emits the per-batch delta as the + // replicated_workflow_count counter. + VerifiedCount int64 +} + +// releaseShardsPayload is the body of the mid-flight ReleaseShards signal +// an activity sends to its parent workflow when the cumulative idle cost +// across its completed-but-not-yet-released shards crosses IdleShardCost. +// The workflow handler clears these shards from shardInFlight + +// heldByBatch[BatchID] so the packer can immediately dispatch new work +// against them while the activity stays running on its still-pending +// shards. Only fires in normal mode — once the activity enters drain mode +// it returns its remaining state via the activity result instead. +type releaseShardsPayload struct { + BatchID int64 + Shards []int32 +} diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go new file mode 100644 index 00000000000..05c85b8c0b4 --- /dev/null +++ b/service/worker/migration/sharded_types_test.go @@ -0,0 +1,91 @@ +package migration + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestRunEntry_JSONRoundTrip pins down the tuple-array wire shape: +// 1-element when ArchetypeID is zero, 2-element otherwise. +func TestRunEntry_JSONRoundTrip(t *testing.T) { + tests := []struct { + name string + entry RunEntry + expect string + }{ + {"zero archetype is omitted", RunEntry{RunID: "r1"}, `["r1"]`}, + {"non-zero archetype is included", RunEntry{RunID: "r1", ArchetypeID: 42}, `["r1",42]`}, + {"escaped runID", RunEntry{RunID: `r"1`}, `["r\"1"]`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + b, err := json.Marshal(tt.entry) + require.NoError(t, err) + require.Equal(t, tt.expect, string(b)) + + var out RunEntry + require.NoError(t, json.Unmarshal(b, &out)) + require.Equal(t, tt.entry, out) + }) + } +} + +// TestRunEntry_UnmarshalRejectsBadShape: the marshaller emits 1- or +// 2-element tuples only; anything else is a protocol violation and +// must surface a clear error rather than a silent zero value. +func TestRunEntry_UnmarshalRejectsBadShape(t *testing.T) { + cases := []string{ + `[]`, + `["r1", 1, 2]`, + `{"r": "r1"}`, + `"r1"`, + } + for _, in := range cases { + t.Run(in, func(t *testing.T) { + var out RunEntry + require.Error(t, json.Unmarshal([]byte(in), &out)) + }) + } +} + +// TestBatchPayload_JSONRoundTrip exercises the nested wire shape end +// to end so a change to either RunEntry or the surrounding map type +// can't silently regress it. +func TestBatchPayload_JSONRoundTrip(t *testing.T) { + p := BatchPayload{ + 7: {"bid-a": {{RunID: "r1"}, {RunID: "r2", ArchetypeID: 5}}}, + 8: {"bid-b": {{RunID: "r3"}}}, + } + b, err := json.Marshal(p) + require.NoError(t, err) + // Go's encoding/json sorts map keys, so this output is stable. + require.JSONEq(t, `{ + "7": {"bid-a": [["r1"], ["r2", 5]]}, + "8": {"bid-b": [["r3"]]} + }`, string(b)) + + var out BatchPayload + require.NoError(t, json.Unmarshal(b, &out)) + require.Equal(t, p, out) +} + +// TestBatchPayload_Flatten orders by shard ascending then BID +// alphabetical; runs within a BID keep input order. The activity +// inner loop depends on this for deterministic replays. +func TestBatchPayload_Flatten(t *testing.T) { + p := BatchPayload{ + 2: {"b-z": {{RunID: "rz"}}, "b-a": {{RunID: "ra1"}, {RunID: "ra2"}}}, + 1: {"b-c": {{RunID: "rc"}}}, + } + got := p.flatten() + require.Len(t, got, 4) + require.Equal(t, int32(1), got[0].Shard) + require.Equal(t, "b-c", got[0].BusinessID) + require.Equal(t, int32(2), got[1].Shard) + require.Equal(t, "b-a", got[1].BusinessID) + require.Equal(t, "ra1", got[1].RunID) + require.Equal(t, "ra2", got[2].RunID) + require.Equal(t, "b-z", got[3].BusinessID) +} diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go new file mode 100644 index 00000000000..5a2ef24cf89 --- /dev/null +++ b/service/worker/migration/sharded_workflow.go @@ -0,0 +1,915 @@ +package migration + +import ( + "fmt" + "slices" + "time" + + "go.temporal.io/api/workflowservice/v1" + sdkclient "go.temporal.io/sdk/client" + "go.temporal.io/sdk/temporal" + "go.temporal.io/sdk/workflow" + "go.temporal.io/server/common" + "go.temporal.io/server/common/metrics" +) + +// ShardedForceReplicationWorkflow runs the sharded design for one CAN +// cycle: dispatch any resume activities carried over from the prior +// cycle, then page through ListWorkflows until either the namespace +// exhausts or workflow.GetContinueAsNewSuggested(ctx) trips, bucketing +// each execution by destination history shard and dispatching a +// paired inject+verify activity once a bucket reaches packing +// eligibility. At cycle end, every remaining bucket flushes as +// packed activities. +// +// If there are more pages, in-flight activities are cancelled (giving +// them DrainGrace to drain), their drain payload arrives via the +// activity return value, and the workflow CANs with NextPageToken + +// ResumeShards in the carry-over. Otherwise it waits for activities +// to finish naturally and returns nil. +func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { + // Page token at workflow entry — returned by the status query as + // PageTokenForRestart so an operator can resume from this run's + // starting position rather than its current (in-flight) position. + startPageToken := params.NextPageToken + + // Register the status query under the same name upstream uses + // (forceReplicationStatusQueryType = "force-replication-status") + // so tooling that polls force-rep progress works across both + // workflow variants. Sharded-irrelevant ForceReplicationStatus + // fields (TaskQueueUserDataReplicationStatus) are left zero — + // sharded only handles the data-replication phase. + if err := workflow.SetQueryHandler(ctx, forceReplicationStatusQueryType, func() (ForceReplicationStatus, error) { + return ForceReplicationStatus{ + ContinuedAsNewCount: params.ContinuedAsNewCount, + TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, + ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, + ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, + PageTokenForRestart: startPageToken, + }, nil + }); err != nil { + return err + } + + state, err := newShardedWorkflowState(ctx, ¶ms) + if err != nil { + return err + } + // Defaults are now applied; reject configurations the packer + // can't honour. MaxExecsPerShard > BatchSize is meaningless — + // each batch caps at BatchSize total, so the per-shard cap + // can't exceed the whole-batch cap. + if params.MaxExecsPerShard > params.BatchSize { + return temporal.NewNonRetryableApplicationError( + fmt.Sprintf("MaxExecsPerShard (%d) must be <= BatchSize (%d)", params.MaxExecsPerShard, params.BatchSize), + "InvalidConfiguration", nil) + } + + // On the first cycle, populate TotalForceReplicateWorkflowCount + // via the same CountWorkflow activity upstream uses. Skipped on + // subsequent CAN cycles — the count carries across via params. + if params.TotalForceReplicateWorkflowCount == 0 { + wfCount, err := shardedCountWorkflowsForReplication(ctx, ¶ms) + if err != nil { + return err + } + params.TotalForceReplicateWorkflowCount = wfCount + } + + return state.run(ctx) +} + +// shardedCountWorkflowsForReplication asks the frontend how many +// workflows match the namespace's force-rep query. Used once at +// workflow start to seed TotalForceReplicateWorkflowCount for the +// status query's progress reporting. No filter is applied — sharded +// currently lists the namespace's entire workflow population. +func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedForceReplicationParams) (int64, error) { + ao := workflow.ActivityOptions{ + StartToCloseTimeout: 2 * time.Minute, + RetryPolicy: forceReplicationActivityRetryPolicy, + } + var a *activities + var output countWorkflowResponse + if err := workflow.ExecuteActivity( + workflow.WithActivityOptions(ctx, ao), + a.CountWorkflow, + &workflowservice.CountWorkflowExecutionsRequest{ + Namespace: params.Namespace, + }).Get(ctx, &output); err != nil { + return 0, err + } + return output.WorkflowCount, nil +} + +// shardedWorkflowState holds the workflow's per-run state. Workflow +// coroutines yield only at SDK calls, so plain maps + ints are safe +// without mutexes — workflow.Await re-evaluates its predicate after +// each yield, which is what makes the shard-in-flight bookkeeping +// drive each dispatch coroutine's wait. +type shardedWorkflowState struct { + params *ShardedForceReplicationParams + + namespaceID string + + // buckets accumulate execs that have been listed but not yet + // dispatched. Nested by destination shard then businessID, so a + // hot BID's many runs share one BID-string-worth of bytes when + // the bucket is shipped over the wire. + buckets BatchPayload + + // bucketCounts mirrors len of all runs across BIDs for each + // shard. Kept as a sidecar so the packer's per-shard ordering + // decisions are O(1) rather than O(#BIDs in shard); it's + // consulted many times per cycle. + bucketCounts map[int32]int + + // shardInFlight is the per-shard exclusivity set: a shard's + // entry is set when it's part of any in-flight batch and + // cleared when that batch returns (either fully or via mid-flight + // signal-release). Concurrent batches are limited only by this + // set — there is no global slot cap. + shardInFlight map[int32]bool + + // heldByBatch tracks per-batch shard ownership. spawnBatch + // populates it with the batch's claimed shards; the signal + // handler removes entries as shards are released mid-flight; + // the dispatch coroutine's defer clears whatever's left after + // the activity returns. Required because a signal-released + // shard may have been re-claimed by a subsequent batch — the + // returning original batch must only clear its own remaining + // claims, not stomp on the new claimant. + heldByBatch map[int64]map[int32]bool + + // batchCancels carries each batch's cancel func keyed by batch + // ID, so the workflow's drain-for-CAN phase can cancel + // in-flight activities individually. + batchCancels map[int64]workflow.CancelFunc + + // batchExecs tracks the input payload of each in-flight batch. + // Cleared on any nil-error return (drained execs are folded + // into drainPayload from the activity result; cleanly completed + // batches return an empty InFlight). Anything left at CAN time + // corresponds to a batch whose activity returned CanceledError + // with no result — i.e. the activity body never ran. Those + // execs are recovered into the next cycle's streaming buckets. + batchExecs map[int64]BatchPayload + + // pendingDispatches counts spawned dispatch coroutines that + // have not yet returned. The main coroutine waits on this + // dropping to zero before issuing CAN or returning. + pendingDispatches int + + // drainPayload accumulates ResumeShard entries from drained + // activities (via the activity result on nil-error return). + // Fed into the CAN-carry-over params at the end. + drainPayload []ResumeShard + + // lastErr stops further dispatch once an activity errors out + // (e.g. ShardNoProgress). Without it the workflow would keep + // paging through ListWorkflows after a broken namespace surfaces, + // burning the rest of the population before returning the + // failure. + lastErr error + + nextBatchID int64 + + // metricsHandler is tagged with the workflow's fixed scope + + // namespace once at state construction; recordVerified reuses it on + // every batch return. + metricsHandler sdkclient.MetricsHandler +} + +func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicationParams) (*shardedWorkflowState, error) { + ao := workflow.ActivityOptions{ + StartToCloseTimeout: 24 * time.Hour, + HeartbeatTimeout: time.Minute, + RetryPolicy: forceReplicationActivityRetryPolicy, + } + metaCtx := workflow.WithActivityOptions(ctx, ao) + var a *activities + var md MetadataResponse + if err := workflow.ExecuteActivity(metaCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { + return nil, err + } + if params.TargetClusterShardCount <= 0 { + params.TargetClusterShardCount = md.ShardCount + } + if params.BatchSize <= 0 { + params.BatchSize = defaultBatchSize + } + if params.MaxExecsPerShard <= 0 { + params.MaxExecsPerShard = defaultMaxExecsPerShard + } + if params.ShardNoProgress <= 0 { + params.ShardNoProgress = defaultShardNoProgress + } + if params.DrainGrace <= 0 { + params.DrainGrace = defaultDrainGrace + } + if params.IdleShardCost <= 0 { + params.IdleShardCost = defaultIdleShardCost + } + if params.ListWorkflowsPageSize <= 0 { + params.ListWorkflowsPageSize = defaultShardedListPageSize + } + if params.PerBatchGenerateRPS <= 0 { + params.PerBatchGenerateRPS = defaultPerBatchGenerateRPS + } + if params.ConcurrentBatchCount <= 0 { + params.ConcurrentBatchCount = defaultConcurrentBatchCount(params.TargetClusterShardCount) + } + if params.EstimationMultiplier <= 0 { + params.EstimationMultiplier = 2 + } + // QPSQueue is sized off ConcurrentBatchCount (one sample slot per + // expected in-flight batch + one for the starting count). Seeded + // with the current ReplicatedWorkflowCount so the very first + // post-CAN batch return has a baseline to compute the rate against. + if params.QPSQueue.Data == nil { + params.QPSQueue = NewQPSQueue(params.ConcurrentBatchCount, params.EstimationMultiplier) + params.QPSQueue.Enqueue(ctx, params.ReplicatedWorkflowCount) + } + s := &shardedWorkflowState{ + params: params, + namespaceID: md.NamespaceID, + buckets: BatchPayload{}, + bucketCounts: map[int32]int{}, + shardInFlight: map[int32]bool{}, + heldByBatch: map[int64]map[int32]bool{}, + batchCancels: map[int64]workflow.CancelFunc{}, + batchExecs: map[int64]BatchPayload{}, + metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ + metrics.OperationTagName: metrics.MigrationWorkflowScope, + NamespaceTagName: params.Namespace, + }), + } + // Restore execs recovered from cancel-before-start batches in + // the prior cycle so the streaming packer picks them up + // alongside any new pages. + s.buckets.merge(params.RecoveredBuckets) + for sh, byBID := range params.RecoveredBuckets { + for _, runs := range byBID { + s.bucketCounts[sh] += len(runs) + } + } + params.RecoveredBuckets = nil + return s, nil +} + +func (s *shardedWorkflowState) run(ctx workflow.Context) error { + // Start the signal handler coroutine first so any signal + // arriving during resume dispatch or page-loop drains is + // processed promptly. + workflow.Go(ctx, s.handleReleaseSignals) + + // Dispatch resume activities carried over from the prior cycle. + // Done before the page loop so their shards are claimed in + // shardInFlight before any new pages arrive — keeps the packer + // from racing to dispatch against them with fresh execs. + s.dispatchResumeBatches(ctx) + + ao := workflow.ActivityOptions{ + StartToCloseTimeout: time.Hour, + RetryPolicy: &temporal.RetryPolicy{ + InitialInterval: time.Second, + BackoffCoefficient: 2.0, + MaximumAttempts: 3, + }, + } + listCtx := workflow.WithActivityOptions(ctx, ao) + + // Drive ListWorkflows until either we exhaust the namespace or + // the SDK signals that history is large enough to CAN. + for !workflow.GetInfo(ctx).GetContinueAsNewSuggested() { + if s.lastErr != nil { + break + } + listReq := &workflowservice.ListWorkflowExecutionsRequest{ + Namespace: s.params.Namespace, + PageSize: int32(s.params.ListWorkflowsPageSize), + NextPageToken: s.params.NextPageToken, + } + var a *activities + var listResp listWorkflowsResponse + if err := workflow.ExecuteActivity(listCtx, a.ListWorkflows, listReq).Get(ctx, &listResp); err != nil { + return err + } + for _, ex := range listResp.Executions { + sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.params.TargetClusterShardCount) + s.addToBucket(sh, ex.BusinessID, RunEntry{ + RunID: ex.RunID, + ArchetypeID: ex.ArchetypeID, + }) + } + s.params.NextPageToken = listResp.NextPageToken + + for s.tryPackStreaming(ctx, false) { //nolint:revive // intentional empty body + } + + if len(listResp.NextPageToken) == 0 { + break + } + } + + // Drain remaining buckets — dispatch every leftover exec. The + // dispatched batches may themselves get cancelled by drainForCAN + // below if we're going to CAN; that's fine, each cancelled + // batch returns its drain payload in its result. + if s.lastErr == nil { + s.drainBuckets(ctx) + } + + // If there are no more pages we're done: wait for activities to + // finish naturally and return. Slow shards hold their claims + // until their per-shard no-progress backstop trips. + if len(s.params.NextPageToken) == 0 || s.lastErr != nil { + _ = workflow.Await(ctx, func() bool { + return s.pendingDispatches == 0 || s.lastErr != nil + }) + if s.lastErr != nil { + return s.lastErr + } + return nil + } + + // More pages → CAN. Cancel in-flight batches, wait for them to + // drain, then harvest any leftover batchExecs entries (batches + // whose activity was cancelled before its body ran — no result + // returned) into RecoveredBuckets so they get re-dispatched as + // fresh inject+verify activities next cycle. + s.drainForCAN(ctx) + + if s.lastErr != nil { + return s.lastErr + } + + next := *s.params + next.ContinuedAsNewCount++ + // Defensive copy so we don't alias s.drainPayload into the + // carry-over params. drainForCAN guarantees every spawnBatch + // coroutine has finished appending before we get here (the + // append happens before the defer that decrements + // pendingDispatches), so this is style — but cheap insurance + // against future code paths that append post-drain. + next.ResumeShards = append([]ResumeShard(nil), s.drainPayload...) + next.RecoveredBuckets = collectRecoveredBuckets(s.batchExecs) + return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) +} + +// recordVerified accumulates one batch's verified-exec delta into the +// workflow's running count, emits the per-batch counter delta, and +// updates the sliding-window RPS gauge. No-op when verified == 0 so a +// batch that ran entirely as drain-no-progress doesn't poison the +// QPSQueue with a zero-delta sample. +func (s *shardedWorkflowState) recordVerified(ctx workflow.Context, verified int64) { + if verified <= 0 { + return + } + s.params.ReplicatedWorkflowCount += verified + + s.metricsHandler.Counter(metrics.ReplicatedWorkflowCount.Name()).Inc(verified) + + s.params.QPSQueue.Enqueue(ctx, s.params.ReplicatedWorkflowCount) + s.params.ReplicatedWorkflowCountPerSecond = s.params.QPSQueue.CalculateQPS() + s.metricsHandler.Gauge(ForceReplicationRpsTagName).Update(s.params.ReplicatedWorkflowCountPerSecond) +} + +// defaultConcurrentBatchCount derives the in-flight-batch ceiling +// from the target cluster's shard count: a quarter of the shards, +// capped at defaultConcurrentBatchCap. The 1/4 fraction keeps the +// workflow inside Temporal Cloud's concurrent-activity suggestions +// (a 4k-shard cell with cap 500 still has spare worker slots for +// unrelated activities); the absolute cap bounds the cluster blast +// radius regardless of cluster size. Returns at least 1. +func defaultConcurrentBatchCount(shards int32) int { + return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) +} + +// collectRecoveredBuckets re-buckets any execs from batches whose +// dispatching activity returned CanceledError without returning a +// result — i.e. the activity body never ran, so its execs were +// never injected. They go back into the next cycle's streaming +// buckets to be dispatched as fresh inject+verify batches. The +// shard is the top-level map key on each batch's payload, so no +// re-hashing here — collectRecoveredBuckets just merges. +func collectRecoveredBuckets(batchExecs map[int64]BatchPayload) BatchPayload { + if len(batchExecs) == 0 { + return nil + } + out := BatchPayload{} + for _, bp := range batchExecs { + out.merge(bp) + } + return out +} + +// addToBucket appends one run to the (shard, BID) bucket and bumps +// the per-shard count. The count is a sidecar so the streaming +// packer's eligibility check stays O(1) per shard. +func (s *shardedWorkflowState) addToBucket(shard int32, businessID string, run RunEntry) { + if s.buckets[shard] == nil { + s.buckets[shard] = map[string][]RunEntry{} + } + s.buckets[shard][businessID] = append(s.buckets[shard][businessID], run) + s.bucketCounts[shard]++ +} + +// takeFromBucket consumes up to n runs from the given shard and +// returns them grouped by BID. Walks BIDs in alphabetical order so +// the resulting payload is deterministic across replays; takes whole +// per-BID runs only as needed to reach n. Empties the shard from +// s.buckets / s.bucketCounts when nothing remains. +func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]RunEntry { + if n <= 0 { + return nil + } + byBID := s.buckets[shard] + if len(byBID) == 0 { + return nil + } + bids := make([]string, 0, len(byBID)) + for bid := range byBID { + bids = append(bids, bid) + } + slices.Sort(bids) + + out := map[string][]RunEntry{} + taken := 0 + for _, bid := range bids { + if taken >= n { + break + } + runs := byBID[bid] + take := min(len(runs), n-taken) + // append([]RunEntry(nil), ...) gives the output its own + // backing array — keeps the workflow's leftover slice + // (byBID[bid][take:]) and the activity's input independent + // in case either side appends later. + out[bid] = append([]RunEntry(nil), runs[:take]...) + if take == len(runs) { + delete(byBID, bid) + } else { + byBID[bid] = runs[take:] + } + taken += take + } + s.bucketCounts[shard] -= taken + if s.bucketCounts[shard] <= 0 { + delete(s.bucketCounts, shard) + delete(s.buckets, shard) + } + return out +} + +// handleReleaseSignals runs as a long-lived workflow coroutine, +// consuming ReleaseShards signals from in-flight activities. Each +// signal lists shards the activity considers complete; the handler +// clears them from heldByBatch[BatchID] (so the dispatch coroutine's +// defer won't double-release) and shardInFlight (so the packer can +// dispatch new work against them while the activity stays running on +// its still-pending shards). +// +// DO NOT add workflow yields (ExecuteActivity, Sleep, Await, etc.) +// between Receive and the next Receive. drainForCAN relies on +// ch.Len() == 0 implying "every delivered signal has been processed"; +// a yield mid-handler would invalidate that, leaving shardInFlight +// stale after a CAN. +func (s *shardedWorkflowState) handleReleaseSignals(ctx workflow.Context) { + ch := workflow.GetSignalChannel(ctx, releaseShardsSignalName) + for ctx.Err() == nil { + var payload releaseShardsPayload + if !ch.Receive(ctx, &payload) { + return + } + held, ok := s.heldByBatch[payload.BatchID] + if !ok { + continue + } + for _, sh := range payload.Shards { + if held[sh] { + delete(held, sh) + delete(s.shardInFlight, sh) + } + } + } +} + +// dispatchSlotAvailable returns true when the workflow is below the +// in-flight batch ceiling and is free to spawn another batch. Callers +// that can defer dispatch (the streaming packer) consult this and +// bail out; callers that must dispatch (resume payloads) pair it +// with waitForDispatchSlot. ConcurrentBatchCount is normalised to +// >= 1 at state construction, so no zero-disable path is needed. +func (s *shardedWorkflowState) dispatchSlotAvailable() bool { + return s.pendingDispatches < s.params.ConcurrentBatchCount +} + +// waitForDispatchSlot blocks the calling workflow coroutine until a +// dispatch slot frees up or lastErr trips. +func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { + _ = workflow.Await(ctx, func() bool { + return s.lastErr != nil || s.pendingDispatches < s.params.ConcurrentBatchCount + }) +} + +// dispatchResumeBatches turns the prior cycle's drain payload into a +// fresh round of resume activities, packed across shards up to +// BatchSize per batch. Each shard appears at most once across the +// payload (shardInFlight enforces that only one batch holds a shard +// at a time, and a shard only lands in a drain return while its +// owning batch still has unverified execs on it), so per-shard +// contributions are taken whole and no MaxExecsPerShard cap applies — +// resume carries no inject load so the per-shard blast-radius the +// streaming packer guards against doesn't exist here. +func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { + if len(s.params.ResumeShards) == 0 { + return + } + entries := make([]ResumeShard, 0, len(s.params.ResumeShards)) + for _, rs := range s.params.ResumeShards { + if runCount(rs.Execs) == 0 { + continue + } + entries = append(entries, rs) + } + slices.SortFunc(entries, func(a, b ResumeShard) int { + return int(a.Shard - b.Shard) + }) + + payload := BatchPayload{} + packed := 0 + packNoProgress := map[int32]time.Duration{} + flush := func() { + if packed == 0 { + return + } + // Block until a dispatch slot is free so resume payloads + // can't overshoot ConcurrentBatchCount on cycles that + // carried many shards across CAN. + s.waitForDispatchSlot(ctx) + if s.lastErr != nil { + return + } + for sh := range payload { + s.shardInFlight[sh] = true + } + s.spawnBatch(ctx, payload, true, packNoProgress) + payload = BatchPayload{} + packed = 0 + packNoProgress = map[int32]time.Duration{} + } + + for _, rs := range entries { + if s.lastErr != nil { + return + } + rsCount := runCount(rs.Execs) + // Flush before this shard if it would push us over the + // batch cap, so each batch stays within BatchSize. A + // single shard's contribution is taken whole — we don't + // split a shard across batches because shardInFlight only + // admits one batch per shard at a time. + if packed+rsCount > s.params.BatchSize && packed > 0 { + flush() + } + payload[rs.Shard] = rs.Execs + packNoProgress[rs.Shard] = rs.NoProgressDuration + packed += rsCount + if packed >= s.params.BatchSize { + flush() + } + } + flush() +} + +// runCount sums runs across BIDs in a single shard's payload entry. +func runCount(byBID map[string][]RunEntry) int { + n := 0 + for _, runs := range byBID { + n += len(runs) + } + return n +} + +// drainBuckets blocks until buckets are empty (success) or lastErr +// trips (failure). Each pass packs everything currently dispatchable, +// then awaits any change in pendingDispatches + shardInFlight so the +// next pass can attempt shards just freed by signal-release. +func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { + for { + if s.lastErr != nil { + return + } + for s.tryPackStreaming(ctx, true) { //nolint:revive + } + if s.bucketsEmpty() || s.lastErr != nil { + return + } + currentPending := s.pendingDispatches + if currentPending == 0 { + // Non-empty buckets but nothing in flight means the + // shard-claim bookkeeping is corrupted: tryPackStreaming + // declined to pack anything yet no batch is running to + // eventually free a shard. Returning silently would + // proceed to CAN (or completion) with execs still in + // buckets that were never dispatched — silent data loss. + // Fail the workflow instead so lastErr propagates out + // through run(). + remaining := 0 + for _, n := range s.bucketCounts { + remaining += n + } + s.lastErr = temporal.NewNonRetryableApplicationError( + fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), + "DrainBucketsStuck", nil) + return + } + // A "free shard" wake-up only counts when there's also a + // dispatch slot to use it, otherwise the outer loop would + // busy-spin on tryPackStreaming returning false against + // the in-flight cap. + _ = workflow.Await(ctx, func() bool { + if s.lastErr != nil { + return true + } + if s.pendingDispatches < currentPending { + return true + } + if !s.dispatchSlotAvailable() { + return false + } + for sh, n := range s.bucketCounts { + if n > 0 && !s.shardInFlight[sh] { + return true + } + } + return false + }) + } +} + +// drainForCAN cancels every in-flight batch and waits for them to +// return AND for the ReleaseShards signal channel to be drained. +// Activities honour cancellation by entering drain mode and returning +// a result whose InFlight carries their still-unverified execs; +// spawnBatch appends those entries to s.drainPayload. The signal +// channel drain is so a final ReleaseShards fired by an activity just +// before it returns doesn't get stranded mid-flight, which would +// leave shardInFlight set for shards the activity already considers +// complete. +// +// Channel.Len() is safe here because handleReleaseSignals has no +// yield points between Receive and the next blocking Receive, so +// Len() == 0 observed across an Await re-evaluation means every +// delivered signal has been processed. +// +// No explicit time bound: a well-behaved activity returns within +// req.DrainGrace (15s default) plus a small idle-cost slack; a +// misbehaved one is bounded by the activity's HeartbeatTimeout (1m). +// In practice this Await unblocks well under a minute. +func (s *shardedWorkflowState) drainForCAN(ctx workflow.Context) { + if s.pendingDispatches == 0 { + return + } + for _, cancel := range s.batchCancels { + cancel() + } + releaseCh := workflow.GetSignalChannel(ctx, releaseShardsSignalName) + _ = workflow.Await(ctx, func() bool { + if s.lastErr != nil { + return true + } + return s.pendingDispatches == 0 && releaseCh.Len() == 0 + }) +} + +// spawnBatch dispatches one batch on a new workflow.Go coroutine. +// Callers must have already marked every shard appearing as a +// top-level key in payload as shardInFlight (the "claim") so the +// packer can see them as busy while picking subsequent batches. The +// coroutine receives a per-batch cancellable ctx; the cancel func is +// stored in batchCancels so drainForCAN can cancel individual +// batches without tearing down the whole workflow. +func (s *shardedWorkflowState) spawnBatch( + ctx workflow.Context, + payload BatchPayload, + resume bool, + noProgressByShard map[int32]time.Duration, +) { + if payload.totalRuns() == 0 { + return + } + s.nextBatchID++ + batchID := s.nextBatchID + + req := &shardedBatchReq{ + BatchID: batchID, + Namespace: s.params.Namespace, + NamespaceID: s.namespaceID, + Executions: payload, + TargetClusterEndpoint: s.params.TargetClusterEndpoint, + TargetClusterName: s.params.TargetClusterName, + Resume: resume, + NoProgressByShard: noProgressByShard, + PerBatchGenerateRPS: s.params.PerBatchGenerateRPS, + ShardNoProgress: s.params.ShardNoProgress, + DrainGrace: s.params.DrainGrace, + IdleShardCost: s.params.IdleShardCost, + } + + batchCtx, cancel := workflow.WithCancel(ctx) + s.batchCancels[batchID] = cancel + held := make(map[int32]bool, len(payload)) + for sh := range payload { + held[sh] = true + } + s.heldByBatch[batchID] = held + s.batchExecs[batchID] = payload + + s.pendingDispatches++ + workflow.Go(ctx, func(coroCtx workflow.Context) { + defer func() { + s.pendingDispatches-- + // Clear any shards we still hold — signal-released + // shards have already been cleared from shardInFlight + // by handleReleaseSignals and may by now belong to a + // subsequent batch's claim. + for sh := range s.heldByBatch[batchID] { + delete(s.shardInFlight, sh) + } + delete(s.heldByBatch, batchID) + delete(s.batchCancels, batchID) + cancel() + }() + ao := workflow.ActivityOptions{ + StartToCloseTimeout: 24 * time.Hour, + HeartbeatTimeout: time.Minute, + // MaxAttempts=1: per-exec backoff inside the activity + // is the retry path; an activity-level retry would + // re-run inject and reset all backoff state, wasting + // the apply work this attempt already drove. + RetryPolicy: &temporal.RetryPolicy{ + MaximumAttempts: 1, + }, + // WaitForCancellation: the cancelled activity needs to + // run its drain logic and return its drain result + // before the dispatch coroutine's defer fires. + WaitForCancellation: true, + } + actx := workflow.WithActivityOptions(batchCtx, ao) + var result replicateBatchResult + err := workflow.ExecuteActivity(actx, shardedBatchActivityName, req).Get(coroCtx, &result) + if err == nil { + // Activity body ran and returned cleanly — either a + // clean completion (empty InFlight) or a drained + // CAN-cancel (InFlight carries the still-unverified + // execs). CompletedShards is informational; the + // defer above clears heldByBatch + shardInFlight + // either way. + if len(result.InFlight) > 0 { + s.drainPayload = append(s.drainPayload, result.InFlight...) + } + s.recordVerified(coroCtx, result.VerifiedCount) + delete(s.batchExecs, batchID) + return + } + if temporal.IsCanceledError(err) { + // Cancel-before-start: the activity body never ran, + // so no result is available. Leaving batchExecs[batchID] + // intact lets the CAN-end recovery path re-bucket the + // execs as fresh inject+verify work next cycle. + return + } + s.lastErr = err + }) +} + +// tryPackStreaming attempts to pack and dispatch one batch from +// s.buckets. Returns true if a batch was dispatched. +// +// No per-shard or total-bucket threshold: as soon as any free shard +// has any execs and a dispatch slot is open, a batch fires. Safety +// is enforced outside the packer — MaxExecsPerShard caps a single +// shard's contribution to a batch (so one hot shard can't dominate), +// ConcurrentBatchCount caps in-flight batches, and PerBatchGenerateRPS +// caps the per-batch source RPS. Within those bounds the packer's +// sole job is to make progress every chance it gets. +// +// relax=false (streaming, during ListWorkflows pages): pack fullest +// free shards first. When work is plentiful, batches land at +// BatchSize across few shards (small in-flight shard set per batch, +// nice and predictable); when it's sparse, the same loop ships a +// smaller batch rather than waiting and burning wall-clock on idle. +// +// relax=true (drain, after listing finishes — no more execs coming): +// hot shards (count > MaxExecsPerShard, i.e. those needing more than +// one round trip) first, fullest within hot. Total drain wall-clock +// is bounded by the heaviest shard's round-trip count × activity +// duration, so getting each hot shard's pipeline started ASAP is the +// dominant lever. Once they're all claimed, remaining batch capacity +// fills from smallest cold buckets so light shards finish and free +// their slots quickly — no point waiting on them to grow, nothing +// will arrive. +func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool) bool { + if s.lastErr != nil || s.params.BatchSize <= 0 || s.params.MaxExecsPerShard <= 0 { + return false + } + if !s.dispatchSlotAvailable() { + return false + } + shardIDs := s.shardIDsByPackPriority(relax) + if len(shardIDs) == 0 { + return false + } + + payload := BatchPayload{} + packed := 0 + for _, sh := range shardIDs { + room := s.params.BatchSize - packed + if room <= 0 { + break + } + take := min(s.params.MaxExecsPerShard, s.bucketCounts[sh], room) + if take == 0 { + continue + } + payload[sh] = s.takeFromBucket(sh, take) + packed += take + s.shardInFlight[sh] = true + } + if packed == 0 { + return false + } + s.spawnBatch(ctx, payload, false, nil) + return true +} + +// bucketsEmpty reports whether every shard's bucket is empty. Reads +// from the sidecar count map so it's O(#shards), not O(#runs). +func (s *shardedWorkflowState) bucketsEmpty() bool { + for _, n := range s.bucketCounts { + if n > 0 { + return false + } + } + return true +} + +// shardIDsByPackPriority returns free, non-empty shard IDs in the +// order the packer should consider them. Deterministic across +// replays: ordering is derived from workflow state (bucketCounts) +// with shard ID as a stable tiebreaker. +// +// relax=false (streaming): fullest first. Packer naturally produces +// large, predictable batches when work is plentiful and small ones +// when it isn't — either way it ships rather than waiting. +// +// relax=true (drain): hot shards (count > MaxExecsPerShard) first, +// fullest within hot. These need >1 round trip to drain, so total +// drain wall-clock is bounded by the heaviest shard; starting their +// pipelines first is the dominant lever. After all hot shards are +// claimed, remaining batch capacity fills from smallest cold buckets +// (ascending count) so light shards clear out quickly — nothing is +// arriving in drain, so waiting for cold buckets to grow is wasted +// wall-clock. +func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { + out := make([]int32, 0, len(s.bucketCounts)) + for sh, n := range s.bucketCounts { + if n == 0 || s.shardInFlight[sh] { + continue + } + out = append(out, sh) + } + if relax { + maxPerShard := s.params.MaxExecsPerShard + slices.SortFunc(out, func(a, b int32) int { + aHot, bHot := s.bucketCounts[a] > maxPerShard, s.bucketCounts[b] > maxPerShard + switch { + case aHot && !bHot: + return -1 + case !aHot && bHot: + return 1 + case aHot && bHot: + // Both hot: fullest first. + if d := s.bucketCounts[b] - s.bucketCounts[a]; d != 0 { + return d + } + default: + // Both cold: smallest first. + if d := s.bucketCounts[a] - s.bucketCounts[b]; d != 0 { + return d + } + } + return int(a - b) + }) + return out + } + slices.SortFunc(out, func(a, b int32) int { + if d := s.bucketCounts[b] - s.bucketCounts[a]; d != 0 { + return d + } + return int(a - b) + }) + return out +} diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go new file mode 100644 index 00000000000..55680a22ecb --- /dev/null +++ b/service/worker/migration/sharded_workflow_test.go @@ -0,0 +1,490 @@ +package migration + +import ( + "context" + "errors" + "fmt" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.temporal.io/api/workflowservice/v1" + "go.temporal.io/sdk/activity" + "go.temporal.io/sdk/converter" + "go.temporal.io/sdk/temporal" + "go.temporal.io/sdk/testsuite" + "go.temporal.io/sdk/workflow" + "go.temporal.io/server/common" +) + +// testNamespaceID is what metadataResponseFor returns. Tests pass it +// through to bidsForShards so the BIDs they hand to makeExecs hash +// to the same shards the workflow will compute during the page loop. +const testNamespaceID = "test-ns-id" + +// ---- Test setup helpers ---- + +// bidsForShards returns, for each shard the hash actually populates +// under common.WorkflowIDToHistoryShard(namespaceID, bid, totalShards), +// a slice of `perShard` BusinessIDs. Brute-force search over candidate +// strings — sufficient for the small shard counts the tests use. +func bidsForShards(namespaceID string, totalShards int32, perShard int) map[int32][]string { + out := make(map[int32][]string, totalShards) + for i := 0; ; i++ { + bid := fmt.Sprintf("wf-%d", i) + sh := common.WorkflowIDToHistoryShard(namespaceID, bid, totalShards) + if len(out[sh]) < perShard { + out[sh] = append(out[sh], bid) + } + if int32(len(out)) == totalShards { + done := true + for sh := range out { + if len(out[sh]) < perShard { + done = false + break + } + } + if done { + return out + } + } + } +} + +// makeExecs builds a slice of ExecutionInfos engineered to hash across +// `shards` distinct shards (`perShard` execs per shard) under the test +// namespace ID + shard count. Shard is left zero — the workflow's page +// loop populates it via common.WorkflowIDToHistoryShard. +func makeExecs(shards int32, perShard int) []*ExecutionInfo { + bids := bidsForShards(testNamespaceID, shards, perShard) + var execs []*ExecutionInfo + idx := 0 + for sh := range bids { + for _, bid := range bids[sh] { + execs = append(execs, &ExecutionInfo{ + BusinessID: bid, + RunID: "run-" + strconv.Itoa(idx), + }) + idx++ + } + } + return execs +} + +// pageThrough returns a function suitable for OnActivity("ListWorkflows") +// that paginates `all` into pages of `pageSize` execs each. The +// workflow populates ex.Shard after this returns, so callers can hand +// over ExecutionInfos with Shard left zero. +func pageThrough(all []*ExecutionInfo, pageSize int) func(context.Context, *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + return func(_ context.Context, req *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + start := 0 + if len(req.NextPageToken) > 0 { + start, _ = strconv.Atoi(string(req.NextPageToken)) + } + end := min(start+pageSize, len(all)) + var nextToken []byte + if end < len(all) { + nextToken = []byte(strconv.Itoa(end)) + } + return &listWorkflowsResponse{ + Executions: all[start:end], + NextPageToken: nextToken, + }, nil + } +} + +// metadataResponseFor returns a function suitable for +// OnActivity("GetMetadata") that yields a fixed shard count + ns ID. +func metadataResponseFor(shardCount int32) func(context.Context, MetadataRequest) (*MetadataResponse, error) { + return func(_ context.Context, _ MetadataRequest) (*MetadataResponse, error) { + return &MetadataResponse{ + ShardCount: shardCount, + NamespaceID: testNamespaceID, + }, nil + } +} + +// registerShardedScaffolding registers GetMetadata + CountWorkflow stubs. +// Every sharded workflow test exercises a different scenario (paging, +// resume, drain, ...) but they all need these two before the page loop +// runs, so factoring them out keeps each test focused on its scenario. +func registerShardedScaffolding(env *testsuite.TestWorkflowEnvironment, shardCount int32) { + env.RegisterActivityWithOptions(metadataResponseFor(shardCount), activity.RegisterOptions{Name: "GetMetadata"}) + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.CountWorkflowExecutionsRequest) (*countWorkflowResponse, error) { + return &countWorkflowResponse{WorkflowCount: 0}, nil + }, activity.RegisterOptions{Name: "CountWorkflow"}) +} + +// ---- Tests ---- + +// TestSharded_HappyPath_SingleCycle: a small workload exhausts in one +// cycle, every batch returns clean completion, no CAN, no resume. +func TestSharded_HappyPath_SingleCycle(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + + execs := makeExecs(4, 5) // 20 execs across 4 shards + registerShardedScaffolding(env, 4) + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + var ( + mu sync.Mutex + batchesSeen int + execsSeen int + shardsSeen = map[int32]struct{}{} + ) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + mu.Lock() + batchesSeen++ + execsSeen += req.Executions.totalRuns() + for _, sh := range req.Executions.sortedShards() { + shardsSeen[sh] = struct{}{} + } + mu.Unlock() + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 4, + }) + + require.True(t, env.IsWorkflowCompleted(), "workflow should complete") + require.NoError(t, env.GetWorkflowError(), "workflow should succeed") + // 20 execs across 4 shards, BatchSize=100, MaxExecsPerShard=50, single + // page — every exec fits in one batch. + require.Equal(t, 1, batchesSeen, "all execs should pack into a single batch") + require.Equal(t, len(execs), execsSeen, "every exec should reach the activity") + require.Len(t, shardsSeen, 4, "every shard should be represented") +} + +// TestSharded_ResumeShards_Packed: a non-empty ResumeShards in params +// gets packed into multi-shard batches up to BatchSize. Asserts that +// no resume batch exceeds BatchSize and the per-shard contributions +// match the input. +func TestSharded_ResumeShards_Packed(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 8) + env.RegisterActivityWithOptions(pageThrough(nil, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + // 8 shards, 15 unverified execs each = 120 total. With BatchSize=100, + // we expect 2 batches: e.g., [0..5] (90 execs) + [5+, 6, 7] (30) or + // similar — depends on greedy packing. + resumeShards := make([]ResumeShard, 8) + for s := range 8 { + resumeShards[s] = ResumeShard{ + Shard: int32(s), + Execs: makeExecsForShard(int32(s), 15), + } + } + + var ( + mu sync.Mutex + batches [][]int32 + ) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + require.True(t, req.Resume, "all resume-dispatched batches must have Resume=true") + require.LessOrEqual(t, req.Executions.totalRuns(), 100, "batch must not exceed BatchSize") + mu.Lock() + batches = append(batches, req.Executions.sortedShards()) + mu.Unlock() + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 8, + ResumeShards: resumeShards, + }) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + // Confirm every shard was covered exactly once across all batches. + covered := map[int32]int{} + for _, b := range batches { + for _, sh := range b { + covered[sh]++ + } + } + for s := range int32(8) { + require.Equal(t, 1, covered[s], "shard %d should be covered exactly once", s) + } + require.GreaterOrEqual(t, len(batches), 2, "should pack into at least 2 batches given 120 execs / BatchSize=100") +} + +// TestSharded_ReleaseShards_FreesShardForReuse: an activity sends a +// ReleaseShards signal mid-flight. The workflow must clear the shard +// from shardInFlight so the packer can dispatch a fresh batch +// targeting that shard *while the original activity is still running* +// — the slot in ConcurrentBatchCount is still claimed by batch 1, so +// batch 2 can only dispatch if signal-release worked. +// +// ConcurrentBatchCount=2 is explicit: defaultConcurrentBatchCount(2)=1 +// would gate batch 2 on the dispatch slot regardless of shard state, so +// the test couldn't distinguish "release-from-flight" from +// "batch 1 returned and freed the slot". +func TestSharded_ReleaseShards_FreesShardForReuse(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + + // Two pages of execs on the same shards. The second batch can dispatch + // only when batch 1 releases its shards. + phase1 := makeExecs(2, 10) + phase2 := makeExecs(2, 10) + all := append(append([]*ExecutionInfo{}, phase1...), phase2...) + env.RegisterActivityWithOptions(pageThrough(all, 20), activity.RegisterOptions{Name: "ListWorkflows"}) + + var ( + mu sync.Mutex + batches []*shardedBatchReq + secondStarted = make(chan struct{}) + secondOnce sync.Once + releaseObserved atomic.Bool + ) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + mu.Lock() + batches = append(batches, req) + idx := len(batches) + mu.Unlock() + + switch idx { + case 1: + // Signal-release, then block here until batch 2 actually + // dispatches. If signal-release wires through, the workflow + // dispatches batch 2 concurrently; if it doesn't, batch 2 + // can't run until this activity returns (the safety timeout + // below). + env.SignalWorkflow("ReleaseShards", releaseShardsPayload{ + BatchID: req.BatchID, + Shards: req.Executions.sortedShards(), + }) + select { + case <-secondStarted: + releaseObserved.Store(true) + case <-time.After(30 * time.Second): + // Safety release so the test fails the assertion rather + // than hanging indefinitely. Generous bound because CI + // can be slow; the happy path returns immediately. + } + case 2: + secondOnce.Do(func() { close(secondStarted) }) + } + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 2, + ConcurrentBatchCount: 2, + }) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + require.Len(t, batches, 2, "expected two batches across the two pages") + require.True(t, releaseObserved.Load(), + "signal-release should let batch 2 dispatch while batch 1 still holds its dispatch slot") +} + +// TestSharded_ShardNoProgress_FailsWorkflow: activity returns a +// non-retryable ShardNoProgress error → workflow fails with that +// error, no CAN. +func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + env.RegisterActivityWithOptions(pageThrough(makeExecs(2, 5), 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + shards := req.Executions.sortedShards() + return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( + "shard "+strconv.Itoa(int(shards[0]))+" stuck", "ShardNoProgress", nil) + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 2, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + var appErr *temporal.ApplicationError + require.True(t, errors.As(err, &appErr)) + require.Equal(t, "ShardNoProgress", appErr.Type()) +} + +// TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover: a +// non-empty InFlight in the activity's replicateBatchResult must +// populate drainPayload and end up as ResumeShards in the CAN +// carry-over. +// +// Tests the workflow plumbing only. In production, an activity +// returns InFlight after entering drain mode and grace-expiring; +// the testsuite has a known bug where activity-cancellation result +// handling loses payload fidelity, so we exercise the same code +// path by returning InFlight from a non-cancelled run instead. The +// dispatch coroutine's err == nil branch is what we're testing — +// it doesn't care whether the activity was cancelled or not, only +// whether the returned result has InFlight to fold into drainPayload. +func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 10) + + // Page 1 returns a multi-shard population large enough to trip + // the streaming packer's mid-cycle dispatch under the pinned + // BatchSize/MaxExecsPerShard the test sets below + // (trigger=MaxExecsPerShard/2=5, minShards=BatchSize/MaxExecsPerShard=10). + // The activity + // flips CAN-suggested from inside its body before returning, so + // by the time it has handed back its InFlight the workflow is + // committed to CAN — but without going through cancel, which + // the testsuite delivers as a CanceledError regardless of any + // result the activity returned. + pageExecs := makeExecs(10, 10) + // Drained exec mirrors a real input row so the simulated drain + // payload would be a valid response from a real activity. drainedBID + // is the first exec; drainedShard is the shard the workflow will + // hash it to. + drainedBID := pageExecs[0].BusinessID + const drainedRunID = "run-drained" + drainedShard := common.WorkflowIDToHistoryShard(testNamespaceID, drainedBID, 10) + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + return &listWorkflowsResponse{ + Executions: pageExecs, + NextPageToken: []byte("more"), + }, nil + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + env.SetContinueAsNewSuggested(true) + return replicateBatchResult{ + InFlight: []ResumeShard{{ + Shard: drainedShard, + Execs: map[string][]RunEntry{drainedBID: {{RunID: drainedRunID}}}, + NoProgressDuration: 42 * time.Second, + }}, + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 10, + BatchSize: 100, + MaxExecsPerShard: 10, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err, "workflow should CAN, not return success") + + var canErr *workflow.ContinueAsNewError + require.True(t, errors.As(err, &canErr), "error should be ContinueAsNewError") + + var nextParams ShardedForceReplicationParams + require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) + require.NotEmpty(t, nextParams.ResumeShards, "InFlight from a returned activity should appear in resume payload") + require.Equal(t, drainedShard, nextParams.ResumeShards[0].Shard) + runs, ok := nextParams.ResumeShards[0].Execs[drainedBID] + require.True(t, ok, "drained business ID should appear in nested resume payload") + require.Equal(t, []RunEntry{{RunID: drainedRunID}}, runs) + require.Equal(t, 42*time.Second, nextParams.ResumeShards[0].NoProgressDuration) +} + +// TestSharded_CancelBeforeStart_NoLostExecs: this is the reproducer +// for the sim bug — when an activity is dispatched and the workflow +// CANs before the activity body runs (testsuite-scheduler race), the +// activity returns CanceledError with no result. The recovery path +// re-buckets the input execs into RecoveredBuckets so the next cycle +// dispatches them as fresh inject+verify batches. This test pins +// down that behaviour. +func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 4) + + execs := makeExecs(4, 5) // 20 execs across 4 shards + pageServed := false + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + if pageServed { + return &listWorkflowsResponse{}, nil + } + pageServed = true + // Trigger CAN as soon as this page is served so drainForCAN + // runs before any dispatched batches can complete. + env.SetContinueAsNewSuggested(true) + return &listWorkflowsResponse{ + Executions: execs, + NextPageToken: []byte("more"), + }, nil + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + // Activity that responds to ctx cancellation by returning a + // CanceledError with no result — simulating the + // "cancelled before any work done" path. + var ( + batchCount atomic.Int32 + cancelledIDs []int64 + muIDs sync.Mutex + ) + env.RegisterActivityWithOptions(func(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + batchCount.Add(1) + <-ctx.Done() + muIDs.Lock() + cancelledIDs = append(cancelledIDs, req.BatchID) + muIDs.Unlock() + return replicateBatchResult{}, temporal.NewCanceledError() + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterShardCount: 4, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + var canErr *workflow.ContinueAsNewError + require.True(t, errors.As(err, &canErr), "expected CAN, got %v", err) + + var nextParams ShardedForceReplicationParams + require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) + + // Count execs the next cycle would re-dispatch. The activity + // body never ran in this race, so execs land in RecoveredBuckets + // (fresh inject+verify) rather than ResumeShards (verify-only). + recovered := nextParams.RecoveredBuckets.totalRuns() + + t.Logf("dispatched %d batches, cancelled %d, recovered execs %d (expected %d)", + batchCount.Load(), len(cancelledIDs), recovered, len(execs)) + + require.Equal(t, len(execs), recovered, "every dispatched exec should land in RecoveredBuckets when its activity is cancelled before it can run") + require.Empty(t, nextParams.ResumeShards, "no ResumeShards — activity never ran, never injected, so no resume work") +} + +// ---- internal helpers ---- + +// makeExecsForShard produces `count` runs for the named shard's +// ResumeShard.Execs payload. Each run gets a distinct businessID so +// the resulting map has `count` entries with one run each — the +// simplest shape for tests that don't care about BID-reuse. +func makeExecsForShard(shard int32, count int) map[string][]RunEntry { + out := map[string][]RunEntry{} + for i := range count { + bid := "wf-" + strconv.Itoa(int(shard)) + "-" + strconv.Itoa(i) + out[bid] = []RunEntry{{RunID: "run-" + strconv.Itoa(int(shard)*1000+i)}} + } + return out +} From 2a05cacf7e3d954d0057744bf76cac4937f57b18 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 07:43:09 +0100 Subject: [PATCH 02/35] Suppress duplicate-registration panic on default worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit worker.go's upgrade-hack pass registers each component's activities on the default worker before the dedicated worker (see the TODO at worker.go:82). The legacy and sharded WorkerComponents share the *activities method set, so whichever runs second hits the SDK's already-registered check and panics on CountWorkflow / ListWorkflows / etc. Both call sites now use activity.RegisterOptions{DisableAlreadyRegisteredCheck: true} — the default worker isn't dispatched to by either workflow (both have dedicated activity workers), so winner-takes-all on those duplicate registrations is harmless. --- service/worker/migration/fx.go | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index dc3021ce08c..d1c8a72014b 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -5,6 +5,7 @@ import ( "fmt" "go.temporal.io/api/workflowservice/v1" + "go.temporal.io/sdk/activity" sdkworker "go.temporal.io/sdk/worker" "go.temporal.io/sdk/workflow" "go.temporal.io/server/api/adminservice/v1" @@ -113,7 +114,16 @@ func (wc *replicationWorkerComponent) DedicatedWorkflowWorkerOptions() *workerco } func (wc *replicationWorkerComponent) RegisterActivities(registry sdkworker.Registry) { - registry.RegisterActivity(wc.activities) + // DisableAlreadyRegisteredCheck because the sharded WorkerComponent + // shares the *activities method set; whichever component registers + // first on the default worker wins (per the worker.go upgrade-hack + // pass), and the second component's reflection-based registration + // would otherwise panic on every method name. The default worker + // isn't dispatched to by either workflow — both have dedicated + // activity workers — so winner-takes-all is fine. + registry.RegisterActivityWithOptions(wc.activities, activity.RegisterOptions{ + DisableAlreadyRegisteredCheck: true, + }) } func (wc *replicationWorkerComponent) DedicatedActivityWorkerOptions() *workercommon.DedicatedWorkerOptions { @@ -155,7 +165,12 @@ func (sc *shardedWorkerComponent) DedicatedWorkflowWorkerOptions() *workercommon } func (sc *shardedWorkerComponent) RegisterActivities(registry sdkworker.Registry) { - registry.RegisterActivity(sc.activities) + // See replicationWorkerComponent.RegisterActivities — both components + // share the *activities method set, so the second registration on the + // default worker would otherwise panic. + registry.RegisterActivityWithOptions(sc.activities, activity.RegisterOptions{ + DisableAlreadyRegisteredCheck: true, + }) } func (sc *shardedWorkerComponent) DedicatedActivityWorkerOptions() *workercommon.DedicatedWorkerOptions { From 707be4ad60a017f4e0650431e077c26b4a0ad90d Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 08:45:54 +0100 Subject: [PATCH 03/35] Lint. Also use a single cancellation context for all activities as we never cancel just one. --- service/worker/migration/sharded_activity.go | 453 +++++++++--------- service/worker/migration/sharded_types.go | 64 +-- service/worker/migration/sharded_workflow.go | 143 +++--- .../worker/migration/sharded_workflow_test.go | 64 ++- 4 files changed, 359 insertions(+), 365 deletions(-) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index 06b901dacec..d65b897e41f 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -21,20 +21,16 @@ import ( ) // ReplicateBatch is the per-batch activity body for the sharded force -// replication workflow. It runs inject (skip on Resume) followed by -// verify, signal-releasing completed shards mid-flight once their -// cumulative idle cost crosses IdleShardCost. On workflow-initiated -// cancellation it enters drain mode: continues verifying for up to -// DrainGrace, then returns a replicateBatchResult carrying any -// still-unverified execs grouped by shard. Drain-mode signal traffic -// is intentionally suppressed — once we know we're about to return, -// there's no point racing a signal against the return value. +// replication workflow. Runs inject (skipped on Resume) then verify, +// signal-releasing completed shards mid-flight as their cumulative +// idle cost crosses IdleShardCost. On workflow-initiated cancellation +// it enters drain mode and returns a replicateBatchResult carrying +// any still-unverified execs. Drain-mode signal traffic is suppressed: +// once we know we're about to return, there's no point racing a +// signal against the return value. func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { - // Flatten the nested wire shape once on entry so the per-exec - // bookkeeping (verified[], attempts[], nextRetryAt[]) can stay - // index-based. flatten() walks shards ascending then BIDs - // alphabetical, so the order is deterministic — replays of the same - // payload produce the same slice. + // Flatten once so per-exec bookkeeping (verified[], attempts[], + // nextRetryAt[]) can stay index-based. execs := req.Executions.flatten() of := len(execs) if of == 0 { @@ -45,40 +41,17 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( req.TargetClusterEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout) // Namespace lookup feeds the verify phase's retention/zombie skip - // check (checkSkipWorkflowExecution needs ns.Retention()). Looked up - // once per activity, since the namespace registry is local and - // immutable across the run. + // check (checkSkipWorkflowExecution needs ns.Retention()). Snapshotted + // once per activity; we don't track config changes mid-batch. ns, err := a.NamespaceRegistry.GetNamespaceByID(namespace.ID(req.NamespaceID)) if err != nil { return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) } - // ---- Inject phase (skipped on Resume: the execs have already been - // injected by an earlier activity that was drained for CAN). + // ---- Inject phase ---- if !req.Resume { - // One limiter per activity invocation, sized for RPS — the - // post-call ReserveN reservation pulls extra tokens proportional - // to history size, so a single limiter shared across this - // batch's execs is what enforces the per-batch RPS budget - // end-to-end. - rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) - for _, ex := range execs { - if err := ctx.Err(); err != nil { - // Worker shutdown or workflow drain mid-inject. Return a - // recognizable CanceledError so spawnBatch's - // IsCanceledError check fires and the batch's batchExecs - // entry is preserved for RecoveredBuckets re-injection - // next cycle. Any execs already injected here are - // re-injected harmlessly — replication dedupes per - // (namespace, wf, run). - return replicateBatchResult{}, temporal.NewCanceledError("inject phase cancelled") - } - if err := a.generateReplicationTaskForExec(ctx, rateLimiter, req, ex); err != nil { - if ctx.Err() != nil { - return replicateBatchResult{}, temporal.NewCanceledError("inject phase cancelled") - } - return replicateBatchResult{}, err - } + if err := a.runInjectPhase(ctx, req, execs); err != nil { + return replicateBatchResult{}, err } } @@ -93,36 +66,25 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( var draining bool var drainStartAt time.Time - // callCtx is what attemptVerifyExec uses for DescribeMutableState. In - // normal mode it's the activity's parent ctx; when drain mode kicks - // in, it swaps to a fresh detached context with a DrainGrace timeout - // so DMS calls keep working long enough for nearly-verified execs to - // land. The parent ctx is already dead by the time we transition - // (that's what triggers the transition), so using it for drain-mode - // RPCs would defeat the grace window — every call would fail - // instantly. + // callCtx is what attemptVerifyExec uses for DescribeMutableState. + // In drain mode we swap to a detached context: the parent ctx is + // already dead (that's what triggered the transition), so reusing + // it would make every drain-mode RPC fail instantly. callCtx := ctx // Pre-create the drain context up front and defer cancel right away - // so go vet's lostcancel pass sees the canonical pattern. The drain - // timer (started below on the cancel transition) plus the unconditional - // defer cancel make the lifetime obvious; the context only matters - // once draining is set, so creating it early costs nothing. + // so go vet's lostcancel pass sees the canonical pattern. drainCtx, drainCancel := context.WithCancel(context.Background()) defer drainCancel() for { // Worker shutdown short-circuits drain mode entirely. The SDK - // closes activity.GetWorkerStopChannel a fixed WorkerStopTimeout - // (10s default) before forcibly returning; burning that window - // on DescribeMutableState calls that won't get to drive their - // results back is worse than returning current state and letting - // the next cycle's ResumeShards / RecoveredBuckets paths recover. - // - // Checked at the top of each outer iteration (not just on initial - // drain transition) because worker shutdown can fire after we've - // already entered workflow-initiated drain — e.g. CAN cancel - // arrives, drain starts under detached ctx, then a deploy hits - // mid-window. The detached ctx wouldn't notice on its own. + // closes WorkerStopChannel WorkerStopTimeout before forcibly + // returning; burning that window on DMS calls that can't drive + // their results back is worse than returning current state and + // letting ResumeShards / RecoveredBuckets recover next cycle. + // Re-checked each iteration because shutdown can fire after + // we've already entered drain — the detached ctx wouldn't + // notice on its own. select { case <-activity.GetWorkerStopChannel(ctx): return replicateBatchResult{ @@ -139,62 +101,22 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( // workflow blocks for us, so the grace window is genuinely // available — swap callCtx onto a detached deadline so // DescribeMutableState keeps working after the parent ctx died. - if !draining { - if err := ctx.Err(); err != nil { - draining = true - drainStartAt = time.Now() - // Start the drain budget timer here rather than at activity - // entry so the grace window measures from drain transition, - // not from activity start. - time.AfterFunc(req.DrainGrace, drainCancel) - callCtx = drainCtx - } + if !draining && ctx.Err() != nil { + draining = true + drainStartAt = time.Now() + // Start the drain budget timer here rather than at activity + // entry so the grace window measures from drain transition, + // not from activity start. + time.AfterFunc(req.DrainGrace, drainCancel) + callCtx = drainCtx } - now := time.Now() - var minNextRetry time.Time - ctxAborted := false - for i := range of { - if verified[i] { - continue - } - if !nextRetryAt[i].IsZero() && nextRetryAt[i].After(now) { - if minNextRetry.IsZero() || nextRetryAt[i].Before(minNextRetry) { - minNextRetry = nextRetryAt[i] - } - continue - } - - ex := execs[i] - shard := ex.Shard - ok, err := a.attemptVerifyExec(callCtx, remoteAdminClient, ns, req, ex) - if err != nil { - if callCtx.Err() != nil { - // callCtx is dead. Two cases: (1) normal-mode parent - // ctx was just cancelled mid-call — the outer-loop top - // will promote to drain on the next iteration; - // (2) drain-mode detached ctx expired — the drain - // exit check below will fire. Either way, drop out of - // the inner loop now. - ctxAborted = true - break - } - return replicateBatchResult{}, err - } - - if ok { - verified[i] = true - doneCount++ - shards.recordVerified(shard, time.Now()) - continue - } - - attempts[i]++ - nextRetryAt[i] = time.Now().Add(backoffDelay(attempts[i])) - if minNextRetry.IsZero() || nextRetryAt[i].Before(minNextRetry) { - minNextRetry = nextRetryAt[i] - } + passDelta, minNextRetry, ctxAborted, vErr := a.runVerifyPass( + callCtx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) + if vErr != nil { + return replicateBatchResult{}, vErr } + doneCount += passDelta activity.RecordHeartbeat(ctx, doneCount) @@ -206,48 +128,25 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( }, nil } - // Per-shard cumulative no-progress backstop. Trips on any shard - // still holding pending execs whose last verified outcome - // (carried across CAN via tracker seeding) is older than - // ShardNoProgress. - if stuckShard, stuckDur, ok := shards.pickStuck(time.Now(), req.ShardNoProgress); ok { - stuckIdx := a.firstUnverifiedOnShard(execs, verified, stuckShard) - stuck := execs[stuckIdx] - return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( - fmt.Sprintf("shard %d no progress for %v on %s/%s (%d/%d done)", - stuckShard, stuckDur, stuck.BusinessID, stuck.RunID, doneCount, of), - "ShardNoProgress", nil) + // Per-shard cumulative no-progress backstop. + if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, of); sErr != nil { + return replicateBatchResult{}, sErr } - // Drain-mode exit checks. Drain mode is entered when the workflow - // has cancelled the activity for CAN. No signals here — the - // return value carries everything the workflow needs (completed - // shards + unverified execs grouped by shard with their cumulative - // no-progress duration). if draining { - elapsedDrain := time.Since(drainStartAt) - if elapsedDrain >= req.DrainGrace || shards.totalIdleCost(time.Now()) >= req.IdleShardCost { + // Drain-mode exit checks. No signals here — the return value + // carries everything the workflow needs (completed shards + + // unverified execs grouped by shard with their cumulative + // no-progress duration). + if a.shouldExitDrain(req, shards, drainStartAt) { return replicateBatchResult{ CompletedShards: shards.allCompleted(), InFlight: a.buildInFlight(execs, verified, shards, time.Now()), VerifiedCount: int64(doneCount), }, nil } - } else { - // Normal mode: if the cumulative idle cost across - // completed-but-not-yet-signaled shards crosses the threshold, - // signal-release them so the workflow can dispatch new batches - // against those shards while this activity keeps draining its - // still-pending ones. - if shards.totalIdleCost(time.Now()) >= req.IdleShardCost { - releaseList := shards.awaitingRelease() - if len(releaseList) > 0 { - if err := a.signalReleaseShards(ctx, req, releaseList); err != nil { - return replicateBatchResult{}, err - } - shards.markReleased(releaseList) - } - } + } else if err := a.maybeSignalRelease(ctx, req, shards); err != nil { + return replicateBatchResult{}, err } // If the inner loop aborted because callCtx died, skip the sleep @@ -257,46 +156,189 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( continue } - // Sleep until the next exec is due for retry. Drain mode honours - // the same scheduling so we don't burn the grace window - // busy-spinning when every remaining exec is in backoff. - sleepDur := 50 * time.Millisecond - if !minNextRetry.IsZero() { - if delta := time.Until(minNextRetry); delta > sleepDur { - sleepDur = delta - } + a.waitNextTick(ctx, callCtx, minNextRetry, draining, drainStartAt, req.DrainGrace) + } +} + +// runInjectPhase walks execs in flattened order, generating one +// replication task per exec under a per-batch RPS limiter. Cancellation +// mid-loop returns a CanceledError so spawnBatch's IsCanceledError +// check fires and the batch's batchExecs entry is preserved for +// RecoveredBuckets re-injection next cycle. Already-injected execs +// are re-injected harmlessly — replication dedupes per (namespace, +// wf, run). +func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, execs []*shardedExecutionInfo) error { + rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) + for _, ex := range execs { + if ctx.Err() != nil { + return temporal.NewCanceledError("inject phase cancelled") } - if draining { - remaining := req.DrainGrace - time.Since(drainStartAt) - if remaining > 0 && remaining < sleepDur { - sleepDur = remaining + if err := a.generateReplicationTaskForExec(ctx, rateLimiter, req, ex); err != nil { + if ctx.Err() != nil { + return temporal.NewCanceledError("inject phase cancelled") } - // Parent ctx is already dead in drain mode, so the - // select-on-ctx.Done() in normal mode would tight-loop here. - // Use the detached drain ctx (and a pure wall-clock fallback) - // instead. - select { - case <-time.After(sleepDur): - case <-callCtx.Done(): - } - } else { - select { - case <-time.After(sleepDur): - case <-ctx.Done(): - // ctx cancel just sets draining on the next iteration; - // don't unwind here. + return err + } + } + return nil +} + +// runVerifyPass runs one pass over every unverified exec, attempting a +// verify on those whose backoff timer has expired. Returns the count of +// execs newly verified this pass, the earliest pending retry deadline +// (for sleep scheduling), and whether callCtx died mid-pass — in which +// case the outer loop's top reassesses (drain transition or exit). +// Returns a non-nil error only for hard errors from the verify path; +// ctx-derived errors set ctxAborted instead so the outer loop owns +// the decision about what to do next. +func (a *activities) runVerifyPass( + callCtx context.Context, + remoteAdminClient adminservice.AdminServiceClient, + ns *namespace.Namespace, + req *shardedBatchReq, + execs []*shardedExecutionInfo, + verified []bool, + attempts []int, + nextRetryAt []time.Time, + shards shardVerifyTracker, +) (int, time.Time, bool, error) { + now := time.Now() + var minNextRetry time.Time + verifiedDelta := 0 + for i, ex := range execs { + if verified[i] { + continue + } + if !nextRetryAt[i].IsZero() && nextRetryAt[i].After(now) { + minNextRetry = earliest(minNextRetry, nextRetryAt[i]) + continue + } + + ok, err := a.attemptVerifyExec(callCtx, remoteAdminClient, ns, req, ex) + if err != nil { + if callCtx.Err() != nil { + // callCtx is dead. Two cases: (1) normal-mode parent ctx + // was just cancelled mid-call — the outer-loop top will + // promote to drain on the next iteration; (2) drain-mode + // detached ctx expired — the drain exit check fires. + return verifiedDelta, minNextRetry, true, nil } + return verifiedDelta, minNextRetry, false, err } + + if ok { + verified[i] = true + verifiedDelta++ + shards.recordVerified(ex.Shard, time.Now()) + continue + } + + attempts[i]++ + nextRetryAt[i] = time.Now().Add(backoffDelay(attempts[i])) + minNextRetry = earliest(minNextRetry, nextRetryAt[i]) + } + return verifiedDelta, minNextRetry, false, nil +} + +// earliest returns the earlier of cur (which may be zero) and candidate. +// Used to track the next-due retry deadline across the verify pass. +func earliest(cur, candidate time.Time) time.Time { + if cur.IsZero() || candidate.Before(cur) { + return candidate + } + return cur +} + +// checkStuckShard fails non-retryably if any shard has gone longer than +// req.ShardNoProgress without a verified outcome. Duration is cumulative +// across CAN cycles via tracker seeding. +func (a *activities) checkStuckShard( + req *shardedBatchReq, + shards shardVerifyTracker, + execs []*shardedExecutionInfo, + verified []bool, + doneCount, total int, +) error { + stuckShard, stuckDur, ok := shards.pickStuck(time.Now(), req.ShardNoProgress) + if !ok { + return nil + } + msg := fmt.Sprintf("shard %d no progress for %v", stuckShard, stuckDur) + if stuckIdx, found := a.firstUnverifiedOnShard(execs, verified, stuckShard); found { + stuck := execs[stuckIdx] + msg = fmt.Sprintf("shard %d no progress for %v on %s/%s (%d/%d done)", + stuckShard, stuckDur, stuck.BusinessID, stuck.RunID, doneCount, total) + } + return temporal.NewNonRetryableApplicationError(msg, "ShardNoProgress", nil) +} + +// shouldExitDrain reports whether the drain-mode exit conditions are +// met: either the grace window has expired, or the cumulative idle cost +// across completed-but-unsignaled shards crossed the threshold. +func (a *activities) shouldExitDrain(req *shardedBatchReq, shards shardVerifyTracker, drainStartAt time.Time) bool { + if time.Since(drainStartAt) >= req.DrainGrace { + return true + } + return shards.totalIdleCost(time.Now()) >= req.IdleShardCost +} + +// maybeSignalRelease signals the workflow to release any +// completed-but-unsignaled shards if their cumulative idle cost crossed +// the threshold. Only fires in normal mode — drain mode rides the +// activity result instead. +func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchReq, shards shardVerifyTracker) error { + if shards.totalIdleCost(time.Now()) < req.IdleShardCost { + return nil } + releaseList := shards.awaitingRelease() + if len(releaseList) == 0 { + return nil + } + if err := a.signalReleaseShards(ctx, req, releaseList); err != nil { + return err + } + shards.markReleased(releaseList) + return nil } -// generateReplicationTaskForExec injects one execution into the -// replication queue. Delegates to generateWorkflowReplicationTask so the -// rateLimiter wait, frontend-vs-history RPC choice, archetype lookup, -// and history-size-proportional token reservation stay in one place — -// sharded only wraps it to supply a single-element TargetClusters slice -// and to thread the dynamic-config-driven generateViaFrontend flag -// through. +// waitNextTick sleeps until the next exec is due for retry, capped by +// DrainGrace remaining when in drain mode. In drain mode the parent ctx +// is already dead, so we wake on the detached drain ctx instead — using +// the parent ctx would tight-loop on its Done channel. +func (a *activities) waitNextTick( + ctx, callCtx context.Context, + minNextRetry time.Time, + draining bool, + drainStartAt time.Time, + drainGrace time.Duration, +) { + sleepDur := 50 * time.Millisecond + if !minNextRetry.IsZero() { + if delta := time.Until(minNextRetry); delta > sleepDur { + sleepDur = delta + } + } + if draining { + if remaining := drainGrace - time.Since(drainStartAt); remaining > 0 && remaining < sleepDur { + sleepDur = remaining + } + select { + case <-time.After(sleepDur): + case <-callCtx.Done(): + } + return + } + select { + case <-time.After(sleepDur): + case <-ctx.Done(): + // ctx cancel just sets draining on the next iteration; don't + // unwind here. + } +} + +// generateReplicationTaskForExec is the per-exec inject wrapper around +// generateWorkflowReplicationTask; supplies the sharded-only +// single-target-clusters slice and the generateViaFrontend flag. func (a *activities) generateReplicationTaskForExec( ctx context.Context, rateLimiter quotas.RateLimiter, @@ -320,23 +362,13 @@ func (a *activities) generateReplicationTaskForExec( } // attemptVerifyExec runs the source-describe + target-applied check for -// a single execution. Mirrors verifySingleReplicationTask's -// classification — DescribeMutableState on the remote followed by -// workflowVerifier content comparison on success or -// checkSkipWorkflowExecution on NotFound — and returns whether the exec -// is now verified. +// a single execution and returns whether it's now verified. // -// Per-classification metric counters are emitted inline (success, -// pending, not-found, busy-workflow, failed). The caller only needs the -// verified bit; not-verified outcomes (missing/busy) drive per-exec -// backoff identically but stay distinct in metrics so a "passive -// cluster apply is in progress" signal stays visible. -// -// We do the DescribeMutableState call inline rather than delegating -// straight to verifySingleReplicationTask so the busy-workflow branch -// keeps its distinct counter. The existing helper folds BUSY_WORKFLOW -// into the generic notVerified result, losing the signal that the -// passive cluster is making progress. +// Why this isn't a delegation to verifySingleReplicationTask: that +// helper folds BUSY_WORKFLOW into the generic notVerified result. We +// inline the DMS call here to keep busy-workflow as a distinct metric +// counter, preserving the "passive cluster apply is in progress" +// signal. func (a *activities) attemptVerifyExec( ctx context.Context, remoteAdminClient adminservice.AdminServiceClient, @@ -449,7 +481,6 @@ type shardVerify struct { lastProgress time.Time // wall time of the most recent verified outcome } -// shardVerifyTracker is keyed by history shard ID. type shardVerifyTracker map[int32]shardVerify func newShardVerifyTracker( @@ -560,10 +591,8 @@ func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (i // buildInFlight groups unverified execs by shard then businessID and // attaches the cumulative no-progress duration per shard, for the -// drain-mode activity return. Shards with zero unverified execs are not -// included — fully verified shards are reported via CompletedShards (and -// any that completed mid-flight have already been signal-released so the -// workflow's packer could reuse them). +// drain-mode activity return. Shards with zero unverified execs are +// reported via CompletedShards instead. func (a *activities) buildInFlight( execs []*shardedExecutionInfo, verified []bool, @@ -604,23 +633,19 @@ func (a *activities) buildInFlight( // firstUnverifiedOnShard returns the index of the first execution in the // flattened execs slice that targets the given shard and hasn't verified -// yet. Used to name a concrete (BusinessID, RunID) in the ShardNoProgress -// failure event. -// -// Panics if no such execution exists — the only caller is the stuck-shard -// backstop, which by construction only fires for shards with at least one -// pending exec. A silent fallback to index 0 would hide a future -// invariant violation behind a misleading error message. -func (a *activities) firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shard int32) int { +// yet, and a found flag. Callers should only invoke this for shards +// with at least one pending exec; the found=false return is a defensive +// fallback so a tracker / verified-slice drift can't crash the activity. +func (a *activities) firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shard int32) (int, bool) { for i, ex := range execs { if verified[i] { continue } if ex.Shard == shard { - return i + return i, true } } - panic(fmt.Sprintf("firstUnverifiedOnShard: no unverified exec on shard %d", shard)) + return 0, false } // backoffDelay returns the per-exec retry delay after `attempt` diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 6bfd3fd2bfb..45e316889e3 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -14,9 +14,8 @@ const ( // time by setting the start request's workflow type. shardedForceReplicationWorkflowName = "force-replication-sharded" - // shardedBatchActivityName is the registered activity name; the - // workflow dispatches by name so the same activity body can be - // referenced across CAN cycles without re-registering. + // shardedBatchActivityName is the registered activity name used by + // name-based dispatch. shardedBatchActivityName = "ReplicateBatch" // releaseShardsSignalName carries mid-flight ReleaseShards signals @@ -28,9 +27,6 @@ const ( // defaultShardedListPageSize is the ListWorkflows page size when // the sharded workflow's params.ListWorkflowsPageSize is unset. - // Named to avoid clashing with the legacy - // defaultListWorkflowsPageSize already declared in - // force_replication_workflow.go. defaultShardedListPageSize = 1000 // defaultBatchSize bounds the total executions in any single @@ -48,8 +44,8 @@ const ( // defaultShardNoProgress is the per-shard cumulative no-progress // backstop. While a shard's pending exec count is non-zero and // no exec on that shard has produced a verified outcome for this - // many seconds (carried across CAN via the resume payload), the - // activity fails non-retryably naming the stuck shard. + // long (carried across CAN via the resume payload), the activity + // fails non-retryably naming the stuck shard. defaultShardNoProgress = 5 * time.Minute // defaultDrainGrace is the wall-budget the activity gets after @@ -65,21 +61,17 @@ const ( defaultIdleShardCost = 30 * time.Second // defaultPerBatchGenerateRPS is the per-batch inject-phase target. - // Sharded dispatches many concurrent batches and each builds its own - // limiter, so this caps the per-batch generate-replication-task rate; - // the workflow does not normalise against a global cap the way the - // existing migration's OverallRps does. Sits near the mid-range of - // what existing force-replication deployments configure per-activity - // (OverallRps 10–100 divided across ConcurrentActivityCount 2–10). + // Sharded dispatches many concurrent batches and each builds its + // own limiter, so this caps the per-batch generate-replication-task + // rate; the workflow does not normalise against a global cap the + // way the existing migration's OverallRps does. defaultPerBatchGenerateRPS = 30.0 // defaultConcurrentBatchCap is the ceiling applied to the derived // default of TargetClusterShardCount/4. Keeps the in-flight batch - // count safely inside per-worker concurrent-activity suggestions - // even on the largest cells (4k+ shards), and absolutely bounds the - // cluster blast radius of a single force-rep run. + // count safely inside per-worker concurrent-activity budgets and + // bounds the cluster blast radius of a single force-rep run. defaultConcurrentBatchCap = 500 - ) // RunEntry is the per-run leaf in the nested batch payload. Carries the @@ -97,10 +89,11 @@ type RunEntry struct { ArchetypeID uint32 } -// MarshalJSON / UnmarshalJSON: custom because Go's default JSON has no -// way to express a heterogeneous tuple, and we want the archetype-omission -// to happen by changing the tuple length rather than emitting an explicit -// zero. +// MarshalJSON serialises RunEntry as a heterogeneous JSON tuple. Custom +// because Go's default JSON can't express a heterogeneous tuple, and +// archetype-omission needs to happen by changing the tuple length +// rather than emitting an explicit zero. UnmarshalJSON below is the +// inverse. func (r RunEntry) MarshalJSON() ([]byte, error) { if r.ArchetypeID == 0 { return fmt.Appendf(nil, `[%q]`, r.RunID), nil @@ -205,8 +198,7 @@ func (p BatchPayload) flatten() []*shardedExecutionInfo { return out } -// merge folds src into p. Used by collectRecoveredBuckets to combine -// per-batch payloads back into a single carry-over map. +// merge folds src into p. func (p BatchPayload) merge(src BatchPayload) { for sh, byBID := range src { if p[sh] == nil { @@ -236,20 +228,16 @@ type ShardedForceReplicationParams struct { IdleShardCost time.Duration // PerBatchGenerateRPS is the inject-phase rate-limiter target inside - // each ReplicateBatch activity. Each batch builds its own - // quotas.RateLimiter at this rate, so N concurrent batches inject at - // N× this rate in aggregate — unlike the existing migration package's - // OverallRps, sharded does not normalise against a workflow-global - // budget. Defaults to defaultPerBatchGenerateRPS. + // each ReplicateBatch activity. See defaultPerBatchGenerateRPS for + // the rationale; defaults to that value. PerBatchGenerateRPS float64 // ConcurrentBatchCount is the absolute ceiling on in-flight // ReplicateBatch activities. Per-shard exclusivity already bounds - // concurrency to TargetClusterShardCount, but at production cell - // sizes (1k–4k shards) that's well past the worker's - // concurrent-activity budget. This cap keeps the workflow inside - // that budget and limits the cluster blast radius of a single - // force-rep run. Defaults to + // concurrency to TargetClusterShardCount, but at large cluster + // sizes that's well past the worker's concurrent-activity budget. + // This cap keeps the workflow inside that budget and limits the + // cluster blast radius of a single force-rep run. Defaults to // min(TargetClusterShardCount/4, defaultConcurrentBatchCap). ConcurrentBatchCount int @@ -314,11 +302,6 @@ type ResumeShard struct { // unverified execs in its result. NoProgressByShard carries the cumulative // pre-resume no-progress duration so the per-shard backstop stays // meaningful across resume cycles. -// -// PerBatchGenerateRPS is the inject-phase rate-limiter target. Mirrors -// generateReplicationTasksRequest.RPS in the existing migration package — -// each batch builds its own quotas.RateLimiter at that rate, so two -// concurrent batches inject at 2× this rate in aggregate. type shardedBatchReq struct { BatchID int64 Namespace string @@ -346,8 +329,7 @@ type shardedBatchReq struct { // // CompletedShards is informational (the dispatch coroutine's defer clears // heldByBatch + shardInFlight regardless), but keeping it in the result -// gives metrics and future bookkeeping a clean handle on "which shards -// this batch finished". +// gives metrics a clean handle on "which shards this batch finished". type replicateBatchResult struct { CompletedShards []int32 InFlight []ResumeShard diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 5a2ef24cf89..f797b7064d4 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -83,7 +83,7 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe // workflows match the namespace's force-rep query. Used once at // workflow start to seed TotalForceReplicateWorkflowCount for the // status query's progress reporting. No filter is applied — sharded -// currently lists the namespace's entire workflow population. +// lists the namespace's entire workflow population. func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedForceReplicationParams) (int64, error) { ao := workflow.ActivityOptions{ StartToCloseTimeout: 2 * time.Minute, @@ -141,10 +141,14 @@ type shardedWorkflowState struct { // claims, not stomp on the new claimant. heldByBatch map[int64]map[int32]bool - // batchCancels carries each batch's cancel func keyed by batch - // ID, so the workflow's drain-for-CAN phase can cancel - // in-flight activities individually. - batchCancels map[int64]workflow.CancelFunc + // activityCtx is a cancellable child of run()'s ctx; every + // dispatched batch activity is derived from it. cancelActivities + // is the matching cancel func — drainForCAN calls it once to + // cancel every in-flight batch at once. Cancelling activityCtx + // leaves the workflow's main ctx alive so the drain loop's + // Await keeps running. + activityCtx workflow.Context + cancelActivities workflow.CancelFunc // batchExecs tracks the input payload of each in-flight batch. // Cleared on any nil-error return (drained execs are folded @@ -237,7 +241,6 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati bucketCounts: map[int32]int{}, shardInFlight: map[int32]bool{}, heldByBatch: map[int64]map[int32]bool{}, - batchCancels: map[int64]workflow.CancelFunc{}, batchExecs: map[int64]BatchPayload{}, metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ metrics.OperationTagName: metrics.MigrationWorkflowScope, @@ -258,6 +261,15 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati } func (s *shardedWorkflowState) run(ctx workflow.Context) error { + // Cancellable child ctx for all dispatched batch activities. + // drainForCAN cancels it once to drain in-flight batches without + // touching the workflow's main ctx (which the drain loop's + // Await still rides on). + actCtx, cancelAll := workflow.WithCancel(ctx) + defer cancelAll() + s.activityCtx = actCtx + s.cancelActivities = cancelAll + // Start the signal handler coroutine first so any signal // arriving during resume dispatch or page-loop drains is // processed promptly. @@ -377,11 +389,9 @@ func (s *shardedWorkflowState) recordVerified(ctx workflow.Context, verified int // defaultConcurrentBatchCount derives the in-flight-batch ceiling // from the target cluster's shard count: a quarter of the shards, -// capped at defaultConcurrentBatchCap. The 1/4 fraction keeps the -// workflow inside Temporal Cloud's concurrent-activity suggestions -// (a 4k-shard cell with cap 500 still has spare worker slots for -// unrelated activities); the absolute cap bounds the cluster blast -// radius regardless of cluster size. Returns at least 1. +// capped at defaultConcurrentBatchCap. The 1/4 fraction leaves worker +// slots free for unrelated activities; the absolute cap bounds the +// cluster blast radius regardless of cluster size. Returns at least 1. func defaultConcurrentBatchCount(shards int32) int { return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) } @@ -405,8 +415,7 @@ func collectRecoveredBuckets(batchExecs map[int64]BatchPayload) BatchPayload { } // addToBucket appends one run to the (shard, BID) bucket and bumps -// the per-shard count. The count is a sidecar so the streaming -// packer's eligibility check stays O(1) per shard. +// the sidecar count. func (s *shardedWorkflowState) addToBucket(shard int32, businessID string, run RunEntry) { if s.buckets[shard] == nil { s.buckets[shard] = map[string][]RunEntry{} @@ -608,44 +617,51 @@ func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { } currentPending := s.pendingDispatches if currentPending == 0 { - // Non-empty buckets but nothing in flight means the - // shard-claim bookkeeping is corrupted: tryPackStreaming - // declined to pack anything yet no batch is running to - // eventually free a shard. Returning silently would - // proceed to CAN (or completion) with execs still in - // buckets that were never dispatched — silent data loss. - // Fail the workflow instead so lastErr propagates out - // through run(). - remaining := 0 - for _, n := range s.bucketCounts { - remaining += n - } - s.lastErr = temporal.NewNonRetryableApplicationError( - fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), - "DrainBucketsStuck", nil) + s.failDrainBucketsStuck() return } - // A "free shard" wake-up only counts when there's also a - // dispatch slot to use it, otherwise the outer loop would - // busy-spin on tryPackStreaming returning false against - // the in-flight cap. - _ = workflow.Await(ctx, func() bool { - if s.lastErr != nil { - return true - } - if s.pendingDispatches < currentPending { + _ = workflow.Await(ctx, s.drainBucketsAwaitPredicate(currentPending)) + } +} + +// failDrainBucketsStuck sets lastErr when buckets are non-empty but no +// batches are in flight — the shard-claim bookkeeping is corrupted, and +// returning silently would proceed to CAN with execs that were never +// dispatched (silent data loss). Failing forces lastErr to propagate +// through run(). +func (s *shardedWorkflowState) failDrainBucketsStuck() { + remaining := 0 + for _, n := range s.bucketCounts { + remaining += n + } + s.lastErr = temporal.NewNonRetryableApplicationError( + fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), + "DrainBucketsStuck", nil) +} + +// drainBucketsAwaitPredicate returns true when the drainBuckets loop +// should wake up: lastErr tripped, a dispatch slot just freed, or a +// new free shard is ready to pack. A "free shard" wake-up only counts +// when there's also a dispatch slot to use it, otherwise the outer +// loop would busy-spin on tryPackStreaming returning false against the +// in-flight cap. +func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { + return func() bool { + if s.lastErr != nil { + return true + } + if s.pendingDispatches < currentPending { + return true + } + if !s.dispatchSlotAvailable() { + return false + } + for sh, n := range s.bucketCounts { + if n > 0 && !s.shardInFlight[sh] { return true } - if !s.dispatchSlotAvailable() { - return false - } - for sh, n := range s.bucketCounts { - if n > 0 && !s.shardInFlight[sh] { - return true - } - } - return false - }) + } + return false } } @@ -672,9 +688,7 @@ func (s *shardedWorkflowState) drainForCAN(ctx workflow.Context) { if s.pendingDispatches == 0 { return } - for _, cancel := range s.batchCancels { - cancel() - } + s.cancelActivities() releaseCh := workflow.GetSignalChannel(ctx, releaseShardsSignalName) _ = workflow.Await(ctx, func() bool { if s.lastErr != nil { @@ -688,9 +702,8 @@ func (s *shardedWorkflowState) drainForCAN(ctx workflow.Context) { // Callers must have already marked every shard appearing as a // top-level key in payload as shardInFlight (the "claim") so the // packer can see them as busy while picking subsequent batches. The -// coroutine receives a per-batch cancellable ctx; the cancel func is -// stored in batchCancels so drainForCAN can cancel individual -// batches without tearing down the whole workflow. +// activity is run on s.activityCtx so drainForCAN can cancel every +// in-flight batch with a single call. func (s *shardedWorkflowState) spawnBatch( ctx workflow.Context, payload BatchPayload, @@ -718,8 +731,6 @@ func (s *shardedWorkflowState) spawnBatch( IdleShardCost: s.params.IdleShardCost, } - batchCtx, cancel := workflow.WithCancel(ctx) - s.batchCancels[batchID] = cancel held := make(map[int32]bool, len(payload)) for sh := range payload { held[sh] = true @@ -739,8 +750,6 @@ func (s *shardedWorkflowState) spawnBatch( delete(s.shardInFlight, sh) } delete(s.heldByBatch, batchID) - delete(s.batchCancels, batchID) - cancel() }() ao := workflow.ActivityOptions{ StartToCloseTimeout: 24 * time.Hour, @@ -757,7 +766,7 @@ func (s *shardedWorkflowState) spawnBatch( // before the dispatch coroutine's defer fires. WaitForCancellation: true, } - actx := workflow.WithActivityOptions(batchCtx, ao) + actx := workflow.WithActivityOptions(s.activityCtx, ao) var result replicateBatchResult err := workflow.ExecuteActivity(actx, shardedBatchActivityName, req).Get(coroCtx, &result) if err == nil { @@ -796,21 +805,7 @@ func (s *shardedWorkflowState) spawnBatch( // caps the per-batch source RPS. Within those bounds the packer's // sole job is to make progress every chance it gets. // -// relax=false (streaming, during ListWorkflows pages): pack fullest -// free shards first. When work is plentiful, batches land at -// BatchSize across few shards (small in-flight shard set per batch, -// nice and predictable); when it's sparse, the same loop ships a -// smaller batch rather than waiting and burning wall-clock on idle. -// -// relax=true (drain, after listing finishes — no more execs coming): -// hot shards (count > MaxExecsPerShard, i.e. those needing more than -// one round trip) first, fullest within hot. Total drain wall-clock -// is bounded by the heaviest shard's round-trip count × activity -// duration, so getting each hot shard's pipeline started ASAP is the -// dominant lever. Once they're all claimed, remaining batch capacity -// fills from smallest cold buckets so light shards finish and free -// their slots quickly — no point waiting on them to grow, nothing -// will arrive. +// See shardIDsByPackPriority for the relax-mode ordering rationale. func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool) bool { if s.lastErr != nil || s.params.BatchSize <= 0 || s.params.MaxExecsPerShard <= 0 { return false @@ -891,12 +886,10 @@ func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { case !aHot && bHot: return 1 case aHot && bHot: - // Both hot: fullest first. if d := s.bucketCounts[b] - s.bucketCounts[a]; d != 0 { return d } default: - // Both cold: smallest first. if d := s.bucketCounts[a] - s.bucketCounts[b]; d != 0 { return d } diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 55680a22ecb..6a0417a01fb 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -2,7 +2,6 @@ package migration import ( "context" - "errors" "fmt" "strconv" "sync" @@ -56,8 +55,8 @@ func bidsForShards(namespaceID string, totalShards int32, perShard int) map[int3 // makeExecs builds a slice of ExecutionInfos engineered to hash across // `shards` distinct shards (`perShard` execs per shard) under the test -// namespace ID + shard count. Shard is left zero — the workflow's page -// loop populates it via common.WorkflowIDToHistoryShard. +// namespace ID + shard count. The workflow's page loop computes the +// destination shard itself via common.WorkflowIDToHistoryShard. func makeExecs(shards int32, perShard int) []*ExecutionInfo { bids := bidsForShards(testNamespaceID, shards, perShard) var execs []*ExecutionInfo @@ -75,9 +74,9 @@ func makeExecs(shards int32, perShard int) []*ExecutionInfo { } // pageThrough returns a function suitable for OnActivity("ListWorkflows") -// that paginates `all` into pages of `pageSize` execs each. The -// workflow populates ex.Shard after this returns, so callers can hand -// over ExecutionInfos with Shard left zero. +// that paginates `all` into pages of `pageSize` execs each. The workflow +// computes each exec's destination shard itself, so callers don't need +// to set anything shard-related on the ExecutionInfos. func pageThrough(all []*ExecutionInfo, pageSize int) func(context.Context, *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { return func(_ context.Context, req *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { start := 0 @@ -107,10 +106,8 @@ func metadataResponseFor(shardCount int32) func(context.Context, MetadataRequest } } -// registerShardedScaffolding registers GetMetadata + CountWorkflow stubs. -// Every sharded workflow test exercises a different scenario (paging, -// resume, drain, ...) but they all need these two before the page loop -// runs, so factoring them out keeps each test focused on its scenario. +// registerShardedScaffolding registers the GetMetadata + CountWorkflow +// stubs every sharded test needs before the page loop runs. func registerShardedScaffolding(env *testsuite.TestWorkflowEnvironment, shardCount int32) { env.RegisterActivityWithOptions(metadataResponseFor(shardCount), activity.RegisterOptions{Name: "GetMetadata"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.CountWorkflowExecutionsRequest) (*countWorkflowResponse, error) { @@ -277,6 +274,7 @@ func TestSharded_ReleaseShards_FreesShardForReuse(t *testing.T) { } case 2: secondOnce.Do(func() { close(secondStarted) }) + default: } return replicateBatchResult{}, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) @@ -319,7 +317,7 @@ func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { err := env.GetWorkflowError() require.Error(t, err) var appErr *temporal.ApplicationError - require.True(t, errors.As(err, &appErr)) + require.ErrorAs(t, err, &appErr) require.Equal(t, "ShardNoProgress", appErr.Type()) } @@ -329,11 +327,11 @@ func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { // carry-over. // // Tests the workflow plumbing only. In production, an activity -// returns InFlight after entering drain mode and grace-expiring; -// the testsuite has a known bug where activity-cancellation result -// handling loses payload fidelity, so we exercise the same code -// path by returning InFlight from a non-cancelled run instead. The -// dispatch coroutine's err == nil branch is what we're testing — +// returns InFlight after entering drain mode and grace-expiring; here +// we exercise the same code path by returning InFlight from a +// non-cancelled run, because the testsuite delivers cancellation as +// a CanceledError without preserving the activity's returned result. +// The dispatch coroutine's err == nil branch is what we're testing — // it doesn't care whether the activity was cancelled or not, only // whether the returned result has InFlight to fold into drainPayload. func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) { @@ -342,16 +340,14 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) env.RegisterWorkflow(ShardedForceReplicationWorkflow) registerShardedScaffolding(env, 10) - // Page 1 returns a multi-shard population large enough to trip - // the streaming packer's mid-cycle dispatch under the pinned - // BatchSize/MaxExecsPerShard the test sets below - // (trigger=MaxExecsPerShard/2=5, minShards=BatchSize/MaxExecsPerShard=10). - // The activity - // flips CAN-suggested from inside its body before returning, so - // by the time it has handed back its InFlight the workflow is - // committed to CAN — but without going through cancel, which - // the testsuite delivers as a CanceledError regardless of any - // result the activity returned. + // Page 1 returns a multi-shard population so the streaming packer + // has something to dispatch under the pinned BatchSize / + // MaxExecsPerShard the test sets below. The activity flips + // CAN-suggested from inside its body before returning, so by the + // time it has handed back its InFlight the workflow is committed + // to CAN — but without going through cancel, which the testsuite + // delivers as a CanceledError regardless of any result the + // activity returned. pageExecs := makeExecs(10, 10) // Drained exec mirrors a real input row so the simulated drain // payload would be a valid response from a real activity. drainedBID @@ -390,7 +386,7 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) require.Error(t, err, "workflow should CAN, not return success") var canErr *workflow.ContinueAsNewError - require.True(t, errors.As(err, &canErr), "error should be ContinueAsNewError") + require.ErrorAs(t, err, &canErr, "error should be ContinueAsNewError") var nextParams ShardedForceReplicationParams require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) @@ -402,13 +398,11 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) require.Equal(t, 42*time.Second, nextParams.ResumeShards[0].NoProgressDuration) } -// TestSharded_CancelBeforeStart_NoLostExecs: this is the reproducer -// for the sim bug — when an activity is dispatched and the workflow -// CANs before the activity body runs (testsuite-scheduler race), the -// activity returns CanceledError with no result. The recovery path -// re-buckets the input execs into RecoveredBuckets so the next cycle -// dispatches them as fresh inject+verify batches. This test pins -// down that behaviour. +// TestSharded_CancelBeforeStart_NoLostExecs pins down recovery when +// an activity is dispatched and the workflow CANs before the activity +// body runs: the activity returns CanceledError with no result, and +// the recovery path re-buckets the input execs into RecoveredBuckets +// so the next cycle dispatches them as fresh inject+verify batches. func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() @@ -457,7 +451,7 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { err := env.GetWorkflowError() require.Error(t, err) var canErr *workflow.ContinueAsNewError - require.True(t, errors.As(err, &canErr), "expected CAN, got %v", err) + require.ErrorAs(t, err, &canErr, "expected CAN, got %v", err) var nextParams ShardedForceReplicationParams require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) From 33cd6b4ba8addd961054df532c26e66e67ae4e1d Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 09:51:36 +0100 Subject: [PATCH 04/35] Add some missing features from the current force replication system. --- service/worker/migration/fx.go | 3 + service/worker/migration/sharded_activity.go | 50 +++++++--- service/worker/migration/sharded_types.go | 18 +++- service/worker/migration/sharded_workflow.go | 95 +++++++++++++++---- .../worker/migration/sharded_workflow_test.go | 56 ++++++++++- 5 files changed, 192 insertions(+), 30 deletions(-) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index d1c8a72014b..b92ca6ca85a 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -139,6 +139,9 @@ func (sc *shardedWorkerComponent) RegisterWorkflow(registry sdkworker.Registry) registry.RegisterWorkflowWithOptions(ShardedForceReplicationWorkflow, workflow.RegisterOptions{ Name: shardedForceReplicationWorkflowName, }) + registry.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{ + Name: forceTaskQueueUserDataReplicationWorkflow, + }) } func (sc *shardedWorkerComponent) DedicatedWorkflowWorkerOptions() *workercommon.DedicatedWorkerOptions { diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index d65b897e41f..5fbd547d8b0 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -15,6 +15,8 @@ import ( "go.temporal.io/sdk/temporal" "go.temporal.io/server/api/adminservice/v1" "go.temporal.io/server/client/admin" + "go.temporal.io/server/common" + "go.temporal.io/server/common/log/tag" "go.temporal.io/server/common/metrics" "go.temporal.io/server/common/namespace" "go.temporal.io/server/common/quotas" @@ -40,6 +42,29 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( remoteAdminClient := a.clientFactory.NewRemoteAdminClientWithTimeout( req.TargetClusterEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout) + var hb replicateBatchHeartbeat + if activity.HasHeartbeatDetails(ctx) { + _ = activity.GetHeartbeatDetails(ctx, &hb) + } + + // ---- Inject phase ---- + if !req.Resume && !hb.InjectDone { + startIdx := hb.NextInjectIdx + if err := a.runInjectPhase(ctx, req, execs, startIdx); err != nil { + return replicateBatchResult{}, err + } + activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) + } + + // inject-only path: skip verify and let the workflow accounting + // release shards on activity return. + if req.DisableVerification { + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: 0, + }, nil + } + // Namespace lookup feeds the verify phase's retention/zombie skip // check (checkSkipWorkflowExecution needs ns.Retention()). Snapshotted // once per activity; we don't track config changes mid-batch. @@ -48,13 +73,6 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) } - // ---- Inject phase ---- - if !req.Resume { - if err := a.runInjectPhase(ctx, req, execs); err != nil { - return replicateBatchResult{}, err - } - } - // ---- Verify phase ---- verified := make([]bool, of) attempts := make([]int, of) @@ -118,7 +136,7 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( } doneCount += passDelta - activity.RecordHeartbeat(ctx, doneCount) + activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) // Clean completion — every exec verified. if doneCount >= of { @@ -167,9 +185,10 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( // RecoveredBuckets re-injection next cycle. Already-injected execs // are re-injected harmlessly — replication dedupes per (namespace, // wf, run). -func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, execs []*shardedExecutionInfo) error { +func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, execs []*shardedExecutionInfo, startIdx int) error { rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) - for _, ex := range execs { + for i := startIdx; i < len(execs); i++ { + ex := execs[i] if ctx.Err() != nil { return temporal.NewCanceledError("inject phase cancelled") } @@ -177,8 +196,17 @@ func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, e if ctx.Err() != nil { return temporal.NewCanceledError("inject phase cancelled") } - return err + if common.IsNotFoundError(err) { + a.Logger.Warn("force-replication-sharded ignore replication task due to NotFoundServiceError", + tag.WorkflowNamespaceID(req.NamespaceID), + tag.WorkflowID(ex.BusinessID), + tag.WorkflowRunID(ex.RunID), + tag.Error(err)) + } else { + return err + } } + activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{NextInjectIdx: i + 1}) } return nil } diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 45e316889e3..e1d0c9d492d 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -216,17 +216,21 @@ func (p BatchPayload) merge(src BatchPayload) { type ShardedForceReplicationParams struct { // ---- Configuration ---- Namespace string + Query string BatchSize int MaxExecsPerShard int ListWorkflowsPageSize int TargetClusterEndpoint string TargetClusterName string TargetClusterShardCount int32 + DisableVerification bool ShardNoProgress time.Duration DrainGrace time.Duration IdleShardCost time.Duration + TaskQueueUserDataReplicationParams TaskQueueUserDataReplicationParams + // PerBatchGenerateRPS is the inject-phase rate-limiter target inside // each ReplicateBatch activity. See defaultPerBatchGenerateRPS for // the rationale; defaults to that value. @@ -270,6 +274,8 @@ type ShardedForceReplicationParams struct { // but never injected, so the new cycle restores them into the // streaming buckets to be dispatched as fresh inject+verify batches. RecoveredBuckets BatchPayload + + TaskQueueUserDataReplicationStatus TaskQueueUserDataReplicationStatus } // ResumeShard carries one shard's worth of unverified execs from a drained @@ -311,8 +317,9 @@ type shardedBatchReq struct { TargetClusterEndpoint string TargetClusterName string - Resume bool - NoProgressByShard map[int32]time.Duration + Resume bool + DisableVerification bool + NoProgressByShard map[int32]time.Duration PerBatchGenerateRPS float64 @@ -342,6 +349,13 @@ type replicateBatchResult struct { VerifiedCount int64 } +type replicateBatchHeartbeat struct { + // NextInjectIdx is the index of the next exec to inject on retry. + NextInjectIdx int + // InjectDone marks the inject phase as complete; retries skip inject. + InjectDone bool +} + // releaseShardsPayload is the body of the mid-flight ReleaseShards signal // an activity sends to its parent workflow when the cumulative idle cost // across its completed-but-not-yet-released shards crosses IdleShardCost. diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index f797b7064d4..25c08d010cc 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -5,6 +5,7 @@ import ( "slices" "time" + enumspb "go.temporal.io/api/enums/v1" "go.temporal.io/api/workflowservice/v1" sdkclient "go.temporal.io/sdk/client" "go.temporal.io/sdk/temporal" @@ -36,21 +37,24 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe // Register the status query under the same name upstream uses // (forceReplicationStatusQueryType = "force-replication-status") // so tooling that polls force-rep progress works across both - // workflow variants. Sharded-irrelevant ForceReplicationStatus - // fields (TaskQueueUserDataReplicationStatus) are left zero — - // sharded only handles the data-replication phase. + // workflow variants. if err := workflow.SetQueryHandler(ctx, forceReplicationStatusQueryType, func() (ForceReplicationStatus, error) { return ForceReplicationStatus{ - ContinuedAsNewCount: params.ContinuedAsNewCount, - TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, - ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, - ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, - PageTokenForRestart: startPageToken, + ContinuedAsNewCount: params.ContinuedAsNewCount, + TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, + ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, + ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, + PageTokenForRestart: startPageToken, + TaskQueueUserDataReplicationStatus: params.TaskQueueUserDataReplicationStatus, }, nil }); err != nil { return err } + if err := validateShardedForceReplicationParams(¶ms); err != nil { + return err + } + state, err := newShardedWorkflowState(ctx, ¶ms) if err != nil { return err @@ -76,14 +80,71 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe params.TotalForceReplicateWorkflowCount = wfCount } - return state.run(ctx) + if !params.TaskQueueUserDataReplicationStatus.Done { + if err := maybeKickoffShardedTaskQueueUserDataReplication(ctx, ¶ms, func(failureReason string) { + params.TaskQueueUserDataReplicationStatus.FailureMessage = failureReason + params.TaskQueueUserDataReplicationStatus.Done = true + }); err != nil { + return err + } + } + + if err := state.run(ctx); err != nil { + return err + } + + // state.run returned nil only on the terminal cycle (no more pages, + // no errors). On CAN cycles it returns the CAN error, so we never + // reach here mid-replication. + if err := workflow.Await(ctx, func() bool { return params.TaskQueueUserDataReplicationStatus.Done }); err != nil { + return err + } + if params.TaskQueueUserDataReplicationStatus.FailureMessage != "" { + return fmt.Errorf("task queue user data replication failed: %v", params.TaskQueueUserDataReplicationStatus.FailureMessage) + } + return nil +} + +func validateShardedForceReplicationParams(params *ShardedForceReplicationParams) error { + if len(params.Namespace) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: Namespace is required", "InvalidArgument", nil) + } + if !params.DisableVerification && len(params.TargetClusterEndpoint) == 0 && len(params.TargetClusterName) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterEndpoint or TargetClusterName is required with verification enabled", "InvalidArgument", nil) + } + return nil +} + +func maybeKickoffShardedTaskQueueUserDataReplication(ctx workflow.Context, params *ShardedForceReplicationParams, onDone func(failureReason string)) error { + workflow.Go(ctx, func(ctx workflow.Context) { + ch := workflow.GetSignalChannel(ctx, taskQueueUserDataReplicationDoneSignalType) + var errStr string + _ = ch.Receive(ctx, &errStr) + onDone(errStr) + }) + + if params.ContinuedAsNewCount > 0 { + return nil + } + + options := workflow.ChildWorkflowOptions{ + WorkflowID: fmt.Sprintf("%s-task-queue-user-data-replicator", workflow.GetInfo(ctx).WorkflowExecution.ID), + ParentClosePolicy: enumspb.PARENT_CLOSE_POLICY_ABANDON, + } + childCtx := workflow.WithChildOptions(ctx, options) + input := TaskQueueUserDataReplicationParamsWithNamespace{ + TaskQueueUserDataReplicationParams: params.TaskQueueUserDataReplicationParams, + Namespace: params.Namespace, + } + child := workflow.ExecuteChildWorkflow(childCtx, ForceTaskQueueUserDataReplicationWorkflow, input) + var childExecution workflow.Execution + return child.GetChildWorkflowExecution().Get(ctx, &childExecution) } // shardedCountWorkflowsForReplication asks the frontend how many // workflows match the namespace's force-rep query. Used once at // workflow start to seed TotalForceReplicateWorkflowCount for the -// status query's progress reporting. No filter is applied — sharded -// lists the namespace's entire workflow population. +// status query's progress reporting. func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedForceReplicationParams) (int64, error) { ao := workflow.ActivityOptions{ StartToCloseTimeout: 2 * time.Minute, @@ -96,6 +157,7 @@ func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedFo a.CountWorkflow, &workflowservice.CountWorkflowExecutionsRequest{ Namespace: params.Namespace, + Query: params.Query, }).Get(ctx, &output); err != nil { return 0, err } @@ -299,6 +361,7 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { } listReq := &workflowservice.ListWorkflowExecutionsRequest{ Namespace: s.params.Namespace, + Query: s.params.Query, PageSize: int32(s.params.ListWorkflowsPageSize), NextPageToken: s.params.NextPageToken, } @@ -724,6 +787,7 @@ func (s *shardedWorkflowState) spawnBatch( TargetClusterEndpoint: s.params.TargetClusterEndpoint, TargetClusterName: s.params.TargetClusterName, Resume: resume, + DisableVerification: s.params.DisableVerification, NoProgressByShard: noProgressByShard, PerBatchGenerateRPS: s.params.PerBatchGenerateRPS, ShardNoProgress: s.params.ShardNoProgress, @@ -754,12 +818,11 @@ func (s *shardedWorkflowState) spawnBatch( ao := workflow.ActivityOptions{ StartToCloseTimeout: 24 * time.Hour, HeartbeatTimeout: time.Minute, - // MaxAttempts=1: per-exec backoff inside the activity - // is the retry path; an activity-level retry would - // re-run inject and reset all backoff state, wasting - // the apply work this attempt already drove. + // 3 attempts: per-exec backoff still owns the per-exec + // retry, but a transient activity failure can recover + // via heartbeat-resume without losing inject progress. RetryPolicy: &temporal.RetryPolicy{ - MaximumAttempts: 1, + MaximumAttempts: 3, }, // WaitForCancellation: the cancelled activity needs to // run its drain logic and return its drain result diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 6a0417a01fb..028feb5a513 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -107,12 +107,18 @@ func metadataResponseFor(shardCount int32) func(context.Context, MetadataRequest } // registerShardedScaffolding registers the GetMetadata + CountWorkflow -// stubs every sharded test needs before the page loop runs. +// stubs every sharded test needs before the page loop runs, plus the +// task-queue-user-data child workflow + its activity so the parent's +// terminal Await on Done resolves. func registerShardedScaffolding(env *testsuite.TestWorkflowEnvironment, shardCount int32) { env.RegisterActivityWithOptions(metadataResponseFor(shardCount), activity.RegisterOptions{Name: "GetMetadata"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.CountWorkflowExecutionsRequest) (*countWorkflowResponse, error) { return &countWorkflowResponse{WorkflowCount: 0}, nil }, activity.RegisterOptions{Name: "CountWorkflow"}) + env.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{Name: forceTaskQueueUserDataReplicationWorkflow}) + env.RegisterActivityWithOptions(func(_ context.Context, _ TaskQueueUserDataReplicationParamsWithNamespace) error { + return nil + }, activity.RegisterOptions{Name: "SeedReplicationQueueWithUserDataEntries"}) } // ---- Tests ---- @@ -147,6 +153,7 @@ func TestSharded_HappyPath_SingleCycle(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 4, }) @@ -196,6 +203,7 @@ func TestSharded_ResumeShards_Packed(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 8, ResumeShards: resumeShards, }) @@ -281,6 +289,7 @@ func TestSharded_ReleaseShards_FreesShardForReuse(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 2, ConcurrentBatchCount: 2, }) @@ -310,6 +319,7 @@ func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 2, }) @@ -376,6 +386,7 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 10, BatchSize: 100, MaxExecsPerShard: 10, @@ -444,6 +455,7 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", + TargetClusterName: "remote_cluster", TargetClusterShardCount: 4, }) @@ -468,6 +480,48 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { require.Empty(t, nextParams.ResumeShards, "no ResumeShards — activity never ran, never injected, so no resume work") } +// TestSharded_DisableVerification_NoVerifiedCount: with verification +// disabled the workflow runs inject-only batches, completes +// successfully, and the status query reports ReplicatedWorkflowCount=0 +// because no verification ran. +func TestSharded_DisableVerification_NoVerifiedCount(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + + execs := makeExecs(4, 5) + registerShardedScaffolding(env, 4) + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + var sawDisable atomic.Bool + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + if req.DisableVerification { + sawDisable.Store(true) + } + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: 0, + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + TargetClusterShardCount: 4, + DisableVerification: true, + }) + + require.True(t, env.IsWorkflowCompleted(), "workflow should complete") + require.NoError(t, env.GetWorkflowError(), "workflow should succeed") + require.True(t, sawDisable.Load(), "activity req should carry DisableVerification=true") + + envValue, err := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, err) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + require.Equal(t, int64(0), status.ReplicatedWorkflowCount, "no verification ran so verified count must stay 0") +} + // ---- internal helpers ---- // makeExecsForShard produces `count` runs for the named shard's From 57e730e74baa80e89170595dc2bb896f5abbddd8 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 10:25:29 +0100 Subject: [PATCH 05/35] Improve test coverage. --- service/worker/migration/activities_test.go | 1 + .../force_replication_workflow_test.go | 3 + .../worker/migration/sharded_activity_test.go | 368 ++++++++++++++++++ .../worker/migration/sharded_workflow_test.go | 156 +++++++- tests/xdc/failover_test.go | 250 ++++++++++++ tests/xdc/user_data_replication_test.go | 143 +++++++ 6 files changed, 918 insertions(+), 3 deletions(-) create mode 100644 service/worker/migration/sharded_activity_test.go diff --git a/service/worker/migration/activities_test.go b/service/worker/migration/activities_test.go index 4ca1d7492c9..6bc0e625024 100644 --- a/service/worker/migration/activities_test.go +++ b/service/worker/migration/activities_test.go @@ -111,6 +111,7 @@ func (s *activitiesSuite) SetupTest() { s.mockNamespaceReplicationQueue = persistence.NewMockNamespaceReplicationQueue(s.controller) s.mockNamespaceRegistry = namespace.NewMockRegistry(s.controller) s.mockClientBean = client.NewMockBean(s.controller) + s.mockClientFactory = client.NewMockFactory(s.controller) s.mockFrontendClient = workflowservicemock.NewMockWorkflowServiceClient(s.controller) s.mockAdminClient = adminservicemock.NewMockAdminServiceClient(s.controller) diff --git a/service/worker/migration/force_replication_workflow_test.go b/service/worker/migration/force_replication_workflow_test.go index 8cb9e9e9188..bea39ac2a7e 100644 --- a/service/worker/migration/force_replication_workflow_test.go +++ b/service/worker/migration/force_replication_workflow_test.go @@ -705,6 +705,7 @@ type heartbeatRecordingInterceptor struct { seedRecordedHeartbeats []seedReplicationQueueWithUserDataEntriesHeartbeatDetails replicationRecordedHeartbeats []replicationTasksHeartbeatDetails generateReplicationRecordedHeartbeats []int + replicateBatchRecordedHeartbeats []replicateBatchHeartbeat T *testing.T } @@ -725,6 +726,8 @@ func (i *heartbeatRecordingInterceptor) RecordHeartbeat(ctx context.Context, det i.replicationRecordedHeartbeats = append(i.replicationRecordedHeartbeats, d) } else if d, ok := details[0].(int); ok { i.generateReplicationRecordedHeartbeats = append(i.generateReplicationRecordedHeartbeats, d) + } else if d, ok := details[0].(replicateBatchHeartbeat); ok { + i.replicateBatchRecordedHeartbeats = append(i.replicateBatchRecordedHeartbeats, d) } else { assert.Fail(i.T, "invalid heartbeat details") } diff --git a/service/worker/migration/sharded_activity_test.go b/service/worker/migration/sharded_activity_test.go new file mode 100644 index 00000000000..4f2e4f9ca98 --- /dev/null +++ b/service/worker/migration/sharded_activity_test.go @@ -0,0 +1,368 @@ +package migration + +import ( + "time" + + commonpb "go.temporal.io/api/common/v1" + enumspb "go.temporal.io/api/enums/v1" + "go.temporal.io/api/serviceerror" + "go.temporal.io/sdk/temporal" + "go.temporal.io/server/api/adminservice/v1" + enumsspb "go.temporal.io/server/api/enums/v1" + "go.temporal.io/server/api/historyservice/v1" + persistencespb "go.temporal.io/server/api/persistence/v1" + "go.temporal.io/server/chasm" + "go.temporal.io/server/client/admin" + "go.temporal.io/server/common/namespace" + "go.temporal.io/server/common/testing/protomock" + "go.uber.org/mock/gomock" + "google.golang.org/protobuf/types/known/durationpb" + "google.golang.org/protobuf/types/known/timestamppb" +) + +// Sharded-activity tests reuse activitiesSuite's SetupTest so they get the +// same mock graph (HistoryClient, AdminClient, ChasmRegistry, etc.) as the +// legacy force-replication activity tests. ReplicateBatch resolves the +// remote admin client via clientFactory.NewRemoteAdminClientWithTimeout +// rather than clientBean.GetRemoteAdminClient, so each test arms a +// NewRemoteAdminClientWithTimeout expectation that hands back the suite's +// mockRemoteAdminClient. + +const remoteEndpoint = "remote.example:7233" + +// expectNewRemoteAdminClient arms the clientFactory mock so ReplicateBatch +// gets the suite's mockRemoteAdminClient back. AnyTimes() because the +// activity may build the client once per attempt and retries can fire. +func (s *activitiesSuite) expectNewRemoteAdminClient() { + s.mockClientFactory.EXPECT(). + NewRemoteAdminClientWithTimeout(remoteEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout). + Return(s.mockRemoteAdminClient).AnyTimes() +} + +// payloadFor wraps a single ExecutionInfo into a BatchPayload on the named +// shard. Tests that need multiple execs across shards build the BatchPayload +// inline. +func payloadFor(shard int32, ex *ExecutionInfo) BatchPayload { + return BatchPayload{ + shard: {ex.BusinessID: {{RunID: ex.RunID, ArchetypeID: ex.ArchetypeID}}}, + } +} + +// newShardedReq builds a shardedBatchReq with sensible defaults for unit +// tests. ShardNoProgress is large so the stuck-shard backstop doesn't trip +// from real wall-clock latency; IdleShardCost is large so maybeSignalRelease +// (which needs the sdkClientFactory, nil here) doesn't fire. +func newShardedReq(execs BatchPayload) *shardedBatchReq { + return &shardedBatchReq{ + BatchID: 1, + Namespace: mockedNamespace, + NamespaceID: mockedNamespaceID, + Executions: execs, + TargetClusterEndpoint: remoteEndpoint, + TargetClusterName: remoteCluster, + PerBatchGenerateRPS: defaultPerBatchGenerateRPS, + ShardNoProgress: time.Hour, + DrainGrace: time.Second, + IdleShardCost: time.Hour, + } +} + +// expectRemoteNotFound primes the remote admin client to return NotFound +// for the given exec — the trigger for the verify-skip code path that +// consults source DMS to decide between zombie/retention skip vs. real +// pending state. +func (s *activitiesSuite) expectRemoteNotFound(ex *ExecutionInfo) { + s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), protomock.Eq(&adminservice.DescribeMutableStateRequest{ + Namespace: mockedNamespace, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: ex.BusinessID, + RunId: ex.RunID, + }, + Archetype: chasm.WorkflowArchetype, + ArchetypeId: ex.ArchetypeID, + SkipForceReload: true, + })).Return(nil, serviceerror.NewNotFound("")).Times(1) +} + +func (s *activitiesSuite) expectSourceDMS(ex *ExecutionInfo, resp *historyservice.DescribeMutableStateResponse, err error) { + s.mockHistoryClient.EXPECT().DescribeMutableState(gomock.Any(), protomock.Eq(&historyservice.DescribeMutableStateRequest{ + NamespaceId: mockedNamespaceID, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: ex.BusinessID, + RunId: ex.RunID, + }, + ArchetypeId: ex.ArchetypeID, + SkipForceReload: true, + })).Return(resp, err).Times(1) +} + +// TestReplicateBatch_Success exercises the full inject+verify happy path +// for one exec: the inject phase calls GenerateLastHistoryReplicationTasks +// against HistoryClient, then the verify phase's DescribeMutableState on +// the remote admin client returns OK so workflowVerifier marks it +// verified. Mirrors the legacy TestVerifyReplicationTasks_Success and +// TestGenerateReplicationTasks_Success. +func (s *activitiesSuite) TestReplicateBatch_Success() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(&testNamespace, nil).Times(1) + + // Inject phase calls HistoryClient (generateMigrationTaskViaFrontend=false). + s.mockHistoryClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), protomock.Eq(&historyservice.GenerateLastHistoryReplicationTasksRequest{ + NamespaceId: mockedNamespaceID, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: execution1.BusinessID, + RunId: execution1.RunID, + }, + ArchetypeId: execution1.ArchetypeID, + TargetClusters: []string{remoteCluster}, + })).Return(&historyservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1) + + // Verify phase: remote DMS returns OK; workflowVerifierProvider + // returns verified=true unconditionally so the exec verifies on the + // first pass. + s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), protomock.Eq(&adminservice.DescribeMutableStateRequest{ + Namespace: mockedNamespace, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: execution1.BusinessID, + RunId: execution1.RunID, + }, + Archetype: chasm.WorkflowArchetype, + ArchetypeId: execution1.ArchetypeID, + SkipForceReload: true, + })).Return(&adminservice.DescribeMutableStateResponse{}, nil).Times(1) + + req := newShardedReq(payloadFor(0, execution1)) + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(1), out.VerifiedCount) + s.Equal([]int32{0}, out.CompletedShards) + s.Empty(out.InFlight) +} + +// TestReplicateBatch_SkipZombie exercises the retention/zombie skip path: +// remote DMS returns NotFound, source DMS returns a zombie state, and +// checkSkipWorkflowExecution marks the exec as verified-via-skip so the +// shard's verify accounting completes. Mirrors the existing +// TestVerifyReplicationTasks_SkipWorkflowExecution. +func (s *activitiesSuite) TestReplicateBatch_SkipZombie() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(&testNamespace, nil).Times(1) + + // Resume=true so we skip inject and only exercise the verify-skip path. + req := newShardedReq(payloadFor(0, execution1)) + req.Resume = true + + s.expectRemoteNotFound(execution1) + s.expectSourceDMS(execution1, zombieState, nil) + + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(1), out.VerifiedCount) + s.Empty(out.InFlight) +} + +// TestReplicateBatch_SkipRetention exercises the close-time/retention skip +// path: remote DMS returns NotFound, source DMS returns a completed +// workflow whose CloseTime+Retention is in the past, so +// checkSkipWorkflowExecution marks it skipped (counted as verified). +// Mirrors the existing Test_verifyReplicationTasksSkipRetention. +func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + + retention := time.Hour + closeTime := time.Now().Add(-2 * retention) // deleteTime is in the past + + // Build a real namespace.Namespace with a retention setting so + // checkSkipWorkflowExecution's `ns.Retention()` returns non-zero. + factory := namespace.NewDefaultReplicationResolverFactory() + detail := &persistencespb.NamespaceDetail{ + Info: &persistencespb.NamespaceInfo{}, + Config: &persistencespb.NamespaceConfig{ + Retention: durationpb.New(retention), + }, + ReplicationConfig: &persistencespb.NamespaceReplicationConfig{}, + } + ns, nsErr := namespace.FromPersistentState(detail, factory(detail)) + s.NoError(nsErr) + // Override the suite-default GetNamespaceByID for this test so the + // activity sees a namespace with retention configured. + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(ns, nil).Times(1) + + s.expectRemoteNotFound(execution1) + s.expectSourceDMS(execution1, &historyservice.DescribeMutableStateResponse{ + DatabaseMutableState: &persistencespb.WorkflowMutableState{ + ExecutionState: &persistencespb.WorkflowExecutionState{ + State: enumsspb.WORKFLOW_EXECUTION_STATE_COMPLETED, + }, + ExecutionInfo: &persistencespb.WorkflowExecutionInfo{ + CloseTime: timestamppb.New(closeTime), + }, + }, + }, nil) + + req := newShardedReq(payloadFor(0, execution1)) + req.Resume = true + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(1), out.VerifiedCount) +} + +// TestReplicateBatch_ShardNoProgress: the per-shard cumulative no-progress +// backstop fires non-retryably when a shard has gone longer than +// req.ShardNoProgress without a verified outcome. Resume=true with a +// pre-seeded NoProgressByShard pushes the shard right at the threshold +// before the first verify pass, so the first failed attempt trips the +// check immediately and we don't have to spin on wall-clock. Mirrors +// the existing TestVerifyReplicationTasks_FailedNotFound. +func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(&testNamespace, nil).Times(1) + + // Remote returns BUSY_WORKFLOW so verify counts the exec as + // pending without consulting source — that keeps the shard's + // lastProgress at its seeded (already-stale) value. + s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), gomock.Any()). + Return(nil, &serviceerror.ResourceExhausted{ + Cause: enumspb.RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW, + }).AnyTimes() + + req := newShardedReq(payloadFor(0, execution1)) + req.Resume = true // skip inject + req.ShardNoProgress = 10 * time.Millisecond // trip almost immediately + req.NoProgressByShard = map[int32]time.Duration{ // seed past threshold + 0: time.Second, + } + + _, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.Error(err) + var appErr *temporal.ApplicationError + s.ErrorAs(err, &appErr) + s.Equal("ShardNoProgress", appErr.Type()) + s.True(appErr.NonRetryable(), "ShardNoProgress should be non-retryable") +} + +// TestReplicateBatch_Resume_SkipsInject: Resume=true should bypass the +// inject phase entirely — no GenerateLastHistoryReplicationTasks call. +// Mirrors the inject-side guarantee that the legacy +// TestVerifyReplicationTasks_AlreadyVerified asserts for verify +// (resume-via-heartbeat skips already-done work). +func (s *activitiesSuite) TestReplicateBatch_Resume_SkipsInject() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(&testNamespace, nil).Times(1) + + // No GenerateLastHistoryReplicationTasks expectation — gomock with + // strict expectations would fail if inject called it. Verify path: + // remote DMS OK so the exec verifies on first pass. + s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), gomock.Any()). + Return(&adminservice.DescribeMutableStateResponse{}, nil).Times(1) + + req := newShardedReq(payloadFor(0, execution1)) + req.Resume = true + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(1), out.VerifiedCount) +} + +// TestReplicateBatch_DisableVerification: with verification disabled the +// activity runs inject, then returns immediately with VerifiedCount=0 and +// every batch shard listed as completed — no DMS calls. Mirrors the +// workflow-level TestSharded_DisableVerification_NoVerifiedCount but +// from the activity side. +func (s *activitiesSuite) TestReplicateBatch_DisableVerification() { + env, _ := s.initEnv() + // No NewRemoteAdminClientWithTimeout expectation — the activity + // builds the client unconditionally even in inject-only mode, so we + // still need the factory to hand back something. Return without + // expecting any DMS calls on it. + s.expectNewRemoteAdminClient() + + s.mockHistoryClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), gomock.Any()). + Return(&historyservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1) + + req := newShardedReq(payloadFor(0, execution1)) + req.DisableVerification = true + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(0), out.VerifiedCount) + s.Equal([]int32{0}, out.CompletedShards) +} + +// TestReplicateBatch_HeartbeatResumesInject: a recorded NextInjectIdx +// heartbeat from a prior attempt causes the inject phase to skip +// already-injected execs. Pre-seeds heartbeat NextInjectIdx=1 across a +// two-exec batch, then asserts only the second exec's +// GenerateLastHistoryReplicationTasks is invoked. Mirrors the legacy +// TestGenerateReplicationTasks_Success_ViaFrontend's heartbeat-resume +// assertion. +func (s *activitiesSuite) TestReplicateBatch_HeartbeatResumesInject() { + env, _ := s.initEnv() + s.expectNewRemoteAdminClient() + s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). + Return(&testNamespace, nil).Times(1) + + // Pre-seed the heartbeat so inject resumes at index 1. + env.SetHeartbeatDetails(replicateBatchHeartbeat{NextInjectIdx: 1, InjectDone: false}) + + // Two execs on shard 0. flatten() orders by BID alphabetically, so + // index 0 = execution1 ("workflow1"), index 1 = execution2 ("workflow2"). + payload := BatchPayload{ + 0: { + execution1.BusinessID: {{RunID: execution1.RunID, ArchetypeID: execution1.ArchetypeID}}, + execution2.BusinessID: {{RunID: execution2.RunID, ArchetypeID: execution2.ArchetypeID}}, + }, + } + + // Only execution2 (index 1) should be injected. Strict Times(1) + // ensures execution1 is NOT injected — gomock would fail an + // unexpected execution1 call. + s.mockHistoryClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), protomock.Eq(&historyservice.GenerateLastHistoryReplicationTasksRequest{ + NamespaceId: mockedNamespaceID, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: execution2.BusinessID, + RunId: execution2.RunID, + }, + ArchetypeId: execution2.ArchetypeID, + TargetClusters: []string{remoteCluster}, + })).Return(&historyservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1) + + // Both execs verify in one pass. + for _, ex := range []*ExecutionInfo{execution1, execution2} { + s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), protomock.Eq(&adminservice.DescribeMutableStateRequest{ + Namespace: mockedNamespace, + Execution: &commonpb.WorkflowExecution{ + WorkflowId: ex.BusinessID, + RunId: ex.RunID, + }, + Archetype: chasm.WorkflowArchetype, + ArchetypeId: ex.ArchetypeID, + SkipForceReload: true, + })).Return(&adminservice.DescribeMutableStateResponse{}, nil).Times(1) + } + + req := newShardedReq(payload) + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(2), out.VerifiedCount) +} diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 028feb5a513..7b2d9f0c604 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -111,14 +111,25 @@ func metadataResponseFor(shardCount int32) func(context.Context, MetadataRequest // task-queue-user-data child workflow + its activity so the parent's // terminal Await on Done resolves. func registerShardedScaffolding(env *testsuite.TestWorkflowEnvironment, shardCount int32) { + registerShardedScaffoldingWithSeed(env, shardCount, func(_ context.Context, _ TaskQueueUserDataReplicationParamsWithNamespace) error { + return nil + }) +} + +// registerShardedScaffoldingWithSeed is like registerShardedScaffolding +// but lets the caller supply the SeedReplicationQueueWithUserDataEntries +// activity body — needed for tests that exercise the seed-failure path. +func registerShardedScaffoldingWithSeed( + env *testsuite.TestWorkflowEnvironment, + shardCount int32, + seed func(context.Context, TaskQueueUserDataReplicationParamsWithNamespace) error, +) { env.RegisterActivityWithOptions(metadataResponseFor(shardCount), activity.RegisterOptions{Name: "GetMetadata"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.CountWorkflowExecutionsRequest) (*countWorkflowResponse, error) { return &countWorkflowResponse{WorkflowCount: 0}, nil }, activity.RegisterOptions{Name: "CountWorkflow"}) env.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{Name: forceTaskQueueUserDataReplicationWorkflow}) - env.RegisterActivityWithOptions(func(_ context.Context, _ TaskQueueUserDataReplicationParamsWithNamespace) error { - return nil - }, activity.RegisterOptions{Name: "SeedReplicationQueueWithUserDataEntries"}) + env.RegisterActivityWithOptions(seed, activity.RegisterOptions{Name: "SeedReplicationQueueWithUserDataEntries"}) } // ---- Tests ---- @@ -522,6 +533,145 @@ func TestSharded_DisableVerification_NoVerifiedCount(t *testing.T) { require.Equal(t, int64(0), status.ReplicatedWorkflowCount, "no verification ran so verified count must stay 0") } +// TestSharded_InvalidInput: validateShardedForceReplicationParams +// rejects an empty Namespace and a missing TargetClusterEndpoint / +// TargetClusterName when verification is enabled. Mirrors the existing +// force-replication TestInvalidInput. +func TestSharded_InvalidInput(t *testing.T) { + for _, tc := range []struct { + name string + params ShardedForceReplicationParams + }{ + { + name: "empty namespace", + params: ShardedForceReplicationParams{}, + }, + { + name: "missing target with verification on", + params: ShardedForceReplicationParams{ + Namespace: "test-ns", + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, tc.params) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + require.Contains(t, err.Error(), "InvalidArgument") + }) + } +} + +// TestSharded_ListWorkflowsError: a hard failure from ListWorkflows +// propagates out as the workflow error. Mirrors the existing +// force-replication TestListWorkflowsError. +func TestSharded_ListWorkflowsError(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + return nil, temporal.NewNonRetryableApplicationError("mock listWorkflows error", "ListFailed", nil) + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + // ReplicateBatch should never be invoked because listing fails up + // front. Register a fail-loud stub so we notice if the workflow + // ever dispatches a batch. + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + t.Fatalf("ReplicateBatch must not be called when listing fails") + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + TargetClusterShardCount: 2, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + require.Contains(t, err.Error(), "mock listWorkflows error") +} + +// TestSharded_ReplicateBatchRetryableError: when ReplicateBatch returns +// a retryable error, the workflow exhausts its 3-attempt retry policy +// and surfaces the error as lastErr. Mirrors the existing +// TestGenerateReplicationTaskRetryableError. +func TestSharded_ReplicateBatchRetryableError(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + env.RegisterActivityWithOptions(pageThrough(makeExecs(2, 5), 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + var attempts atomic.Int32 + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + attempts.Add(1) + return replicateBatchResult{}, temporal.NewApplicationError("transient backend error", "Transient") + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + TargetClusterShardCount: 2, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + require.Contains(t, err.Error(), "transient backend error") + // MaximumAttempts: 3 in spawnBatch's activity options — assert at + // least 2 retries actually happened so a future change that drops + // the retry policy fails this test. + require.GreaterOrEqual(t, attempts.Load(), int32(2), + "expected ReplicateBatch to be retried at least twice before failing") +} + +// TestSharded_TaskQueueReplicationFailure: when the +// SeedReplicationQueueWithUserDataEntries activity returns a +// non-retryable error, the child workflow signals the failure back +// and the parent fails with the seed error message; the status +// query reports the failure reason. Mirrors the existing +// TestTaskQueueReplicationFailure. +func TestSharded_TaskQueueReplicationFailure(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffoldingWithSeed(env, 2, + func(_ context.Context, _ TaskQueueUserDataReplicationParamsWithNamespace) error { + return temporal.NewNonRetryableApplicationError("namespace is required", "InvalidArgument", nil) + }) + env.RegisterActivityWithOptions(pageThrough(nil, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + TargetClusterShardCount: 2, + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + require.Contains(t, err.Error(), "namespace is required") + + envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + require.True(t, status.TaskQueueUserDataReplicationStatus.Done) + require.Contains(t, status.TaskQueueUserDataReplicationStatus.FailureMessage, "namespace is required") +} + // ---- internal helpers ---- // makeExecsForShard produces `count` runs for the named shard's diff --git a/tests/xdc/failover_test.go b/tests/xdc/failover_test.go index 7795083a4e1..28086bfa5fe 100644 --- a/tests/xdc/failover_test.go +++ b/tests/xdc/failover_test.go @@ -2707,6 +2707,256 @@ func (s *FunctionalClustersTestSuite) TestForceMigration_ResetWorkflow() { verifyHistory(workflowID, resp.GetRunId()) } +// TestForceMigration_ClosedWorkflow_Sharded is the sharded-variant +// duplicate of TestForceMigration_ClosedWorkflow. Kept intentionally +// close to the original — same workflow IDs prefixed with "sharded-", +// same assertions — so a future swap from legacy to sharded +// force-replication can drop the legacy version and keep behaviour +// coverage intact. +func (s *FunctionalClustersTestSuite) TestForceMigration_ClosedWorkflow_Sharded() { + testCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + namespace := s.createNamespaceInCluster0(true) + + taskqueue := "functional-local-force-replication-sharded-task-queue" + client0, worker0 := s.newClientAndWorker(s.clusters[0].Host().FrontendGRPCAddress(), namespace, taskqueue, "worker0") + + testWorkflowFn := func(ctx workflow.Context) error { + return nil + } + + worker0.RegisterWorkflow(testWorkflowFn) + s.NoError(worker0.Start()) + defer worker0.Stop() + + // Start wf1 + workflowID := "sharded-force-replication-test-wf-1" + run1, err := client0.ExecuteWorkflow(testCtx, sdkclient.StartWorkflowOptions{ + ID: workflowID, + TaskQueue: taskqueue, + WorkflowRunTimeout: time.Second * 30, + }, testWorkflowFn) + + s.NoError(err) + s.NotEmpty(run1.GetRunID()) + s.logger.Info("start wf1", tag.WorkflowRunID(run1.GetRunID())) + // wait until wf1 complete + err = run1.Get(testCtx, nil) + s.NoError(err) + + // Update ns to have 2 clusters + s.updateNamespaceClusters(namespace, 0, s.clusters) + + // Wait for wf1 to be indexed before force-replication. + s.waitForVisibilityCount(testCtx, namespace, 1) + + // Start force-replicate wf — sharded variant lives on its own + // dedicated task queue (MigrationShardedActivityTQ) and is + // registered under workflow name "force-replication-sharded". + sysClient, err := sdkclient.Dial(sdkclient.Options{ + HostPort: s.clusters[0].Host().FrontendGRPCAddress(), + Namespace: "temporal-system", + }) + s.NoError(err) + forceReplicationWorkflowID := "sharded-force-replication-wf" + sysWfRun, err := sysClient.ExecuteWorkflow(testCtx, sdkclient.StartWorkflowOptions{ + ID: forceReplicationWorkflowID, + TaskQueue: primitives.MigrationShardedActivityTQ, + WorkflowRunTimeout: time.Second * 30, + }, "force-replication-sharded", migration.ShardedForceReplicationParams{ + Namespace: namespace, + TargetClusterName: s.clusters[1].ClusterName(), + }) + s.NoError(err) + err = sysWfRun.Get(testCtx, nil) + s.NoError(err) + + // Verify all wf in ns is now available in cluster2 + client1, worker1 := s.newClientAndWorker(s.clusters[1].Host().FrontendGRPCAddress(), namespace, taskqueue, "worker1") + verify := func(wfID string, expectedRunID string) { + s.Eventually(func() bool { + desc1, err := client1.DescribeWorkflowExecution(testCtx, wfID, "") + if err != nil { + return false + } + return desc1.WorkflowExecutionInfo.Execution.RunId == expectedRunID && + desc1.WorkflowExecutionInfo.Status == enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED + }, 15*time.Second, 200*time.Millisecond, "workflow %s should be replicated to cluster2", wfID) + } + verify(workflowID, run1.GetRunID()) + + s.failover(namespace, 0, s.clusters[1].ClusterName(), 2) + + worker1.RegisterWorkflow(testWorkflowFn) + s.NoError(worker1.Start()) + defer worker1.Stop() + + // Test reset workflow in cluster1 + resetResp, err := client1.ResetWorkflowExecution(testCtx, &workflowservice.ResetWorkflowExecutionRequest{ + Namespace: namespace, + WorkflowExecution: &commonpb.WorkflowExecution{ + WorkflowId: workflowID, + RunId: run1.GetRunID(), + }, + Reason: "force-replication-sharded-test", + WorkflowTaskFinishEventId: 3, + RequestId: uuid.NewString(), + }) + s.NoError(err) + + resetRun := client1.GetWorkflow(testCtx, workflowID, resetResp.GetRunId()) + err = resetRun.Get(testCtx, nil) + s.NoError(err) + + s.Eventually(func() bool { + descResp, err := client1.DescribeWorkflowExecution(testCtx, workflowID, resetResp.GetRunId()) + if err != nil { + return false + } + return descResp.GetWorkflowExecutionInfo().Status == enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED + }, 15*time.Second, 200*time.Millisecond, "reset workflow should be visible on cluster2") +} + +// TestForceMigration_ResetWorkflow_Sharded is the sharded-variant +// duplicate of TestForceMigration_ResetWorkflow. Same intent as the +// legacy version: replicate a (reset → completed) workflow pair across +// clusters and confirm both runs are visible on the target. Asserts +// the activity-level verification count by walking the workflow's +// history for "ReplicateBatch" activity completions (the sharded +// activity name, replacing legacy "VerifyReplicationTasks"). +func (s *FunctionalClustersTestSuite) TestForceMigration_ResetWorkflow_Sharded() { + testCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + namespace := s.createNamespaceInCluster0(true) + + taskqueue := "functional-force-replication-sharded-reset-task-queue" + client0, worker0 := s.newClientAndWorker(s.clusters[0].Host().FrontendGRPCAddress(), namespace, taskqueue, "worker0") + + testWorkflowFn := func(ctx workflow.Context) error { + return nil + } + + worker0.RegisterWorkflow(testWorkflowFn) + s.NoError(worker0.Start()) + defer worker0.Stop() + + // Start wf1 + workflowID := "sharded-force-replication-test-reset-wf-1" + run1, err := client0.ExecuteWorkflow(testCtx, sdkclient.StartWorkflowOptions{ + ID: workflowID, + TaskQueue: taskqueue, + WorkflowRunTimeout: time.Second * 30, + }, testWorkflowFn) + + s.NoError(err) + s.NotEmpty(run1.GetRunID()) + s.logger.Info("start wf1", tag.WorkflowRunID(run1.GetRunID())) + // wait until wf1 complete + err = run1.Get(testCtx, nil) + s.NoError(err) + + resp, err := client0.ResetWorkflowExecution(testCtx, &workflowservice.ResetWorkflowExecutionRequest{ + Namespace: namespace, + WorkflowExecution: &commonpb.WorkflowExecution{ + WorkflowId: workflowID, + RunId: run1.GetRunID(), + }, + Reason: "test", + WorkflowTaskFinishEventId: 3, + RequestId: uuid.NewString(), + }) + s.NoError(err) + resetRun := client0.GetWorkflow(testCtx, workflowID, resp.GetRunId()) + err = resetRun.Get(testCtx, nil) + s.NoError(err) + + // Update ns to have 2 clusters + s.updateNamespaceClusters(namespace, 0, s.clusters) + + // Wait for both workflow runs (original + reset) to be indexed before force-replication. + s.waitForVisibilityCount(testCtx, namespace, 2) + + // Start force-replicate wf + sysClient, err := sdkclient.Dial(sdkclient.Options{ + HostPort: s.clusters[0].Host().FrontendGRPCAddress(), + Namespace: "temporal-system", + }) + s.NoError(err) + forceReplicationWorkflowID := "sharded-force-replication-wf" + sysWfRun, err := sysClient.ExecuteWorkflow(testCtx, sdkclient.StartWorkflowOptions{ + ID: forceReplicationWorkflowID, + TaskQueue: primitives.MigrationShardedActivityTQ, + WorkflowRunTimeout: time.Second * 30, + }, "force-replication-sharded", migration.ShardedForceReplicationParams{ + Namespace: namespace, + TargetClusterName: s.clusters[1].ClusterName(), + }) + s.NoError(err) + err = sysWfRun.Get(testCtx, nil) + s.NoError(err) + + // Verify the force-replication workflow actually ran ReplicateBatch + // activities (the sharded activity name; legacy is + // VerifyReplicationTasks) and that VerifiedCount sums to the + // expected number of workflow runs. + var totalVerifiedCount int64 + scheduledActivityTypes := make(map[int64]string) // scheduledEventId -> activity type name + histIter := sysClient.GetWorkflowHistory(testCtx, forceReplicationWorkflowID, sysWfRun.GetRunID(), + false, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT) + for histIter.HasNext() { + event, err := histIter.Next() + s.NoError(err) + switch event.GetEventType() { + case enumspb.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED: + attrs := event.GetActivityTaskScheduledEventAttributes() + scheduledActivityTypes[event.GetEventId()] = attrs.GetActivityType().GetName() + case enumspb.EVENT_TYPE_ACTIVITY_TASK_COMPLETED: + attrs := event.GetActivityTaskCompletedEventAttributes() + activityType := scheduledActivityTypes[attrs.GetScheduledEventId()] + if activityType != "ReplicateBatch" { + continue + } + result := attrs.GetResult() + if result != nil && len(result.GetPayloads()) > 0 { + // Mirrors replicateBatchResult.VerifiedCount on the + // activity-side struct. Anonymous shape avoids + // importing the activity package's internal type. + var resp struct { + VerifiedCount int64 + } + s.NoError(payloads.Decode(result, &resp)) + totalVerifiedCount += resp.VerifiedCount + } + default: + } + } + // Expect exactly 2 verified workflow runs: original run + reset run + s.Equal(int64(2), totalVerifiedCount, + "sharded force-replication should have verified exactly 2 workflow runs (original + reset run)") + + s.waitForClusterSynced() + + // Verify all wf in ns is now available in cluster2 + client1, _ := s.newClientAndWorker(s.clusters[1].Host().FrontendGRPCAddress(), namespace, taskqueue, "worker1") + verifyHistory := func(wfID string, runID string) { + iter1 := client0.GetWorkflowHistory(testCtx, wfID, runID, false, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT) + iter2 := client1.GetWorkflowHistory(testCtx, wfID, runID, false, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT) + for iter1.HasNext() && iter2.HasNext() { + event1, err := iter1.Next() + s.NoError(err) + event2, err := iter2.Next() + s.NoError(err) + s.Equal(event1, event2) + } + s.False(iter1.HasNext()) + s.False(iter2.HasNext()) + } + verifyHistory(workflowID, run1.GetRunID()) + verifyHistory(workflowID, resp.GetRunId()) +} + func (s *FunctionalClustersTestSuite) TestBlockNamespaceDeleteInPassiveCluster() { namespace := s.createGlobalNamespace() diff --git a/tests/xdc/user_data_replication_test.go b/tests/xdc/user_data_replication_test.go index b49f6be6698..52195666d0d 100644 --- a/tests/xdc/user_data_replication_test.go +++ b/tests/xdc/user_data_replication_test.go @@ -499,6 +499,149 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand( } } +// TestUserDataEntriesAreReplicatedOnDemand_Sharded is the +// sharded-variant duplicate of TestUserDataEntriesAreReplicatedOnDemand. +// Same intent — confirm that running the force-replication workflow +// pushes every task-queue user data entry onto the namespace +// replication queue — but exercised through the sharded workflow +// (registered name "force-replication-sharded" on +// MigrationShardedActivityTQ). DisableVerification:true mirrors the +// legacy test's EnableVerification:false default and lets the workflow +// run without a TargetClusterName since this test only exercises the +// task-queue-user-data side of force-replication. +func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_Sharded() { + ctx := testcore.NewContext() + activeFrontendClient := s.clusters[0].FrontendClient() + adminClient := s.clusters[0].AdminClient() + numTaskQueues := 10 + + replicationResponse, err := adminClient.GetNamespaceReplicationMessages(ctx, &adminservice.GetNamespaceReplicationMessagesRequest{ + ClusterName: "follower", + LastRetrievedMessageId: -1, + LastProcessedMessageId: -1, + }) + s.NoError(err) + lastMessageId := replicationResponse.GetMessages().GetLastRetrievedMessageId() + + namespace := s.createNamespaceInCluster0(true) + description, err := activeFrontendClient.DescribeNamespace(testcore.NewContext(), &workflowservice.DescribeNamespaceRequest{Namespace: namespace}) + s.NoError(err) + + expectedReplicatedTaskQueues := make(map[string]struct{}, numTaskQueues) + for i := range numTaskQueues { + taskQueue := fmt.Sprintf("v1q%v", i) + res, err := activeFrontendClient.UpdateWorkerBuildIdCompatibility(ctx, &workflowservice.UpdateWorkerBuildIdCompatibilityRequest{ + Namespace: namespace, + TaskQueue: taskQueue, + Operation: &workflowservice.UpdateWorkerBuildIdCompatibilityRequest_AddNewBuildIdInNewDefaultSet{ + AddNewBuildIdInNewDefaultSet: "v0.1", + }, + }) + s.NoError(err) + s.NotNil(res) + expectedReplicatedTaskQueues[taskQueue] = struct{}{} + + taskQueue2 := fmt.Sprintf("v2q%v", i) + rules, err := activeFrontendClient.GetWorkerVersioningRules(ctx, &workflowservice.GetWorkerVersioningRulesRequest{ + Namespace: namespace, + TaskQueue: taskQueue2, + }) + s.NoError(err) + s.NotNil(rules) + + rulesRes, err := activeFrontendClient.UpdateWorkerVersioningRules(ctx, &workflowservice.UpdateWorkerVersioningRulesRequest{ + Namespace: namespace, + TaskQueue: taskQueue2, + ConflictToken: rules.ConflictToken, + Operation: &workflowservice.UpdateWorkerVersioningRulesRequest_InsertAssignmentRule{ + InsertAssignmentRule: &workflowservice.UpdateWorkerVersioningRulesRequest_InsertBuildIdAssignmentRule{ + Rule: &taskqueuepb.BuildIdAssignmentRule{ + TargetBuildId: "asdf", + }, + }, + }, + }) + s.NoError(err) + s.NotNil(rulesRes) + expectedReplicatedTaskQueues[taskQueue2] = struct{}{} + } + + // update namespace to cross clusters + s.updateNamespaceClusters(namespace, 0, s.clusters) + + // we should see one new namespace task in the replication queue + replicationResponse, err = adminClient.GetNamespaceReplicationMessages(ctx, &adminservice.GetNamespaceReplicationMessagesRequest{ + ClusterName: "follower", + LastRetrievedMessageId: lastMessageId, + LastProcessedMessageId: -1, + }) + s.NoError(err) + lastMessageId = replicationResponse.GetMessages().GetLastRetrievedMessageId() + s.Len(replicationResponse.GetMessages().ReplicationTasks, 1) + task := replicationResponse.GetMessages().ReplicationTasks[0] + s.Equal(namespace, task.GetNamespaceTaskAttributes().GetInfo().GetName()) + + // start sharded force-replicate wf + sysClient, err := sdkclient.Dial(sdkclient.Options{ + HostPort: s.clusters[0].Host().FrontendGRPCAddress(), + Namespace: primitives.SystemLocalNamespace, + }) + s.NoError(err) + run, err := sysClient.ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + ID: "sharded-force-replication-wf", + TaskQueue: primitives.MigrationShardedActivityTQ, + WorkflowRunTimeout: time.Second * 30, + }, "force-replication-sharded", migration.ShardedForceReplicationParams{ + Namespace: namespace, + DisableVerification: true, // mirrors legacy test's EnableVerification:false (no target needed) + }) + s.NoError(err) + err = run.Get(ctx, nil) + s.NoError(err) + + replicationResponse, err = adminClient.GetNamespaceReplicationMessages(ctx, &adminservice.GetNamespaceReplicationMessagesRequest{ + ClusterName: "follower", + LastRetrievedMessageId: lastMessageId, + LastProcessedMessageId: -1, + }) + s.NoError(err) + + // we should see a user data task for all task queues + seenTaskQueues := make(map[string]struct{}, numTaskQueues) + for _, task := range replicationResponse.GetMessages().ReplicationTasks { + if attrs := task.GetTaskQueueUserDataAttributes(); attrs.GetNamespaceId() == description.GetNamespaceInfo().Id { + seenTaskQueues[attrs.GetTaskQueueName()] = struct{}{} + } + } + s.Equal(expectedReplicatedTaskQueues, seenTaskQueues) + + // failover and check on the other side + s.failover(namespace, 0, s.clusters[1].ClusterName(), 2) + + activeFrontendClient = s.clusters[1].FrontendClient() + for i := range numTaskQueues { + taskQueue := fmt.Sprintf("v1q%v", i) + + get, err := activeFrontendClient.GetWorkerBuildIdCompatibility(ctx, &workflowservice.GetWorkerBuildIdCompatibilityRequest{ + Namespace: namespace, + TaskQueue: taskQueue, + }) + s.NoError(err) + s.NotNil(get) + + s.NotEmpty(get.MajorVersionSets) + + taskQueue2 := fmt.Sprintf("v2q%v", i) + rules, err := activeFrontendClient.GetWorkerVersioningRules(ctx, &workflowservice.GetWorkerVersioningRulesRequest{ + Namespace: namespace, + TaskQueue: taskQueue2, + }) + s.NoError(err) + s.NotNil(rules) + s.NotEmpty(rules.AssignmentRules) + } +} + func (s *UserDataReplicationTestSuite) TestUserDataTombstonesAreReplicated() { s.T().SkipNow() // flaky test ctx := testcore.NewContext() From 6216ce65b90f2543e73fb4455ce8df49d60e14c0 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 10:52:38 +0100 Subject: [PATCH 06/35] Lint. --- service/worker/migration/sharded_activity.go | 108 ++++++++++++------ .../worker/migration/sharded_activity_test.go | 4 +- .../worker/migration/sharded_workflow_test.go | 2 +- tests/xdc/failover_test.go | 4 +- tests/xdc/user_data_replication_test.go | 8 +- 5 files changed, 83 insertions(+), 43 deletions(-) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index 5fbd547d8b0..e95f07e862c 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -73,7 +73,22 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) } - // ---- Verify phase ---- + return a.runVerifyPhase(ctx, req, execs, of, remoteAdminClient, ns) +} + +// runVerifyPhase is the verify-phase loop body of ReplicateBatch. It +// owns per-exec bookkeeping, the drain-transition handoff, and the +// per-iteration completion / stuck-shard / signal-release decisions +// — extracted from ReplicateBatch to keep its cognitive complexity +// under the linter cap. +func (a *activities) runVerifyPhase( + ctx context.Context, + req *shardedBatchReq, + execs []*shardedExecutionInfo, + of int, + remoteAdminClient adminservice.AdminServiceClient, + ns *namespace.Namespace, +) (replicateBatchResult, error) { verified := make([]bool, of) attempts := make([]int, of) nextRetryAt := make([]time.Time, of) @@ -138,33 +153,11 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) - // Clean completion — every exec verified. - if doneCount >= of { - return replicateBatchResult{ - CompletedShards: shards.allCompleted(), - VerifiedCount: int64(doneCount), - }, nil - } - - // Per-shard cumulative no-progress backstop. - if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, of); sErr != nil { - return replicateBatchResult{}, sErr - } - - if draining { - // Drain-mode exit checks. No signals here — the return value - // carries everything the workflow needs (completed shards + - // unverified execs grouped by shard with their cumulative - // no-progress duration). - if a.shouldExitDrain(req, shards, drainStartAt) { - return replicateBatchResult{ - CompletedShards: shards.allCompleted(), - InFlight: a.buildInFlight(execs, verified, shards, time.Now()), - VerifiedCount: int64(doneCount), - }, nil - } - } else if err := a.maybeSignalRelease(ctx, req, shards); err != nil { + if done, result, err := a.evaluateVerifyIteration( + ctx, req, execs, verified, shards, doneCount, of, draining, drainStartAt); err != nil { return replicateBatchResult{}, err + } else if done { + return result, nil } // If the inner loop aborted because callCtx died, skip the sleep @@ -178,6 +171,54 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( } } +// evaluateVerifyIteration runs the post-pass checks (clean completion, +// stuck-shard backstop, drain-exit decision, mid-flight signal release) +// for a single verify-loop iteration. Returns done=true with the result +// when the loop should exit; otherwise (false, _, nil) means continue. +func (a *activities) evaluateVerifyIteration( + ctx context.Context, + req *shardedBatchReq, + execs []*shardedExecutionInfo, + verified []bool, + shards shardVerifyTracker, + doneCount, of int, + draining bool, + drainStartAt time.Time, +) (bool, replicateBatchResult, error) { + // Clean completion — every exec verified. + if doneCount >= of { + return true, replicateBatchResult{ + CompletedShards: shards.allCompleted(), + VerifiedCount: int64(doneCount), + }, nil + } + + // Per-shard cumulative no-progress backstop. + if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, of); sErr != nil { + return false, replicateBatchResult{}, sErr + } + + if draining { + // Drain-mode exit checks. No signals here — the return value + // carries everything the workflow needs (completed shards + + // unverified execs grouped by shard with their cumulative + // no-progress duration). + if a.shouldExitDrain(req, shards, drainStartAt) { + return true, replicateBatchResult{ + CompletedShards: shards.allCompleted(), + InFlight: a.buildInFlight(execs, verified, shards, time.Now()), + VerifiedCount: int64(doneCount), + }, nil + } + return false, replicateBatchResult{}, nil + } + + if err := a.maybeSignalRelease(ctx, req, shards); err != nil { + return false, replicateBatchResult{}, err + } + return false, replicateBatchResult{}, nil +} + // runInjectPhase walks execs in flattened order, generating one // replication task per exec under a per-batch RPS limiter. Cancellation // mid-loop returns a CanceledError so spawnBatch's IsCanceledError @@ -196,15 +237,14 @@ func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, e if ctx.Err() != nil { return temporal.NewCanceledError("inject phase cancelled") } - if common.IsNotFoundError(err) { - a.Logger.Warn("force-replication-sharded ignore replication task due to NotFoundServiceError", - tag.WorkflowNamespaceID(req.NamespaceID), - tag.WorkflowID(ex.BusinessID), - tag.WorkflowRunID(ex.RunID), - tag.Error(err)) - } else { + if !common.IsNotFoundError(err) { return err } + a.Logger.Warn("force-replication-sharded ignore replication task due to NotFoundServiceError", + tag.WorkflowNamespaceID(req.NamespaceID), + tag.WorkflowID(ex.BusinessID), + tag.WorkflowRunID(ex.RunID), + tag.Error(err)) } activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{NextInjectIdx: i + 1}) } diff --git a/service/worker/migration/sharded_activity_test.go b/service/worker/migration/sharded_activity_test.go index 4f2e4f9ca98..af3037dfd94 100644 --- a/service/worker/migration/sharded_activity_test.go +++ b/service/worker/migration/sharded_activity_test.go @@ -241,8 +241,8 @@ func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { }).AnyTimes() req := newShardedReq(payloadFor(0, execution1)) - req.Resume = true // skip inject - req.ShardNoProgress = 10 * time.Millisecond // trip almost immediately + req.Resume = true // skip inject + req.ShardNoProgress = 10 * time.Millisecond // trip almost immediately req.NoProgressByShard = map[int32]time.Duration{ // seed past threshold 0: time.Second, } diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 7b2d9f0c604..8971b41c692 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -584,7 +584,7 @@ func TestSharded_ListWorkflowsError(t *testing.T) { // front. Register a fail-loud stub so we notice if the workflow // ever dispatches a batch. env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - t.Fatalf("ReplicateBatch must not be called when listing fails") + t.Fatal("ReplicateBatch must not be called when listing fails") return replicateBatchResult{}, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) diff --git a/tests/xdc/failover_test.go b/tests/xdc/failover_test.go index 28086bfa5fe..08e944d6322 100644 --- a/tests/xdc/failover_test.go +++ b/tests/xdc/failover_test.go @@ -2775,7 +2775,7 @@ func (s *FunctionalClustersTestSuite) TestForceMigration_ClosedWorkflow_Sharded( // Verify all wf in ns is now available in cluster2 client1, worker1 := s.newClientAndWorker(s.clusters[1].Host().FrontendGRPCAddress(), namespace, taskqueue, "worker1") verify := func(wfID string, expectedRunID string) { - s.Eventually(func() bool { + await.RequireTruef(s.T(), func() bool { desc1, err := client1.DescribeWorkflowExecution(testCtx, wfID, "") if err != nil { return false @@ -2809,7 +2809,7 @@ func (s *FunctionalClustersTestSuite) TestForceMigration_ClosedWorkflow_Sharded( err = resetRun.Get(testCtx, nil) s.NoError(err) - s.Eventually(func() bool { + await.RequireTruef(s.T(), func() bool { descResp, err := client1.DescribeWorkflowExecution(testCtx, workflowID, resetResp.GetRunId()) if err != nil { return false diff --git a/tests/xdc/user_data_replication_test.go b/tests/xdc/user_data_replication_test.go index 52195666d0d..fcf0dd26d29 100644 --- a/tests/xdc/user_data_replication_test.go +++ b/tests/xdc/user_data_replication_test.go @@ -521,7 +521,7 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_ LastProcessedMessageId: -1, }) s.NoError(err) - lastMessageId := replicationResponse.GetMessages().GetLastRetrievedMessageId() + lastMessageID := replicationResponse.GetMessages().GetLastRetrievedMessageId() namespace := s.createNamespaceInCluster0(true) description, err := activeFrontendClient.DescribeNamespace(testcore.NewContext(), &workflowservice.DescribeNamespaceRequest{Namespace: namespace}) @@ -572,11 +572,11 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_ // we should see one new namespace task in the replication queue replicationResponse, err = adminClient.GetNamespaceReplicationMessages(ctx, &adminservice.GetNamespaceReplicationMessagesRequest{ ClusterName: "follower", - LastRetrievedMessageId: lastMessageId, + LastRetrievedMessageId: lastMessageID, LastProcessedMessageId: -1, }) s.NoError(err) - lastMessageId = replicationResponse.GetMessages().GetLastRetrievedMessageId() + lastMessageID = replicationResponse.GetMessages().GetLastRetrievedMessageId() s.Len(replicationResponse.GetMessages().ReplicationTasks, 1) task := replicationResponse.GetMessages().ReplicationTasks[0] s.Equal(namespace, task.GetNamespaceTaskAttributes().GetInfo().GetName()) @@ -601,7 +601,7 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_ replicationResponse, err = adminClient.GetNamespaceReplicationMessages(ctx, &adminservice.GetNamespaceReplicationMessagesRequest{ ClusterName: "follower", - LastRetrievedMessageId: lastMessageId, + LastRetrievedMessageId: lastMessageID, LastProcessedMessageId: -1, }) s.NoError(err) From a4219cc25b08c6c2505ad1a4896a97f999294196 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 12:04:07 +0100 Subject: [PATCH 07/35] Remove unused TargetClusterEndpoint param. Fetch target shard count via Describe call rather than passing as a param or defaulting to source side's shard count. --- service/worker/migration/activities.go | 36 ++++++++-- .../migration/force_replication_workflow.go | 27 ++++--- .../force_replication_workflow_test.go | 13 ++-- service/worker/migration/sharded_activity.go | 14 ++-- .../worker/migration/sharded_activity_test.go | 44 +++--------- service/worker/migration/sharded_types.go | 35 +++++----- service/worker/migration/sharded_workflow.go | 61 +++++++++------- .../worker/migration/sharded_workflow_test.go | 70 ++++++++----------- tests/xdc/user_data_replication_test.go | 11 +-- 9 files changed, 155 insertions(+), 156 deletions(-) diff --git a/service/worker/migration/activities.go b/service/worker/migration/activities.go index 0818670daaa..e609bf65d32 100644 --- a/service/worker/migration/activities.go +++ b/service/worker/migration/activities.go @@ -72,12 +72,11 @@ type ( } verifyReplicationTasksRequest struct { - Namespace string - NamespaceID string - TargetClusterEndpoint string - TargetClusterName string - VerifyInterval time.Duration `validate:"gte=0"` - Executions []*ExecutionInfo + Namespace string + NamespaceID string + TargetClusterName string + VerifyInterval time.Duration `validate:"gte=0"` + Executions []*ExecutionInfo } verifyReplicationTasksResponse struct { @@ -93,6 +92,14 @@ type ( NamespaceID string } + DescribeTargetClusterRequest struct { + TargetClusterName string + } + + DescribeTargetClusterResponse struct { + ShardCount int32 + } + ReplicationStatus struct { MaxReplicationTaskIds map[int32]int64 } @@ -194,6 +201,23 @@ func (a *activities) GetMetadata(_ context.Context, request MetadataRequest) (*M }, nil } +// DescribeTargetCluster fetches the remote cluster's history shard count via +// its admin DescribeCluster RPC. The remote must be registered with the +// source cluster's cluster metadata (the cluster name doubles as the +// adminClient cache key) — which is already a prerequisite for force +// replication, since the source generates replication tasks against it. +func (a *activities) DescribeTargetCluster(ctx context.Context, req DescribeTargetClusterRequest) (*DescribeTargetClusterResponse, error) { + remoteAdminClient, err := a.clientBean.GetRemoteAdminClient(req.TargetClusterName) + if err != nil { + return nil, err + } + resp, err := remoteAdminClient.DescribeCluster(ctx, &adminservice.DescribeClusterRequest{}) + if err != nil { + return nil, err + } + return &DescribeTargetClusterResponse{ShardCount: resp.GetHistoryShardCount()}, nil +} + // GetMaxReplicationTaskIDs returns max replication task id per shard func (a *activities) GetMaxReplicationTaskIDs(ctx context.Context) (*ReplicationStatus, error) { ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo) diff --git a/service/worker/migration/force_replication_workflow.go b/service/worker/migration/force_replication_workflow.go index 70a46222692..be2c330b45b 100644 --- a/service/worker/migration/force_replication_workflow.go +++ b/service/worker/migration/force_replication_workflow.go @@ -38,7 +38,6 @@ type ( // Used for verifying workflow executions were replicated successfully on target cluster. EnableVerification bool - TargetClusterEndpoint string TargetClusterName string VerifyIntervalInSeconds int `validate:"gte=0"` @@ -342,8 +341,8 @@ func validateAndSetForceReplicationParams(ctx workflow.Context, params *ForceRep return temporal.NewNonRetryableApplicationError("InvalidArgument: Namespace is required", "InvalidArgument", nil) } - if params.EnableVerification && len(params.TargetClusterEndpoint) == 0 && len(params.TargetClusterName) == 0 { - return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterEndpoint or TargetClusterName is required with verification enabled", "InvalidArgument", nil) + if params.EnableVerification && len(params.TargetClusterName) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterName is required with verification enabled", "InvalidArgument", nil) } if params.ConcurrentActivityCount <= 0 { @@ -512,12 +511,11 @@ func enqueueReplicationTasks(ctx workflow.Context, executionsCh workflow.Channel actx, a.VerifyReplicationTasks, &verifyReplicationTasksRequest{ - TargetClusterEndpoint: params.TargetClusterEndpoint, - TargetClusterName: params.TargetClusterName, - Namespace: params.Namespace, - NamespaceID: namespaceID, - Executions: migrationExecutions, - VerifyInterval: time.Duration(params.VerifyIntervalInSeconds) * time.Second, + TargetClusterName: params.TargetClusterName, + Namespace: params.Namespace, + NamespaceID: namespaceID, + Executions: migrationExecutions, + VerifyInterval: time.Duration(params.VerifyIntervalInSeconds) * time.Second, }) pendingVerifyTasks++ @@ -615,12 +613,11 @@ func enqueueReplicationTasksLocal( lactx, a.VerifyReplicationTasks, &verifyReplicationTasksRequest{ - TargetClusterEndpoint: params.TargetClusterEndpoint, - TargetClusterName: params.TargetClusterName, - Namespace: params.Namespace, - NamespaceID: namespaceID, - Executions: executions, - VerifyInterval: time.Duration(params.VerifyIntervalInSeconds) * time.Second, + TargetClusterName: params.TargetClusterName, + Namespace: params.Namespace, + NamespaceID: namespaceID, + Executions: executions, + VerifyInterval: time.Duration(params.VerifyIntervalInSeconds) * time.Second, }) pendingVerifyTasks++ diff --git a/service/worker/migration/force_replication_workflow_test.go b/service/worker/migration/force_replication_workflow_test.go index bea39ac2a7e..479724c2da8 100644 --- a/service/worker/migration/force_replication_workflow_test.go +++ b/service/worker/migration/force_replication_workflow_test.go @@ -108,7 +108,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestForceReplicationWorkflow() { ListWorkflowsPageSize: 1, PageCountPerExecution: 4, EnableVerification: true, - TargetClusterEndpoint: "test-target", + TargetClusterName: "test-target", }) s.True(env.IsWorkflowCompleted()) @@ -167,8 +167,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestContinueAsNew() { PageCountPerExecution: testMaxPageCountPerExecution, NextPageToken: []byte("fake-page-token-2"), EnableVerification: true, - TargetClusterEndpoint: "test-target", - TargetClusterName: "", + TargetClusterName: "test-target", VerifyIntervalInSeconds: defaultVerifyIntervalInSeconds, LastCloseTime: closeTime, LastStartTime: startTime, @@ -194,7 +193,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestContinueAsNew() { ListWorkflowsPageSize: 1, PageCountPerExecution: testMaxPageCountPerExecution, EnableVerification: true, - TargetClusterEndpoint: "test-target", + TargetClusterName: "test-target", NextPageToken: []byte("fake-initial-page-token"), }, expectContinueAsNew, @@ -295,7 +294,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestInvalidInput() { // Empty namespace }, { - // Empty TargetClusterEndpoint + // Empty TargetClusterName Namespace: uuid.NewString(), EnableVerification: true, }, @@ -438,7 +437,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestGenerateReplicationTaskNonRetrya ListWorkflowsPageSize: 1, PageCountPerExecution: 4, EnableVerification: true, - TargetClusterEndpoint: "test-target", + TargetClusterName: "test-target", }) s.True(env.IsWorkflowCompleted()) @@ -495,7 +494,7 @@ func (s *ForceReplicationWorkflowTestSuite) TestVerifyReplicationTaskNonRetryabl ListWorkflowsPageSize: 1, PageCountPerExecution: 4, EnableVerification: true, - TargetClusterEndpoint: "test-target", + TargetClusterName: "test-target", }) s.True(env.IsWorkflowCompleted()) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index e95f07e862c..955c40e3af2 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -14,7 +14,6 @@ import ( "go.temporal.io/sdk/activity" "go.temporal.io/sdk/temporal" "go.temporal.io/server/api/adminservice/v1" - "go.temporal.io/server/client/admin" "go.temporal.io/server/common" "go.temporal.io/server/common/log/tag" "go.temporal.io/server/common/metrics" @@ -39,8 +38,10 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, nil } - remoteAdminClient := a.clientFactory.NewRemoteAdminClientWithTimeout( - req.TargetClusterEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout) + remoteAdminClient, err := a.clientBean.GetRemoteAdminClient(req.TargetClusterName) + if err != nil { + return replicateBatchResult{}, fmt.Errorf("get remote admin client for %s: %w", req.TargetClusterName, err) + } var hb replicateBatchHeartbeat if activity.HasHeartbeatDetails(ctx) { @@ -456,10 +457,9 @@ func (a *activities) attemptVerifyExec( } vreq := &verifyReplicationTasksRequest{ - Namespace: req.Namespace, - NamespaceID: req.NamespaceID, - TargetClusterEndpoint: req.TargetClusterEndpoint, - TargetClusterName: req.TargetClusterName, + Namespace: req.Namespace, + NamespaceID: req.NamespaceID, + TargetClusterName: req.TargetClusterName, } describeStart := time.Now() diff --git a/service/worker/migration/sharded_activity_test.go b/service/worker/migration/sharded_activity_test.go index af3037dfd94..05110158755 100644 --- a/service/worker/migration/sharded_activity_test.go +++ b/service/worker/migration/sharded_activity_test.go @@ -12,7 +12,6 @@ import ( "go.temporal.io/server/api/historyservice/v1" persistencespb "go.temporal.io/server/api/persistence/v1" "go.temporal.io/server/chasm" - "go.temporal.io/server/client/admin" "go.temporal.io/server/common/namespace" "go.temporal.io/server/common/testing/protomock" "go.uber.org/mock/gomock" @@ -23,21 +22,8 @@ import ( // Sharded-activity tests reuse activitiesSuite's SetupTest so they get the // same mock graph (HistoryClient, AdminClient, ChasmRegistry, etc.) as the // legacy force-replication activity tests. ReplicateBatch resolves the -// remote admin client via clientFactory.NewRemoteAdminClientWithTimeout -// rather than clientBean.GetRemoteAdminClient, so each test arms a -// NewRemoteAdminClientWithTimeout expectation that hands back the suite's -// mockRemoteAdminClient. - -const remoteEndpoint = "remote.example:7233" - -// expectNewRemoteAdminClient arms the clientFactory mock so ReplicateBatch -// gets the suite's mockRemoteAdminClient back. AnyTimes() because the -// activity may build the client once per attempt and retries can fire. -func (s *activitiesSuite) expectNewRemoteAdminClient() { - s.mockClientFactory.EXPECT(). - NewRemoteAdminClientWithTimeout(remoteEndpoint, admin.DefaultTimeout, admin.DefaultLargeTimeout). - Return(s.mockRemoteAdminClient).AnyTimes() -} +// remote admin client via clientBean.GetRemoteAdminClient, which the suite +// already arms with mockRemoteAdminClient — no per-test setup needed. // payloadFor wraps a single ExecutionInfo into a BatchPayload on the named // shard. Tests that need multiple execs across shards build the BatchPayload @@ -54,16 +40,15 @@ func payloadFor(shard int32, ex *ExecutionInfo) BatchPayload { // (which needs the sdkClientFactory, nil here) doesn't fire. func newShardedReq(execs BatchPayload) *shardedBatchReq { return &shardedBatchReq{ - BatchID: 1, - Namespace: mockedNamespace, - NamespaceID: mockedNamespaceID, - Executions: execs, - TargetClusterEndpoint: remoteEndpoint, - TargetClusterName: remoteCluster, - PerBatchGenerateRPS: defaultPerBatchGenerateRPS, - ShardNoProgress: time.Hour, - DrainGrace: time.Second, - IdleShardCost: time.Hour, + BatchID: 1, + Namespace: mockedNamespace, + NamespaceID: mockedNamespaceID, + Executions: execs, + TargetClusterName: remoteCluster, + PerBatchGenerateRPS: defaultPerBatchGenerateRPS, + ShardNoProgress: time.Hour, + DrainGrace: time.Second, + IdleShardCost: time.Hour, } } @@ -104,7 +89,6 @@ func (s *activitiesSuite) expectSourceDMS(ex *ExecutionInfo, resp *historyservic // TestGenerateReplicationTasks_Success. func (s *activitiesSuite) TestReplicateBatch_Success() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) @@ -150,7 +134,6 @@ func (s *activitiesSuite) TestReplicateBatch_Success() { // TestVerifyReplicationTasks_SkipWorkflowExecution. func (s *activitiesSuite) TestReplicateBatch_SkipZombie() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) @@ -176,7 +159,6 @@ func (s *activitiesSuite) TestReplicateBatch_SkipZombie() { // Mirrors the existing Test_verifyReplicationTasksSkipRetention. func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() retention := time.Hour closeTime := time.Now().Add(-2 * retention) // deleteTime is in the past @@ -228,7 +210,6 @@ func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { // the existing TestVerifyReplicationTasks_FailedNotFound. func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) @@ -262,7 +243,6 @@ func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { // (resume-via-heartbeat skips already-done work). func (s *activitiesSuite) TestReplicateBatch_Resume_SkipsInject() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) @@ -292,7 +272,6 @@ func (s *activitiesSuite) TestReplicateBatch_DisableVerification() { // builds the client unconditionally even in inject-only mode, so we // still need the factory to hand back something. Return without // expecting any DMS calls on it. - s.expectNewRemoteAdminClient() s.mockHistoryClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), gomock.Any()). Return(&historyservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1) @@ -316,7 +295,6 @@ func (s *activitiesSuite) TestReplicateBatch_DisableVerification() { // assertion. func (s *activitiesSuite) TestReplicateBatch_HeartbeatResumesInject() { env, _ := s.initEnv() - s.expectNewRemoteAdminClient() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index e1d0c9d492d..32460089e3f 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -68,9 +68,9 @@ const ( defaultPerBatchGenerateRPS = 30.0 // defaultConcurrentBatchCap is the ceiling applied to the derived - // default of TargetClusterShardCount/4. Keeps the in-flight batch - // count safely inside per-worker concurrent-activity budgets and - // bounds the cluster blast radius of a single force-rep run. + // default of targetShardCount/4. Keeps the in-flight batch count + // safely inside per-worker concurrent-activity budgets and bounds + // the cluster blast radius of a single force-rep run. defaultConcurrentBatchCap = 500 ) @@ -215,15 +215,13 @@ func (p BatchPayload) merge(src BatchPayload) { // bottom is mutated each cycle. type ShardedForceReplicationParams struct { // ---- Configuration ---- - Namespace string - Query string - BatchSize int - MaxExecsPerShard int - ListWorkflowsPageSize int - TargetClusterEndpoint string - TargetClusterName string - TargetClusterShardCount int32 - DisableVerification bool + Namespace string + Query string + BatchSize int + MaxExecsPerShard int + ListWorkflowsPageSize int + TargetClusterName string + DisableVerification bool ShardNoProgress time.Duration DrainGrace time.Duration @@ -238,11 +236,11 @@ type ShardedForceReplicationParams struct { // ConcurrentBatchCount is the absolute ceiling on in-flight // ReplicateBatch activities. Per-shard exclusivity already bounds - // concurrency to TargetClusterShardCount, but at large cluster - // sizes that's well past the worker's concurrent-activity budget. - // This cap keeps the workflow inside that budget and limits the - // cluster blast radius of a single force-rep run. Defaults to - // min(TargetClusterShardCount/4, defaultConcurrentBatchCap). + // concurrency to the target shard count, but at large cluster sizes + // that's well past the worker's concurrent-activity budget. This + // cap keeps the workflow inside that budget and limits the cluster + // blast radius of a single force-rep run. Defaults to + // min(targetShardCount/4, defaultConcurrentBatchCap). ConcurrentBatchCount int // EstimationMultiplier sizes the QPSQueue's initial slice capacity @@ -314,8 +312,7 @@ type shardedBatchReq struct { NamespaceID string Executions BatchPayload - TargetClusterEndpoint string - TargetClusterName string + TargetClusterName string Resume bool DisableVerification bool diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 25c08d010cc..11f42735c4a 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -109,8 +109,8 @@ func validateShardedForceReplicationParams(params *ShardedForceReplicationParams if len(params.Namespace) == 0 { return temporal.NewNonRetryableApplicationError("InvalidArgument: Namespace is required", "InvalidArgument", nil) } - if !params.DisableVerification && len(params.TargetClusterEndpoint) == 0 && len(params.TargetClusterName) == 0 { - return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterEndpoint or TargetClusterName is required with verification enabled", "InvalidArgument", nil) + if len(params.TargetClusterName) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterName is required", "InvalidArgument", nil) } return nil } @@ -174,6 +174,12 @@ type shardedWorkflowState struct { namespaceID string + // targetShardCount is the target cluster's history shard count, + // fetched once via DescribeTargetCluster at state construction. + // Drives the per-exec shard hash (so packing groups execs by their + // destination shard) and the default ConcurrentBatchCount. + targetShardCount int32 + // buckets accumulate execs that have been listed but not yet // dispatched. Nested by destination shard then businessID, so a // hot BID's many runs share one BID-string-worth of bytes when @@ -258,8 +264,11 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati if err := workflow.ExecuteActivity(metaCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { return nil, err } - if params.TargetClusterShardCount <= 0 { - params.TargetClusterShardCount = md.ShardCount + var targetMd DescribeTargetClusterResponse + if err := workflow.ExecuteActivity(metaCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ + TargetClusterName: params.TargetClusterName, + }).Get(ctx, &targetMd); err != nil { + return nil, err } if params.BatchSize <= 0 { params.BatchSize = defaultBatchSize @@ -283,7 +292,7 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati params.PerBatchGenerateRPS = defaultPerBatchGenerateRPS } if params.ConcurrentBatchCount <= 0 { - params.ConcurrentBatchCount = defaultConcurrentBatchCount(params.TargetClusterShardCount) + params.ConcurrentBatchCount = defaultConcurrentBatchCount(targetMd.ShardCount) } if params.EstimationMultiplier <= 0 { params.EstimationMultiplier = 2 @@ -297,13 +306,14 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati params.QPSQueue.Enqueue(ctx, params.ReplicatedWorkflowCount) } s := &shardedWorkflowState{ - params: params, - namespaceID: md.NamespaceID, - buckets: BatchPayload{}, - bucketCounts: map[int32]int{}, - shardInFlight: map[int32]bool{}, - heldByBatch: map[int64]map[int32]bool{}, - batchExecs: map[int64]BatchPayload{}, + params: params, + namespaceID: md.NamespaceID, + targetShardCount: targetMd.ShardCount, + buckets: BatchPayload{}, + bucketCounts: map[int32]int{}, + shardInFlight: map[int32]bool{}, + heldByBatch: map[int64]map[int32]bool{}, + batchExecs: map[int64]BatchPayload{}, metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ metrics.OperationTagName: metrics.MigrationWorkflowScope, NamespaceTagName: params.Namespace, @@ -371,7 +381,7 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { return err } for _, ex := range listResp.Executions { - sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.params.TargetClusterShardCount) + sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.targetShardCount) s.addToBucket(sh, ex.BusinessID, RunEntry{ RunID: ex.RunID, ArchetypeID: ex.ArchetypeID, @@ -780,19 +790,18 @@ func (s *shardedWorkflowState) spawnBatch( batchID := s.nextBatchID req := &shardedBatchReq{ - BatchID: batchID, - Namespace: s.params.Namespace, - NamespaceID: s.namespaceID, - Executions: payload, - TargetClusterEndpoint: s.params.TargetClusterEndpoint, - TargetClusterName: s.params.TargetClusterName, - Resume: resume, - DisableVerification: s.params.DisableVerification, - NoProgressByShard: noProgressByShard, - PerBatchGenerateRPS: s.params.PerBatchGenerateRPS, - ShardNoProgress: s.params.ShardNoProgress, - DrainGrace: s.params.DrainGrace, - IdleShardCost: s.params.IdleShardCost, + BatchID: batchID, + Namespace: s.params.Namespace, + NamespaceID: s.namespaceID, + Executions: payload, + TargetClusterName: s.params.TargetClusterName, + Resume: resume, + DisableVerification: s.params.DisableVerification, + NoProgressByShard: noProgressByShard, + PerBatchGenerateRPS: s.params.PerBatchGenerateRPS, + ShardNoProgress: s.params.ShardNoProgress, + DrainGrace: s.params.DrainGrace, + IdleShardCost: s.params.IdleShardCost, } held := make(map[int32]bool, len(payload)) diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 8971b41c692..81cc7087173 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -125,6 +125,9 @@ func registerShardedScaffoldingWithSeed( seed func(context.Context, TaskQueueUserDataReplicationParamsWithNamespace) error, ) { env.RegisterActivityWithOptions(metadataResponseFor(shardCount), activity.RegisterOptions{Name: "GetMetadata"}) + env.RegisterActivityWithOptions(func(_ context.Context, _ DescribeTargetClusterRequest) (*DescribeTargetClusterResponse, error) { + return &DescribeTargetClusterResponse{ShardCount: shardCount}, nil + }, activity.RegisterOptions{Name: "DescribeTargetCluster"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.CountWorkflowExecutionsRequest) (*countWorkflowResponse, error) { return &countWorkflowResponse{WorkflowCount: 0}, nil }, activity.RegisterOptions{Name: "CountWorkflow"}) @@ -163,9 +166,8 @@ func TestSharded_HappyPath_SingleCycle(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 4, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted(), "workflow should complete") @@ -213,10 +215,9 @@ func TestSharded_ResumeShards_Packed(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 8, - ResumeShards: resumeShards, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + ResumeShards: resumeShards, }) require.True(t, env.IsWorkflowCompleted()) @@ -299,10 +300,9 @@ func TestSharded_ReleaseShards_FreesShardForReuse(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 2, - ConcurrentBatchCount: 2, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + ConcurrentBatchCount: 2, }) require.True(t, env.IsWorkflowCompleted()) @@ -329,9 +329,8 @@ func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 2, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted()) @@ -396,11 +395,10 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 10, - BatchSize: 100, - MaxExecsPerShard: 10, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + BatchSize: 100, + MaxExecsPerShard: 10, }) require.True(t, env.IsWorkflowCompleted()) @@ -465,9 +463,8 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 4, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted()) @@ -516,10 +513,9 @@ func TestSharded_DisableVerification_NoVerifiedCount(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 4, - DisableVerification: true, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + DisableVerification: true, }) require.True(t, env.IsWorkflowCompleted(), "workflow should complete") @@ -534,9 +530,8 @@ func TestSharded_DisableVerification_NoVerifiedCount(t *testing.T) { } // TestSharded_InvalidInput: validateShardedForceReplicationParams -// rejects an empty Namespace and a missing TargetClusterEndpoint / -// TargetClusterName when verification is enabled. Mirrors the existing -// force-replication TestInvalidInput. +// rejects an empty Namespace and a missing TargetClusterName. Mirrors +// the existing force-replication TestInvalidInput. func TestSharded_InvalidInput(t *testing.T) { for _, tc := range []struct { name string @@ -547,7 +542,7 @@ func TestSharded_InvalidInput(t *testing.T) { params: ShardedForceReplicationParams{}, }, { - name: "missing target with verification on", + name: "missing target cluster name", params: ShardedForceReplicationParams{ Namespace: "test-ns", }, @@ -589,9 +584,8 @@ func TestSharded_ListWorkflowsError(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 2, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted()) @@ -618,9 +612,8 @@ func TestSharded_ReplicateBatchRetryableError(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 2, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted()) @@ -654,9 +647,8 @@ func TestSharded_TaskQueueReplicationFailure(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - TargetClusterShardCount: 2, + Namespace: "test-ns", + TargetClusterName: "remote_cluster", }) require.True(t, env.IsWorkflowCompleted()) diff --git a/tests/xdc/user_data_replication_test.go b/tests/xdc/user_data_replication_test.go index fcf0dd26d29..a8b18bc915f 100644 --- a/tests/xdc/user_data_replication_test.go +++ b/tests/xdc/user_data_replication_test.go @@ -506,9 +506,11 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand( // replication queue — but exercised through the sharded workflow // (registered name "force-replication-sharded" on // MigrationShardedActivityTQ). DisableVerification:true mirrors the -// legacy test's EnableVerification:false default and lets the workflow -// run without a TargetClusterName since this test only exercises the -// task-queue-user-data side of force-replication. +// legacy test's EnableVerification:false default; TargetClusterName is +// still set (validation requires it, and the workflow fetches the +// remote shard count even when there are no execs to replicate) but +// the inject path is never reached here since this test only exercises +// the task-queue-user-data side of force-replication. func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_Sharded() { ctx := testcore.NewContext() activeFrontendClient := s.clusters[0].FrontendClient() @@ -593,7 +595,8 @@ func (s *UserDataReplicationTestSuite) TestUserDataEntriesAreReplicatedOnDemand_ WorkflowRunTimeout: time.Second * 30, }, "force-replication-sharded", migration.ShardedForceReplicationParams{ Namespace: namespace, - DisableVerification: true, // mirrors legacy test's EnableVerification:false (no target needed) + TargetClusterName: s.clusters[1].ClusterName(), + DisableVerification: true, }) s.NoError(err) err = run.Get(ctx, nil) From e42d5f4f4b0a9020b5711fd23cfbadeee6b90f1e Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 15:17:25 +0100 Subject: [PATCH 08/35] Improve error handling so that executions aren't lost. If we exit due to an error, ensure that an operator can restart the workflow without missing any executions from previous pages that had not yet been handled. There were a few edge cases where some might be dropped. --- .../migration/force_replication_workflow.go | 11 + service/worker/migration/sharded_activity.go | 54 +++- service/worker/migration/sharded_workflow.go | 298 ++++++++++++++---- .../worker/migration/sharded_workflow_test.go | 251 +++++++++++++++ 4 files changed, 546 insertions(+), 68 deletions(-) diff --git a/service/worker/migration/force_replication_workflow.go b/service/worker/migration/force_replication_workflow.go index be2c330b45b..e7a6899369d 100644 --- a/service/worker/migration/force_replication_workflow.go +++ b/service/worker/migration/force_replication_workflow.go @@ -86,6 +86,17 @@ type ( ReplicatedWorkflowCount int64 ReplicatedWorkflowCountPerSecond float64 PageTokenForRestart []byte + + // Sharded-workflow-only recovery bundle: feed these three + // fields back into a fresh ShardedForceReplicationWorkflow's + // NextPageToken / ResumeShards / RecoveredBuckets params to + // resume from a failed run without missing executions. Left + // zero by the legacy ForceReplicationWorkflow variants — + // their PageTokenForRestart is the start-of-run token and + // already covers all in-flight execs at restart cost. + RecoveryNextPageToken []byte + RecoveryResumeShards []ResumeShard + RecoveryBuckets BatchPayload } ) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index 955c40e3af2..ad46be7d2dd 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -147,16 +147,21 @@ func (a *activities) runVerifyPhase( passDelta, minNextRetry, ctxAborted, vErr := a.runVerifyPass( callCtx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) + // Fold partial progress in before the error check — the SDK + // discards the activity result on failure, so the only way + // the workflow learns about partially-verified execs on the + // error path is via wrapBatchVerifyError encoding the count + // as ApplicationError details below. + doneCount += passDelta if vErr != nil { - return replicateBatchResult{}, vErr + return replicateBatchResult{}, wrapBatchVerifyError(vErr, int64(doneCount)) } - doneCount += passDelta activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) if done, result, err := a.evaluateVerifyIteration( ctx, req, execs, verified, shards, doneCount, of, draining, drainStartAt); err != nil { - return replicateBatchResult{}, err + return replicateBatchResult{}, wrapBatchVerifyError(err, int64(doneCount)) } else if done { return result, nil } @@ -355,6 +360,14 @@ func (a *activities) shouldExitDrain(req *shardedBatchReq, shards shardVerifyTra // completed-but-unsignaled shards if their cumulative idle cost crossed // the threshold. Only fires in normal mode — drain mode rides the // activity result instead. +// +// Ctx-canceled errors from signalReleaseShards are suppressed: a +// workflow-initiated cancel arriving mid-signal would otherwise +// surface as a wrapped ctx-canceled error (not temporal.CanceledError) +// that the workflow side wouldn't recognise via IsCanceledError — +// turning a clean CAN into an error exit. Suppressing here lets the +// outer loop see ctx.Err() at its top and promote to drain mode +// normally. func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchReq, shards shardVerifyTracker) error { if shards.totalIdleCost(time.Now()) < req.IdleShardCost { return nil @@ -364,6 +377,9 @@ func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchRe return nil } if err := a.signalReleaseShards(ctx, req, releaseList); err != nil { + if ctx.Err() != nil { + return nil + } return err } shards.markReleased(releaseList) @@ -716,6 +732,38 @@ func (a *activities) firstUnverifiedOnShard(execs []*shardedExecutionInfo, verif return 0, false } +// batchVerifyPartialErrorType is the ApplicationError Type stamped on +// wrappers produced by wrapBatchVerifyError. The workflow keys off +// this Type via extractVerifiedCountFromError to disambiguate "the +// wrapper we made" from any other ApplicationError carrying an +// int64. The original error is reachable via Unwrap on the wrapper. +const batchVerifyPartialErrorType = "BatchVerifyPartial" + +// wrapBatchVerifyError wraps the verify-phase error so the partial +// VerifiedCount survives the activity boundary — the SDK discards +// the activity result on failure, so the count would otherwise be +// lost. The original error is attached as Cause; the workflow side +// reaches it via errors.As / Unwrap as usual. Returns the cause +// unchanged when there's no progress to report. +func wrapBatchVerifyError(cause error, verifiedCount int64) error { + if cause == nil || verifiedCount <= 0 { + return cause + } + nonRetryable := false + if appErr, ok := errors.AsType[*temporal.ApplicationError](cause); ok { + nonRetryable = appErr.NonRetryable() + } + return temporal.NewApplicationErrorWithOptions( + cause.Error(), + batchVerifyPartialErrorType, + temporal.ApplicationErrorOptions{ + Cause: cause, + Details: []any{verifiedCount}, + NonRetryable: nonRetryable, + }, + ) +} + // backoffDelay returns the per-exec retry delay after `attempt` // consecutive failed verify attempts: 100ms × 2^(attempt-1), capped at // 5s. The cap bounds how long after the apply pipeline recovers we'd diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 11f42735c4a..9740d393853 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -1,6 +1,7 @@ package migration import ( + "errors" "fmt" "slices" "time" @@ -30,23 +31,41 @@ import ( // to finish naturally and returns nil. func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { // Page token at workflow entry — returned by the status query as - // PageTokenForRestart so an operator can resume from this run's - // starting position rather than its current (in-flight) position. + // PageTokenForRestart so tooling that already knows the legacy + // "restart from the starting position" semantic keeps working. + // The richer Recovery* fields below carry the current page token + // plus in-flight execs and are what the sharded restart flow + // actually uses. startPageToken := params.NextPageToken + // state is assigned after newShardedWorkflowState below; the + // query handler closes over the pointer so it sees the live state + // once setup completes. Queries that arrive during setup return + // the static fields without the recovery bundle, which matches + // the prior behaviour. + var state *shardedWorkflowState + // Register the status query under the same name upstream uses // (forceReplicationStatusQueryType = "force-replication-status") // so tooling that polls force-rep progress works across both // workflow variants. if err := workflow.SetQueryHandler(ctx, forceReplicationStatusQueryType, func() (ForceReplicationStatus, error) { - return ForceReplicationStatus{ + status := ForceReplicationStatus{ ContinuedAsNewCount: params.ContinuedAsNewCount, TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, PageTokenForRestart: startPageToken, TaskQueueUserDataReplicationStatus: params.TaskQueueUserDataReplicationStatus, - }, nil + RecoveryNextPageToken: params.NextPageToken, + RecoveryResumeShards: params.ResumeShards, + RecoveryBuckets: params.RecoveredBuckets, + } + if state != nil { + status.RecoveryResumeShards = state.collectResumeShardsForCarryover() + status.RecoveryBuckets = state.collectRecoveredBucketsForCarryover() + } + return status, nil }); err != nil { return err } @@ -55,7 +74,8 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe return err } - state, err := newShardedWorkflowState(ctx, ¶ms) + var err error + state, err = newShardedWorkflowState(ctx, ¶ms) if err != nil { return err } @@ -364,7 +384,9 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { listCtx := workflow.WithActivityOptions(ctx, ao) // Drive ListWorkflows until either we exhaust the namespace or - // the SDK signals that history is large enough to CAN. + // the SDK signals that history is large enough to CAN. Errors + // here latch into lastErr and fall through to the unified exit + // funnel — same drain-and-decide path as activity-driven errors. for !workflow.GetInfo(ctx).GetContinueAsNewSuggested() { if s.lastErr != nil { break @@ -378,7 +400,8 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { var a *activities var listResp listWorkflowsResponse if err := workflow.ExecuteActivity(listCtx, a.ListWorkflows, listReq).Get(ctx, &listResp); err != nil { - return err + s.setLastErr(err) + break } for _, ex := range listResp.Executions { sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.targetShardCount) @@ -397,51 +420,120 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { } } - // Drain remaining buckets — dispatch every leftover exec. The - // dispatched batches may themselves get cancelled by drainForCAN - // below if we're going to CAN; that's fine, each cancelled - // batch returns its drain payload in its result. + // Drain remaining buckets only when no error has latched — on + // error we deliberately stop scheduling new work and let the + // already-dispatched batches finish via awaitInFlightCompletion. if s.lastErr == nil { s.drainBuckets(ctx) } - // If there are no more pages we're done: wait for activities to - // finish naturally and return. Slow shards hold their claims - // until their per-shard no-progress backstop trips. - if len(s.params.NextPageToken) == 0 || s.lastErr != nil { - _ = workflow.Await(ctx, func() bool { - return s.pendingDispatches == 0 || s.lastErr != nil - }) - if s.lastErr != nil { - return s.lastErr - } - return nil - } - - // More pages → CAN. Cancel in-flight batches, wait for them to - // drain, then harvest any leftover batchExecs entries (batches - // whose activity was cancelled before its body ran — no result - // returned) into RecoveredBuckets so they get re-dispatched as - // fresh inject+verify activities next cycle. - s.drainForCAN(ctx) + // Wait for in-flight activities. Cancels them when we already + // know we're failing or CAN-ing; on the clean success path + // (no error, no more pages) waits naturally so a healthy + // activity isn't cancelled into a CanceledError that masquerades + // as carry-over state. + s.awaitInFlightCompletion(ctx) + // Single exit decision. Recovery state has the same shape on + // either path — drainPayload + undispatched ResumeShards + + // batchExecs + leftover buckets — so the only difference between + // "fail with state" and "CAN with state" is the return value. if s.lastErr != nil { return s.lastErr } + if !s.hasCarryover() { + return nil + } next := *s.params next.ContinuedAsNewCount++ - // Defensive copy so we don't alias s.drainPayload into the - // carry-over params. drainForCAN guarantees every spawnBatch - // coroutine has finished appending before we get here (the - // append happens before the defer that decrements - // pendingDispatches), so this is style — but cheap insurance - // against future code paths that append post-drain. - next.ResumeShards = append([]ResumeShard(nil), s.drainPayload...) - next.RecoveredBuckets = collectRecoveredBuckets(s.batchExecs) + next.ResumeShards = s.collectResumeShardsForCarryover() + next.RecoveredBuckets = s.collectRecoveredBucketsForCarryover() return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) } +// awaitInFlightCompletion drains in-flight batches before the workflow +// exits. The strategy depends on what we already know: +// +// - Error latched or more pages remain (we're going to fail or CAN +// either way): cancel immediately via drainForCAN, bounded by +// DrainGrace + IdleShardCost. No point waiting for activities +// that are going to be discarded. +// - Clean success path (no error, no more pages): wait for natural +// completion so a healthy activity's clean result isn't masked +// as a CanceledError. If an activity hits its ShardNoProgress +// backstop mid-wait, latch the error then cancel the rest fast +// rather than waiting on every shard's backstop too. +func (s *shardedWorkflowState) awaitInFlightCompletion(ctx workflow.Context) { + if s.pendingDispatches == 0 { + return + } + if s.lastErr != nil || len(s.params.NextPageToken) > 0 { + s.drainForCAN(ctx) + return + } + _ = workflow.Await(ctx, func() bool { + return s.pendingDispatches == 0 || s.lastErr != nil + }) + if s.pendingDispatches > 0 { + s.drainForCAN(ctx) + } +} + +// hasCarryover reports whether the workflow has any state worth +// preserving across an exit — either a remaining page token, drained +// execs from in-flight batches, cancel-before-start batches that +// never injected, undispatched resume entries, or listed-but-unpacked +// execs. Drives both the "CAN vs return nil" decision and the +// recovery bundle exposed in the status query. +func (s *shardedWorkflowState) hasCarryover() bool { + if len(s.params.NextPageToken) > 0 { + return true + } + if len(s.drainPayload) > 0 { + return true + } + if len(s.batchExecs) > 0 { + return true + } + if len(s.params.ResumeShards) > 0 { + return true + } + return !s.bucketsEmpty() +} + +// collectResumeShardsForCarryover concatenates this cycle's drained +// execs with any prior-cycle ResumeShards that didn't get dispatched +// (left in params.ResumeShards by dispatchResumeBatches when it bailed +// out on lastErr). Both groups are already shard-keyed; the next +// cycle's dispatchResumeBatches sorts and re-packs them. +func (s *shardedWorkflowState) collectResumeShardsForCarryover() []ResumeShard { + if len(s.drainPayload) == 0 && len(s.params.ResumeShards) == 0 { + return nil + } + out := make([]ResumeShard, 0, len(s.drainPayload)+len(s.params.ResumeShards)) + out = append(out, s.drainPayload...) + out = append(out, s.params.ResumeShards...) + return out +} + +// collectRecoveredBucketsForCarryover merges the two sources of +// "execs that never made it through a verify activity this cycle": +// batches that returned CanceledError without running a body, and +// listed-but-unpacked execs still sitting in s.buckets when the +// workflow exited (either lastErr stopped the streaming packer or +// drainBuckets bailed out on lastErr partway through). +func (s *shardedWorkflowState) collectRecoveredBucketsForCarryover() BatchPayload { + out := collectRecoveredBuckets(s.batchExecs) + if !s.bucketsEmpty() { + if out == nil { + out = BatchPayload{} + } + out.merge(s.buckets) + } + return out +} + // recordVerified accumulates one batch's verified-exec delta into the // workflow's running count, emits the per-batch counter delta, and // updates the sliding-window RPS gauge. No-op when verified == 0 so a @@ -577,6 +669,38 @@ func (s *shardedWorkflowState) handleReleaseSignals(ctx workflow.Context) { } } +// extractVerifiedCountFromError pulls the partial VerifiedCount that +// wrapBatchVerifyError encoded into a BatchVerifyPartial-typed +// ApplicationError's Details on the activity side. Returns 0 when +// the error didn't come through the verify-phase wrapper (e.g. +// inject-phase failures, ctx errors, non-ApplicationError types), so +// callers can unconditionally fold the result into recordVerified. +func extractVerifiedCountFromError(err error) int64 { + if err == nil { + return 0 + } + appErr, ok := errors.AsType[*temporal.ApplicationError](err) + if !ok || appErr.Type() != batchVerifyPartialErrorType { + return 0 + } + var count int64 + if appErr.Details(&count) != nil { + return 0 + } + return count +} + +// setLastErr latches the first error encountered. Subsequent errors +// are dropped so the root cause is preserved for the workflow's +// returned failure — without the latch, a stuck-shard backstop firing +// on every batch as the workflow tears down would overwrite the +// genuinely interesting first failure. +func (s *shardedWorkflowState) setLastErr(err error) { + if s.lastErr == nil { + s.lastErr = err + } +} + // dispatchSlotAvailable returns true when the workflow is below the // in-flight batch ceiling and is free to spawn another batch. Callers // that can defer dispatch (the streaming packer) consult this and @@ -595,6 +719,16 @@ func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { }) } +// resumeBatch is one packed dispatch plan: the BatchPayload that will +// become a batch's input, plus the matching per-shard no-progress +// durations. Built up front by packResumeBatchPlan so the dispatch +// loop can unpack any remainder back into ResumeShards if lastErr +// trips mid-dispatch. +type resumeBatch struct { + payload BatchPayload + noProgress map[int32]time.Duration +} + // dispatchResumeBatches turns the prior cycle's drain payload into a // fresh round of resume activities, packed across shards up to // BatchSize per batch. Each shard appears at most once across the @@ -604,6 +738,13 @@ func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { // contributions are taken whole and no MaxExecsPerShard cap applies — // resume carries no inject load so the per-shard blast-radius the // streaming packer guards against doesn't exist here. +// +// Plans every batch up front, then dispatches one at a time. If +// lastErr latches mid-dispatch, the remaining planned batches are +// unpacked back into s.params.ResumeShards so the recovery bundle +// (and the next CAN cycle) sees them — without the unpack step, a +// failing first resume batch would silently strand all subsequent +// resume entries. func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { if len(s.params.ResumeShards) == 0 { return @@ -615,15 +756,17 @@ func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { } entries = append(entries, rs) } + // We've taken ownership of these entries — anything not + // dispatched gets restored below. + s.params.ResumeShards = nil slices.SortFunc(entries, func(a, b ResumeShard) int { return int(a.Shard - b.Shard) }) - payload := BatchPayload{} - packed := 0 - packNoProgress := map[int32]time.Duration{} - flush := func() { - if packed == 0 { + batches := s.packResumeBatchPlan(entries) + for i, batch := range batches { + if s.lastErr != nil { + s.params.ResumeShards = unpackResumeBatches(batches[i:]) return } // Block until a dispatch slot is free so resume payloads @@ -631,38 +774,57 @@ func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { // carried many shards across CAN. s.waitForDispatchSlot(ctx) if s.lastErr != nil { + s.params.ResumeShards = unpackResumeBatches(batches[i:]) return } - for sh := range payload { + for sh := range batch.payload { s.shardInFlight[sh] = true } - s.spawnBatch(ctx, payload, true, packNoProgress) - payload = BatchPayload{} - packed = 0 - packNoProgress = map[int32]time.Duration{} + s.spawnBatch(ctx, batch.payload, true, batch.noProgress) } +} +// packResumeBatchPlan groups ResumeShard entries into batches sized +// at or below BatchSize. Entries are taken whole — shardInFlight only +// admits one batch per shard at a time, so a single shard's payload +// can't be split. +func (s *shardedWorkflowState) packResumeBatchPlan(entries []ResumeShard) []resumeBatch { + var batches []resumeBatch + current := resumeBatch{payload: BatchPayload{}, noProgress: map[int32]time.Duration{}} + packed := 0 for _, rs := range entries { - if s.lastErr != nil { - return - } rsCount := runCount(rs.Execs) - // Flush before this shard if it would push us over the - // batch cap, so each batch stays within BatchSize. A - // single shard's contribution is taken whole — we don't - // split a shard across batches because shardInFlight only - // admits one batch per shard at a time. if packed+rsCount > s.params.BatchSize && packed > 0 { - flush() + batches = append(batches, current) + current = resumeBatch{payload: BatchPayload{}, noProgress: map[int32]time.Duration{}} + packed = 0 } - payload[rs.Shard] = rs.Execs - packNoProgress[rs.Shard] = rs.NoProgressDuration + current.payload[rs.Shard] = rs.Execs + current.noProgress[rs.Shard] = rs.NoProgressDuration packed += rsCount - if packed >= s.params.BatchSize { - flush() + } + if packed > 0 { + batches = append(batches, current) + } + return batches +} + +// unpackResumeBatches reverses packResumeBatchPlan, turning planned +// batches back into a flat ResumeShard slice. Used when the dispatch +// loop aborts on lastErr so the undispatched remainder can be carried +// into the recovery bundle / next CAN cycle. +func unpackResumeBatches(batches []resumeBatch) []ResumeShard { + var out []ResumeShard + for _, b := range batches { + for sh, execs := range b.payload { + out = append(out, ResumeShard{ + Shard: sh, + Execs: execs, + NoProgressDuration: b.noProgress[sh], + }) } } - flush() + return out } // runCount sums runs across BIDs in a single shard's payload entry. @@ -707,9 +869,9 @@ func (s *shardedWorkflowState) failDrainBucketsStuck() { for _, n := range s.bucketCounts { remaining += n } - s.lastErr = temporal.NewNonRetryableApplicationError( + s.setLastErr(temporal.NewNonRetryableApplicationError( fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), - "DrainBucketsStuck", nil) + "DrainBucketsStuck", nil)) } // drainBucketsAwaitPredicate returns true when the drainBuckets loop @@ -862,7 +1024,13 @@ func (s *shardedWorkflowState) spawnBatch( // execs as fresh inject+verify work next cycle. return } - s.lastErr = err + // Activity errored after partial verify — the SDK discards + // the result on failure, but wrapBatchVerifyError on the + // activity side carries the partial doneCount through as + // ApplicationError details. Fold it into the running count + // so ReplicatedWorkflowCount reflects work actually done. + s.recordVerified(coroCtx, extractVerifiedCountFromError(err)) + s.setLastErr(err) }) } diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 81cc7087173..bb7e852b92c 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -2,6 +2,7 @@ package migration import ( "context" + "errors" "fmt" "strconv" "sync" @@ -664,6 +665,256 @@ func TestSharded_TaskQueueReplicationFailure(t *testing.T) { require.Contains(t, status.TaskQueueUserDataReplicationStatus.FailureMessage, "namespace is required") } +// TestSharded_RecoveryBundle_OnBatchError: a batch returns a +// non-retryable error mid-cycle; the workflow latches lastErr, drains +// in-flight via cancellation, returns the error, and the status query +// reports a non-empty recovery bundle so an operator can start a +// fresh run with all unverified execs preserved. +func TestSharded_RecoveryBundle_OnBatchError(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 4) + + // One page of execs across 4 shards. + execs := makeExecs(4, 5) // 20 execs total + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + // Every batch fails non-retryably so lastErr latches on the first + // return and subsequent in-flight batches are cancelled. + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( + "batch failed", "BatchFailed", nil) + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + }) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err) + var appErr *temporal.ApplicationError + require.ErrorAs(t, err, &appErr) + require.Equal(t, "BatchFailed", appErr.Type()) + + envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + + // Recovery bundle: the cancelled in-flight batches go into + // RecoveryBuckets (collectRecoveredBuckets on batchExecs). The + // failed batch's execs land there too — its activity attempt + // errored after running, so batchExecs[id] is still populated. + require.NotEmpty(t, status.RecoveryBuckets, + "failed run must expose its in-flight execs as RecoveryBuckets") + recovered := 0 + for _, byBID := range status.RecoveryBuckets { + for _, runs := range byBID { + recovered += len(runs) + } + } + require.Equal(t, len(execs), recovered, + "every listed exec should be recoverable via the bundle") +} + +// TestSharded_RecoveryBundle_PreservesUndispatchedResumeShards: when +// the first resume batch errors and latches lastErr, the remaining +// undispatched resume entries must be preserved in the recovery +// bundle rather than silently dropped. +func TestSharded_RecoveryBundle_PreservesUndispatchedResumeShards(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 8) + env.RegisterActivityWithOptions(pageThrough(nil, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + // 8 shards × 60 execs each = 480 unverified. BatchSize=100 → + // dispatchResumeBatches plans ~5 batches; the first one fails, + // latches lastErr, and the remaining 4 must be preserved. + resumeShards := make([]ResumeShard, 8) + for s := range 8 { + resumeShards[s] = ResumeShard{ + Shard: int32(s), + Execs: makeExecsForShard(int32(s), 60), + } + } + + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( + "resume batch failed", "ResumeFailed", nil) + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + ConcurrentBatchCount: 1, // serialize so we can observe the early-bail behaviour + ResumeShards: resumeShards, + }) + + require.True(t, env.IsWorkflowCompleted()) + require.Error(t, env.GetWorkflowError()) + + envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + + // Every shard must show up exactly once across RecoveryResumeShards + // (still-undispatched, batched-but-cancelled, or the failed batch + // itself — all paths fold back into the recovery bundle). + recoveredRuns := 0 + for _, rs := range status.RecoveryResumeShards { + for _, runs := range rs.Execs { + recoveredRuns += len(runs) + } + } + for _, byBID := range status.RecoveryBuckets { + for _, runs := range byBID { + recoveredRuns += len(runs) + } + } + require.Equal(t, 8*60, recoveredRuns, + "all 480 resume execs should be recoverable; got %d", recoveredRuns) +} + +// TestSharded_RecoveryBundle_TracksCurrentPageToken: a List error +// after the first page should preserve the page token of the next +// page to read, not the start-of-run token. +func TestSharded_RecoveryBundle_TracksCurrentPageToken(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + + // First call succeeds; second call errors. Workflow processes page + // 1 successfully, then fails listing page 2 with NextPageToken + // pointing past page 1. + var listCalls atomic.Int32 + page1 := makeExecs(2, 3) + env.RegisterActivityWithOptions(func(_ context.Context, req *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + n := listCalls.Add(1) + if n == 1 { + return &listWorkflowsResponse{ + Executions: page1, + NextPageToken: []byte("page-2"), + }, nil + } + return nil, temporal.NewNonRetryableApplicationError( + "list page 2 failed", "ListFailed", nil) + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + // Batches succeed so page 1 doesn't pollute the recovery bundle — + // we want the page-token assertion isolated. + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{}, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + }) + + require.True(t, env.IsWorkflowCompleted()) + require.Error(t, env.GetWorkflowError()) + + envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + + require.Equal(t, []byte("page-2"), status.RecoveryNextPageToken, + "RecoveryNextPageToken should reflect where listing was about to resume, not start-of-run") +} + +// TestWrapBatchVerifyError_RoundTrip: the workflow must be able to +// recover the partial VerifiedCount that the activity encoded on a +// failed batch return. Covers retryability preservation, the +// no-progress short-circuit, and inner-error reachability via +// Unwrap (so consumers that care about the underlying Type can walk +// past the wrapper). +func TestWrapBatchVerifyError_RoundTrip(t *testing.T) { + t.Run("non-zero count survives wrap; inner reachable via Unwrap", func(t *testing.T) { + cause := temporal.NewNonRetryableApplicationError("stuck", "ShardNoProgress", nil) + wrapped := wrapBatchVerifyError(cause, 42) + + // Outer wrapper carries the partial-verify tag and the count. + var appErr *temporal.ApplicationError + require.ErrorAs(t, wrapped, &appErr) + require.Equal(t, batchVerifyPartialErrorType, appErr.Type()) + require.True(t, appErr.NonRetryable()) + require.Equal(t, int64(42), extractVerifiedCountFromError(wrapped)) + + // Inner identity is preserved via Cause / Unwrap — consumers + // that key off the underlying type still get there. + var inner *temporal.ApplicationError + require.ErrorAs(t, appErr.Unwrap(), &inner) + require.Equal(t, "ShardNoProgress", inner.Type()) + }) + + t.Run("zero count returns cause unchanged", func(t *testing.T) { + cause := temporal.NewApplicationError("trivial", "X") + require.Same(t, cause, wrapBatchVerifyError(cause, 0)) + require.Equal(t, int64(0), extractVerifiedCountFromError(cause)) + }) + + t.Run("non-ApplicationError cause is wrapped and remains reachable", func(t *testing.T) { + cause := errors.New("plain failure") + wrapped := wrapBatchVerifyError(cause, 7) + var appErr *temporal.ApplicationError + require.ErrorAs(t, wrapped, &appErr) + require.Equal(t, batchVerifyPartialErrorType, appErr.Type()) + require.Equal(t, int64(7), extractVerifiedCountFromError(wrapped)) + require.ErrorIs(t, wrapped, cause) + }) + + t.Run("unrelated ApplicationError returns 0", func(t *testing.T) { + // An ApplicationError that wasn't produced by wrapBatchVerifyError + // — extractVerifiedCountFromError must not pull garbage out of it. + other := temporal.NewApplicationError("plain", "Other") + require.Equal(t, int64(0), extractVerifiedCountFromError(other)) + }) +} + +// TestSharded_PartialVerifiedCount_RecordedOnError: when a batch +// errors out after partial verify, the wrapped error's count must be +// folded into ReplicatedWorkflowCount so the status query reflects +// work actually done rather than zeroing out a partially-successful +// batch. +func TestSharded_PartialVerifiedCount_RecordedOnError(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 2) + env.RegisterActivityWithOptions(pageThrough(makeExecs(2, 5), 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + const partialCount int64 = 6 + env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{}, wrapBatchVerifyError( + temporal.NewNonRetryableApplicationError("simulated mid-verify failure", "Simulated", nil), + partialCount, + ) + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + }) + + require.True(t, env.IsWorkflowCompleted()) + require.Error(t, env.GetWorkflowError()) + + envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, envValue.Get(&status)) + + require.GreaterOrEqual(t, status.ReplicatedWorkflowCount, partialCount, + "failed batch's partial count should still be reflected in ReplicatedWorkflowCount") +} + // ---- internal helpers ---- // makeExecsForShard produces `count` runs for the named shard's From f2947279458755a5885900893bfdb3b08e872a11 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 15:43:03 +0100 Subject: [PATCH 09/35] Style. --- service/worker/migration/sharded_types.go | 4 - service/worker/migration/sharded_workflow.go | 145 +++++++++++-------- 2 files changed, 82 insertions(+), 67 deletions(-) diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 32460089e3f..1de0ba06ce8 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -14,10 +14,6 @@ const ( // time by setting the start request's workflow type. shardedForceReplicationWorkflowName = "force-replication-sharded" - // shardedBatchActivityName is the registered activity name used by - // name-based dispatch. - shardedBatchActivityName = "ReplicateBatch" - // releaseShardsSignalName carries mid-flight ReleaseShards signals // from active replicate-batch activities back to their parent // workflow. Drain-mode shard completions ride the activity return diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 9740d393853..2ef230c4b69 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -373,16 +373,6 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { // from racing to dispatch against them with fresh execs. s.dispatchResumeBatches(ctx) - ao := workflow.ActivityOptions{ - StartToCloseTimeout: time.Hour, - RetryPolicy: &temporal.RetryPolicy{ - InitialInterval: time.Second, - BackoffCoefficient: 2.0, - MaximumAttempts: 3, - }, - } - listCtx := workflow.WithActivityOptions(ctx, ao) - // Drive ListWorkflows until either we exhaust the namespace or // the SDK signals that history is large enough to CAN. Errors // here latch into lastErr and fall through to the unified exit @@ -391,31 +381,24 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { if s.lastErr != nil { break } - listReq := &workflowservice.ListWorkflowExecutionsRequest{ - Namespace: s.params.Namespace, - Query: s.params.Query, - PageSize: int32(s.params.ListWorkflowsPageSize), - NextPageToken: s.params.NextPageToken, - } - var a *activities - var listResp listWorkflowsResponse - if err := workflow.ExecuteActivity(listCtx, a.ListWorkflows, listReq).Get(ctx, &listResp); err != nil { + executions, nextPageToken, err := s.listWorkflowPage(ctx) + if err != nil { s.setLastErr(err) break } - for _, ex := range listResp.Executions { + for _, ex := range executions { sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.targetShardCount) s.addToBucket(sh, ex.BusinessID, RunEntry{ RunID: ex.RunID, ArchetypeID: ex.ArchetypeID, }) } - s.params.NextPageToken = listResp.NextPageToken + s.params.NextPageToken = nextPageToken for s.tryPackStreaming(ctx, false) { //nolint:revive // intentional empty body } - if len(listResp.NextPageToken) == 0 { + if len(nextPageToken) == 0 { break } } @@ -452,6 +435,58 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) } +var ( + shardedListWorkflowsRetryPolicy = &temporal.RetryPolicy{ + InitialInterval: time.Second, + BackoffCoefficient: 2.0, + MaximumAttempts: 3, + } + shardedListWorkflowsActivityOptions = workflow.ActivityOptions{ + StartToCloseTimeout: time.Hour, + RetryPolicy: shardedListWorkflowsRetryPolicy, + } + + // Per-exec backoff still owns the per-exec retry; MaximumAttempts + // lets a transient activity failure recover via heartbeat-resume + // without losing inject progress. WaitForCancellation lets a + // cancelled activity run drain logic and return its drain result. + shardedReplicateBatchRetryPolicy = &temporal.RetryPolicy{ + MaximumAttempts: 3, + } + shardedReplicateBatchActivityOptions = workflow.ActivityOptions{ + StartToCloseTimeout: 24 * time.Hour, + HeartbeatTimeout: time.Minute, + RetryPolicy: shardedReplicateBatchRetryPolicy, + WaitForCancellation: true, + } +) + +func (s *shardedWorkflowState) listWorkflowPage(ctx workflow.Context) ([]*ExecutionInfo, []byte, error) { + listCtx := workflow.WithActivityOptions(ctx, shardedListWorkflowsActivityOptions) + listReq := &workflowservice.ListWorkflowExecutionsRequest{ + Namespace: s.params.Namespace, + Query: s.params.Query, + PageSize: int32(s.params.ListWorkflowsPageSize), + NextPageToken: s.params.NextPageToken, + } + var a *activities + var listResp listWorkflowsResponse + if err := workflow.ExecuteActivity(listCtx, a.ListWorkflows, listReq).Get(ctx, &listResp); err != nil { + return nil, nil, err + } + return listResp.Executions, listResp.NextPageToken, nil +} + +func (s *shardedWorkflowState) replicateBatch(ctx, activityParentCtx workflow.Context, req *shardedBatchReq) (replicateBatchResult, error) { + actx := workflow.WithActivityOptions(activityParentCtx, shardedReplicateBatchActivityOptions) + var a *activities + var result replicateBatchResult + if err := workflow.ExecuteActivity(actx, a.ReplicateBatch, req).Get(ctx, &result); err != nil { + return replicateBatchResult{}, err + } + return result, nil +} + // awaitInFlightCompletion drains in-flight batches before the workflow // exits. The strategy depends on what we already know: // @@ -986,51 +1021,35 @@ func (s *shardedWorkflowState) spawnBatch( } delete(s.heldByBatch, batchID) }() - ao := workflow.ActivityOptions{ - StartToCloseTimeout: 24 * time.Hour, - HeartbeatTimeout: time.Minute, - // 3 attempts: per-exec backoff still owns the per-exec - // retry, but a transient activity failure can recover - // via heartbeat-resume without losing inject progress. - RetryPolicy: &temporal.RetryPolicy{ - MaximumAttempts: 3, - }, - // WaitForCancellation: the cancelled activity needs to - // run its drain logic and return its drain result - // before the dispatch coroutine's defer fires. - WaitForCancellation: true, - } - actx := workflow.WithActivityOptions(s.activityCtx, ao) - var result replicateBatchResult - err := workflow.ExecuteActivity(actx, shardedBatchActivityName, req).Get(coroCtx, &result) - if err == nil { - // Activity body ran and returned cleanly — either a - // clean completion (empty InFlight) or a drained - // CAN-cancel (InFlight carries the still-unverified - // execs). CompletedShards is informational; the - // defer above clears heldByBatch + shardInFlight - // either way. - if len(result.InFlight) > 0 { - s.drainPayload = append(s.drainPayload, result.InFlight...) + result, err := s.replicateBatch(coroCtx, s.activityCtx, req) + if err != nil { + if temporal.IsCanceledError(err) { + // Cancel-before-start: the activity body never ran, + // so no result is available. Leaving batchExecs[batchID] + // intact lets the CAN-end recovery path re-bucket the + // execs as fresh inject+verify work next cycle. + return } - s.recordVerified(coroCtx, result.VerifiedCount) - delete(s.batchExecs, batchID) + // Activity errored after partial verify — the SDK discards + // the result on failure, but wrapBatchVerifyError on the + // activity side carries the partial doneCount through as + // ApplicationError details. Fold it into the running count + // so ReplicatedWorkflowCount reflects work actually done. + s.recordVerified(coroCtx, extractVerifiedCountFromError(err)) + s.setLastErr(err) return } - if temporal.IsCanceledError(err) { - // Cancel-before-start: the activity body never ran, - // so no result is available. Leaving batchExecs[batchID] - // intact lets the CAN-end recovery path re-bucket the - // execs as fresh inject+verify work next cycle. - return + // Activity body ran and returned cleanly — either a + // clean completion (empty InFlight) or a drained + // CAN-cancel (InFlight carries the still-unverified + // execs). CompletedShards is informational; the + // defer above clears heldByBatch + shardInFlight + // either way. + if len(result.InFlight) > 0 { + s.drainPayload = append(s.drainPayload, result.InFlight...) } - // Activity errored after partial verify — the SDK discards - // the result on failure, but wrapBatchVerifyError on the - // activity side carries the partial doneCount through as - // ApplicationError details. Fold it into the running count - // so ReplicatedWorkflowCount reflects work actually done. - s.recordVerified(coroCtx, extractVerifiedCountFromError(err)) - s.setLastErr(err) + s.recordVerified(coroCtx, result.VerifiedCount) + delete(s.batchExecs, batchID) }) } From b8e7dadddfbde4790b3e82c9531e95b482f3b1c0 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 15:52:11 +0100 Subject: [PATCH 10/35] Tidy up. --- service/worker/migration/sharded_activity.go | 40 ++++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index ad46be7d2dd..42f1837aebe 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -33,8 +33,8 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( // Flatten once so per-exec bookkeeping (verified[], attempts[], // nextRetryAt[]) can stay index-based. execs := req.Executions.flatten() - of := len(execs) - if of == 0 { + execCount := len(execs) + if execCount == 0 { return replicateBatchResult{}, nil } @@ -74,7 +74,7 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) } - return a.runVerifyPhase(ctx, req, execs, of, remoteAdminClient, ns) + return a.runVerifyPhase(ctx, req, execs, execCount, remoteAdminClient, ns) } // runVerifyPhase is the verify-phase loop body of ReplicateBatch. It @@ -86,13 +86,13 @@ func (a *activities) runVerifyPhase( ctx context.Context, req *shardedBatchReq, execs []*shardedExecutionInfo, - of int, + execCount int, remoteAdminClient adminservice.AdminServiceClient, ns *namespace.Namespace, ) (replicateBatchResult, error) { - verified := make([]bool, of) - attempts := make([]int, of) - nextRetryAt := make([]time.Time, of) + verified := make([]bool, execCount) + attempts := make([]int, execCount) + nextRetryAt := make([]time.Time, execCount) doneCount := 0 shards := newShardVerifyTracker(execs, req.Resume, req.NoProgressByShard) @@ -123,7 +123,7 @@ func (a *activities) runVerifyPhase( case <-activity.GetWorkerStopChannel(ctx): return replicateBatchResult{ CompletedShards: shards.allCompleted(), - InFlight: a.buildInFlight(execs, verified, shards, time.Now()), + InFlight: buildInFlight(execs, verified, shards, time.Now()), VerifiedCount: int64(doneCount), }, nil default: @@ -160,7 +160,7 @@ func (a *activities) runVerifyPhase( activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) if done, result, err := a.evaluateVerifyIteration( - ctx, req, execs, verified, shards, doneCount, of, draining, drainStartAt); err != nil { + ctx, req, execs, verified, shards, doneCount, execCount, draining, drainStartAt); err != nil { return replicateBatchResult{}, wrapBatchVerifyError(err, int64(doneCount)) } else if done { return result, nil @@ -173,7 +173,7 @@ func (a *activities) runVerifyPhase( continue } - a.waitNextTick(ctx, callCtx, minNextRetry, draining, drainStartAt, req.DrainGrace) + waitNextTick(ctx, callCtx, minNextRetry, draining, drainStartAt, req.DrainGrace) } } @@ -187,12 +187,12 @@ func (a *activities) evaluateVerifyIteration( execs []*shardedExecutionInfo, verified []bool, shards shardVerifyTracker, - doneCount, of int, + doneCount, execCount int, draining bool, drainStartAt time.Time, ) (bool, replicateBatchResult, error) { // Clean completion — every exec verified. - if doneCount >= of { + if doneCount >= execCount { return true, replicateBatchResult{ CompletedShards: shards.allCompleted(), VerifiedCount: int64(doneCount), @@ -200,7 +200,7 @@ func (a *activities) evaluateVerifyIteration( } // Per-shard cumulative no-progress backstop. - if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, of); sErr != nil { + if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, execCount); sErr != nil { return false, replicateBatchResult{}, sErr } @@ -209,10 +209,10 @@ func (a *activities) evaluateVerifyIteration( // carries everything the workflow needs (completed shards + // unverified execs grouped by shard with their cumulative // no-progress duration). - if a.shouldExitDrain(req, shards, drainStartAt) { + if shouldExitDrain(req, shards, drainStartAt) { return true, replicateBatchResult{ CompletedShards: shards.allCompleted(), - InFlight: a.buildInFlight(execs, verified, shards, time.Now()), + InFlight: buildInFlight(execs, verified, shards, time.Now()), VerifiedCount: int64(doneCount), }, nil } @@ -338,7 +338,7 @@ func (a *activities) checkStuckShard( return nil } msg := fmt.Sprintf("shard %d no progress for %v", stuckShard, stuckDur) - if stuckIdx, found := a.firstUnverifiedOnShard(execs, verified, stuckShard); found { + if stuckIdx, found := firstUnverifiedOnShard(execs, verified, stuckShard); found { stuck := execs[stuckIdx] msg = fmt.Sprintf("shard %d no progress for %v on %s/%s (%d/%d done)", stuckShard, stuckDur, stuck.BusinessID, stuck.RunID, doneCount, total) @@ -349,7 +349,7 @@ func (a *activities) checkStuckShard( // shouldExitDrain reports whether the drain-mode exit conditions are // met: either the grace window has expired, or the cumulative idle cost // across completed-but-unsignaled shards crossed the threshold. -func (a *activities) shouldExitDrain(req *shardedBatchReq, shards shardVerifyTracker, drainStartAt time.Time) bool { +func shouldExitDrain(req *shardedBatchReq, shards shardVerifyTracker, drainStartAt time.Time) bool { if time.Since(drainStartAt) >= req.DrainGrace { return true } @@ -390,7 +390,7 @@ func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchRe // DrainGrace remaining when in drain mode. In drain mode the parent ctx // is already dead, so we wake on the detached drain ctx instead — using // the parent ctx would tight-loop on its Done channel. -func (a *activities) waitNextTick( +func waitNextTick( ctx, callCtx context.Context, minNextRetry time.Time, draining bool, @@ -677,7 +677,7 @@ func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (i // attaches the cumulative no-progress duration per shard, for the // drain-mode activity return. Shards with zero unverified execs are // reported via CompletedShards instead. -func (a *activities) buildInFlight( +func buildInFlight( execs []*shardedExecutionInfo, verified []bool, shards shardVerifyTracker, @@ -720,7 +720,7 @@ func (a *activities) buildInFlight( // yet, and a found flag. Callers should only invoke this for shards // with at least one pending exec; the found=false return is a defensive // fallback so a tracker / verified-slice drift can't crash the activity. -func (a *activities) firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shard int32) (int, bool) { +func firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shard int32) (int, bool) { for i, ex := range execs { if verified[i] { continue From 9b77454fd233737ef9c1a5ad60426af781feaeb1 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 16:00:17 +0100 Subject: [PATCH 11/35] Avoid activity heartbeat timeout on long verify passes. --- service/worker/migration/sharded_activity.go | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index 42f1837aebe..8f28929cce1 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -146,7 +146,7 @@ func (a *activities) runVerifyPhase( } passDelta, minNextRetry, ctxAborted, vErr := a.runVerifyPass( - callCtx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) + ctx, callCtx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) // Fold partial progress in before the error check — the SDK // discards the activity result on failure, so the only way // the workflow learns about partially-verified execs on the @@ -265,7 +265,13 @@ func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, e // Returns a non-nil error only for hard errors from the verify path; // ctx-derived errors set ctxAborted instead so the outer loop owns // the decision about what to do next. +// +// ctx is the activity ctx, used only for heartbeating — a single pass +// over a large batch can outlast HeartbeatTimeout if we only heartbeat +// once at the end, so we tick per attempted exec. callCtx is what the +// DMS call rides on (the detached drain ctx in drain mode). func (a *activities) runVerifyPass( + ctx context.Context, callCtx context.Context, remoteAdminClient adminservice.AdminServiceClient, ns *namespace.Namespace, @@ -304,12 +310,13 @@ func (a *activities) runVerifyPass( verified[i] = true verifiedDelta++ shards.recordVerified(ex.Shard, time.Now()) - continue + } else { + attempts[i]++ + nextRetryAt[i] = time.Now().Add(backoffDelay(attempts[i])) + minNextRetry = earliest(minNextRetry, nextRetryAt[i]) } - attempts[i]++ - nextRetryAt[i] = time.Now().Add(backoffDelay(attempts[i])) - minNextRetry = earliest(minNextRetry, nextRetryAt[i]) + activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) } return verifiedDelta, minNextRetry, false, nil } From bc4cec0030543fc786602bdba79834346fbfba21 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 5 Jun 2026 16:14:24 +0100 Subject: [PATCH 12/35] Correct an exit flow and catch some edge cases. --- service/worker/migration/sharded_activity.go | 15 +++++++++------ service/worker/migration/sharded_workflow.go | 12 +++++++++--- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activity.go index 8f28929cce1..336bfa60db1 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activity.go @@ -199,16 +199,14 @@ func (a *activities) evaluateVerifyIteration( }, nil } - // Per-shard cumulative no-progress backstop. - if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, execCount); sErr != nil { - return false, replicateBatchResult{}, sErr - } - if draining { // Drain-mode exit checks. No signals here — the return value // carries everything the workflow needs (completed shards + // unverified execs grouped by shard with their cumulative - // no-progress duration). + // no-progress duration). The per-shard no-progress backstop is + // deliberately skipped: drain is bounded by DrainGrace and the + // outstanding execs need to flow back via InFlight for CAN + // carry-over, not surface as a ShardNoProgress failure. if shouldExitDrain(req, shards, drainStartAt) { return true, replicateBatchResult{ CompletedShards: shards.allCompleted(), @@ -219,6 +217,11 @@ func (a *activities) evaluateVerifyIteration( return false, replicateBatchResult{}, nil } + // Per-shard cumulative no-progress backstop. + if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, execCount); sErr != nil { + return false, replicateBatchResult{}, sErr + } + if err := a.maybeSignalRelease(ctx, req, shards); err != nil { return false, replicateBatchResult{}, err } diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 2ef230c4b69..5e856878c7d 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -290,6 +290,11 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati }).Get(ctx, &targetMd); err != nil { return nil, err } + if targetMd.ShardCount <= 0 { + return nil, temporal.NewNonRetryableApplicationError( + fmt.Sprintf("DescribeTargetCluster returned non-positive ShardCount (%d) for target %q", targetMd.ShardCount, params.TargetClusterName), + "InvalidTargetShardCount", nil) + } if params.BatchSize <= 0 { params.BatchSize = defaultBatchSize } @@ -960,10 +965,11 @@ func (s *shardedWorkflowState) drainForCAN(ctx workflow.Context) { } s.cancelActivities() releaseCh := workflow.GetSignalChannel(ctx, releaseShardsSignalName) + // Wait unconditionally for pendingDispatches to drain — lastErr may + // already be set on entry, but drainPayload, batchExecs, and status + // recovery fields only finalise once every in-flight goroutine has + // returned. _ = workflow.Await(ctx, func() bool { - if s.lastErr != nil { - return true - } return s.pendingDispatches == 0 && releaseCh.Len() == 0 }) } From 9031ae836977b239b9d79797248bbbcfbd19ff97 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 12:14:30 +0100 Subject: [PATCH 13/35] Use cluster metadata for target shard count. Use local activities for metadata/shard count as they are in-memory calls. --- service/worker/migration/activities.go | 27 ++------------- service/worker/migration/fx.go | 1 + ...rded_activity.go => sharded_activities.go} | 33 +++++++++++++++++++ ...ity_test.go => sharded_activities_test.go} | 0 service/worker/migration/sharded_workflow.go | 18 ++++------ 5 files changed, 42 insertions(+), 37 deletions(-) rename service/worker/migration/{sharded_activity.go => sharded_activities.go} (95%) rename service/worker/migration/{sharded_activity_test.go => sharded_activities_test.go} (100%) diff --git a/service/worker/migration/activities.go b/service/worker/migration/activities.go index e609bf65d32..451fda0864a 100644 --- a/service/worker/migration/activities.go +++ b/service/worker/migration/activities.go @@ -22,6 +22,7 @@ import ( chasmactivity "go.temporal.io/server/chasm/lib/activity" serverClient "go.temporal.io/server/client" "go.temporal.io/server/common" + "go.temporal.io/server/common/cluster" "go.temporal.io/server/common/dynamicconfig" "go.temporal.io/server/common/headers" "go.temporal.io/server/common/log" @@ -92,14 +93,6 @@ type ( NamespaceID string } - DescribeTargetClusterRequest struct { - TargetClusterName string - } - - DescribeTargetClusterResponse struct { - ShardCount int32 - } - ReplicationStatus struct { MaxReplicationTaskIds map[int32]int64 } @@ -129,6 +122,7 @@ type ( adminClient adminservice.AdminServiceClient clientFactory serverClient.Factory clientBean serverClient.Bean + clusterMetadata cluster.Metadata Logger log.Logger MetricsHandler metrics.Handler forceReplicationMetricsHandler metrics.Handler @@ -201,23 +195,6 @@ func (a *activities) GetMetadata(_ context.Context, request MetadataRequest) (*M }, nil } -// DescribeTargetCluster fetches the remote cluster's history shard count via -// its admin DescribeCluster RPC. The remote must be registered with the -// source cluster's cluster metadata (the cluster name doubles as the -// adminClient cache key) — which is already a prerequisite for force -// replication, since the source generates replication tasks against it. -func (a *activities) DescribeTargetCluster(ctx context.Context, req DescribeTargetClusterRequest) (*DescribeTargetClusterResponse, error) { - remoteAdminClient, err := a.clientBean.GetRemoteAdminClient(req.TargetClusterName) - if err != nil { - return nil, err - } - resp, err := remoteAdminClient.DescribeCluster(ctx, &adminservice.DescribeClusterRequest{}) - if err != nil { - return nil, err - } - return &DescribeTargetClusterResponse{ShardCount: resp.GetHistoryShardCount()}, nil -} - // GetMaxReplicationTaskIDs returns max replication task id per shard func (a *activities) GetMaxReplicationTaskIDs(ctx context.Context) (*ReplicationStatus, error) { ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index b92ca6ca85a..e06dbac623c 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -228,6 +228,7 @@ func newActivitiesFromParams(params initParams, workflowTypeName string) (*activ adminClient: localAdmin, clientFactory: params.ClientFactory, clientBean: params.ClientBean, + clusterMetadata: params.ClusterMetadata, namespaceReplicationQueue: params.NamespaceReplicationQueue, taskManager: params.TaskManager, Logger: params.Logger, diff --git a/service/worker/migration/sharded_activity.go b/service/worker/migration/sharded_activities.go similarity index 95% rename from service/worker/migration/sharded_activity.go rename to service/worker/migration/sharded_activities.go index 336bfa60db1..8ba6cb24d26 100644 --- a/service/worker/migration/sharded_activity.go +++ b/service/worker/migration/sharded_activities.go @@ -21,6 +21,39 @@ import ( "go.temporal.io/server/common/quotas" ) +type ( + DescribeTargetClusterRequest struct { + TargetClusterName string + } + + DescribeTargetClusterResponse struct { + ShardCount int32 + } +) + +// DescribeTargetCluster returns the remote cluster's history shard count +// from the locally-cached cluster_metadata table. The remote is registered +// via AdminService.AddOrUpdateRemoteCluster (which fetches HistoryShardCount +// from the remote at registration time and stores it) — a prerequisite for +// force replication, since the source generates replication tasks against +// it. ClusterMetadata refreshes the cache every minute, so the value can be +// at most that stale; shard count never changes for a live cluster so this +// is fine. +func (a *activities) DescribeTargetCluster(_ context.Context, req DescribeTargetClusterRequest) (*DescribeTargetClusterResponse, error) { + info, ok := a.clusterMetadata.GetAllClusterInfo()[req.TargetClusterName] + if !ok { + return nil, temporal.NewNonRetryableApplicationError( + fmt.Sprintf("target cluster %q not registered in cluster metadata", req.TargetClusterName), + "TargetClusterNotRegistered", nil) + } + if info.ShardCount <= 0 { + return nil, temporal.NewNonRetryableApplicationError( + fmt.Sprintf("target cluster %q has non-positive ShardCount (%d) in cluster metadata", req.TargetClusterName, info.ShardCount), + "InvalidTargetShardCount", nil) + } + return &DescribeTargetClusterResponse{ShardCount: info.ShardCount}, nil +} + // ReplicateBatch is the per-batch activity body for the sharded force // replication workflow. Runs inject (skipped on Resume) then verify, // signal-releasing completed shards mid-flight as their cumulative diff --git a/service/worker/migration/sharded_activity_test.go b/service/worker/migration/sharded_activities_test.go similarity index 100% rename from service/worker/migration/sharded_activity_test.go rename to service/worker/migration/sharded_activities_test.go diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 5e856878c7d..18ec0c86af8 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -167,7 +167,7 @@ func maybeKickoffShardedTaskQueueUserDataReplication(ctx workflow.Context, param // status query's progress reporting. func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedForceReplicationParams) (int64, error) { ao := workflow.ActivityOptions{ - StartToCloseTimeout: 2 * time.Minute, + StartToCloseTimeout: 30 * time.Second, RetryPolicy: forceReplicationActivityRetryPolicy, } var a *activities @@ -273,28 +273,22 @@ type shardedWorkflowState struct { } func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicationParams) (*shardedWorkflowState, error) { - ao := workflow.ActivityOptions{ - StartToCloseTimeout: 24 * time.Hour, - HeartbeatTimeout: time.Minute, + lao := workflow.LocalActivityOptions{ + StartToCloseTimeout: 1 * time.Second, RetryPolicy: forceReplicationActivityRetryPolicy, } - metaCtx := workflow.WithActivityOptions(ctx, ao) + localCtx := workflow.WithLocalActivityOptions(ctx, lao) var a *activities var md MetadataResponse - if err := workflow.ExecuteActivity(metaCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { + if err := workflow.ExecuteLocalActivity(localCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { return nil, err } var targetMd DescribeTargetClusterResponse - if err := workflow.ExecuteActivity(metaCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ + if err := workflow.ExecuteLocalActivity(localCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ TargetClusterName: params.TargetClusterName, }).Get(ctx, &targetMd); err != nil { return nil, err } - if targetMd.ShardCount <= 0 { - return nil, temporal.NewNonRetryableApplicationError( - fmt.Sprintf("DescribeTargetCluster returned non-positive ShardCount (%d) for target %q", targetMd.ShardCount, params.TargetClusterName), - "InvalidTargetShardCount", nil) - } if params.BatchSize <= 0 { params.BatchSize = defaultBatchSize } From ed18f0885b2494191af583602085d1342f882cb3 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 12:37:42 +0100 Subject: [PATCH 14/35] Register local activities. --- service/worker/migration/fx.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index e06dbac623c..40cecd13c4a 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -142,6 +142,15 @@ func (sc *shardedWorkerComponent) RegisterWorkflow(registry sdkworker.Registry) registry.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{ Name: forceTaskQueueUserDataReplicationWorkflow, }) + // Local activities dispatch from the workflow worker's own registry, so + // the ones invoked via ExecuteLocalActivity (GetMetadata, + // DescribeTargetCluster) need to be visible here too. Registering the + // whole *activities set mirrors the activity-worker registration; the + // workflow worker has LocalActivityWorkerOnly=true so this doesn't + // race the activity worker for regular activity tasks. + registry.RegisterActivityWithOptions(sc.activities, activity.RegisterOptions{ + DisableAlreadyRegisteredCheck: true, + }) } func (sc *shardedWorkerComponent) DedicatedWorkflowWorkerOptions() *workercommon.DedicatedWorkerOptions { From 6a27f091a96ba5cc3fdd66d3d1ad6a0b6b1a20ac Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 14:35:23 +0100 Subject: [PATCH 15/35] Improve behaviour during worker shutdown by retrying. --- .../worker/migration/sharded_activities.go | 26 +++++++++---------- service/worker/migration/sharded_workflow.go | 8 +++++- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 8ba6cb24d26..dad815da593 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -144,21 +144,21 @@ func (a *activities) runVerifyPhase( defer drainCancel() for { - // Worker shutdown short-circuits drain mode entirely. The SDK - // closes WorkerStopChannel WorkerStopTimeout before forcibly - // returning; burning that window on DMS calls that can't drive - // their results back is worse than returning current state and - // letting ResumeShards / RecoveredBuckets recover next cycle. - // Re-checked each iteration because shutdown can fire after - // we've already entered drain — the detached ctx wouldn't - // notice on its own. + // Worker shutdown short-circuits with a retryable error so the + // SDK reschedules on another worker. Folding partial state into + // the workflow's drain bucket would conflate worker shutdown + // (deploys, host loss — orthogonal to migration progress) with + // drain-for-CAN. Returning here also avoids the ~HeartbeatTimeout + // wait that the alternative (silent worker death) would incur + // before the server retries the attempt. Inject is already + // heartbeat-preserved (InjectDone), so retry skips it; verify + // re-runs from scratch but DMS reads are idempotent. select { case <-activity.GetWorkerStopChannel(ctx): - return replicateBatchResult{ - CompletedShards: shards.allCompleted(), - InFlight: buildInFlight(execs, verified, shards, time.Now()), - VerifiedCount: int64(doneCount), - }, nil + return replicateBatchResult{}, temporal.NewApplicationErrorWithOptions( + "worker shutdown", "WorkerShutdown", + temporal.ApplicationErrorOptions{NextRetryDelay: time.Second}, + ) default: } diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 18ec0c86af8..f05451a1d2c 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -449,8 +449,14 @@ var ( // lets a transient activity failure recover via heartbeat-resume // without losing inject progress. WaitForCancellation lets a // cancelled activity run drain logic and return its drain result. + // + // 10 attempts is sized for fleet rollouts: a rolling deploy of the + // activity workers can burn several attempts per batch (each + // shutdown surfaces as a retryable WorkerShutdown error from the + // activity). 3 was tight enough that two unlucky deploys could + // exhaust the budget on a long-running CAN cycle. shardedReplicateBatchRetryPolicy = &temporal.RetryPolicy{ - MaximumAttempts: 3, + MaximumAttempts: 10, } shardedReplicateBatchActivityOptions = workflow.ActivityOptions{ StartToCloseTimeout: 24 * time.Hour, From 5af7995f215a341399c3421962b13f276c2c805f Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 15:09:22 +0100 Subject: [PATCH 16/35] Don't fetch adminClient until we need it. Highly unlikely to matter, but does allow us to make any progress we can without it. --- service/worker/migration/sharded_activities.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index dad815da593..47e600a7795 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -71,11 +71,6 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, nil } - remoteAdminClient, err := a.clientBean.GetRemoteAdminClient(req.TargetClusterName) - if err != nil { - return replicateBatchResult{}, fmt.Errorf("get remote admin client for %s: %w", req.TargetClusterName, err) - } - var hb replicateBatchHeartbeat if activity.HasHeartbeatDetails(ctx) { _ = activity.GetHeartbeatDetails(ctx, &hb) @@ -107,6 +102,11 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, fmt.Errorf("look up namespace %s: %w", req.NamespaceID, err) } + remoteAdminClient, err := a.clientBean.GetRemoteAdminClient(req.TargetClusterName) + if err != nil { + return replicateBatchResult{}, fmt.Errorf("get remote admin client for %s: %w", req.TargetClusterName, err) + } + return a.runVerifyPhase(ctx, req, execs, execCount, remoteAdminClient, ns) } From 9b1c682503d888ede468b5033d2f1c16cc0a3266 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 15:34:24 +0100 Subject: [PATCH 17/35] Fix a test-only data race. --- service/worker/migration/sharded_workflow_test.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index bb7e852b92c..33da659a648 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -385,7 +385,10 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) }, activity.RegisterOptions{Name: "ListWorkflows"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - env.SetContinueAsNewSuggested(true) + // Hop to the main loop goroutine to flip CAN-suggested; the + // activity goroutine writing workflowInfo directly would race + // the workflow coroutine reading it at the top of its page loop. + env.RegisterDelayedCallback(func() { env.SetContinueAsNewSuggested(true) }, 0) return replicateBatchResult{ InFlight: []ResumeShard{{ Shard: drainedShard, @@ -438,8 +441,10 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { } pageServed = true // Trigger CAN as soon as this page is served so drainForCAN - // runs before any dispatched batches can complete. - env.SetContinueAsNewSuggested(true) + // runs before any dispatched batches can complete. Hop to the + // main loop goroutine — writing workflowInfo directly from the + // activity goroutine would race the workflow coroutine's read. + env.RegisterDelayedCallback(func() { env.SetContinueAsNewSuggested(true) }, 0) return &listWorkflowsResponse{ Executions: execs, NextPageToken: []byte("more"), From 93e909c49046337dfeca013dd8dca28ced9b359e Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 17:59:54 +0100 Subject: [PATCH 18/35] Comment fixes. --- service/worker/migration/sharded_workflow.go | 4 ++-- service/worker/migration/sharded_workflow_test.go | 5 +---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index f05451a1d2c..cf29b2419f9 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -215,8 +215,8 @@ type shardedWorkflowState struct { // shardInFlight is the per-shard exclusivity set: a shard's // entry is set when it's part of any in-flight batch and // cleared when that batch returns (either fully or via mid-flight - // signal-release). Concurrent batches are limited only by this - // set — there is no global slot cap. + // signal-release). The packer uses this set to ensure that each + // shard can only be present in one in-flight batch at a time. shardInFlight map[int32]bool // heldByBatch tracks per-batch shard ownership. spawnBatch diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 33da659a648..efb6759d463 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -601,7 +601,7 @@ func TestSharded_ListWorkflowsError(t *testing.T) { } // TestSharded_ReplicateBatchRetryableError: when ReplicateBatch returns -// a retryable error, the workflow exhausts its 3-attempt retry policy +// a retryable error, the workflow exhausts its configured retry policy // and surfaces the error as lastErr. Mirrors the existing // TestGenerateReplicationTaskRetryableError. func TestSharded_ReplicateBatchRetryableError(t *testing.T) { @@ -626,9 +626,6 @@ func TestSharded_ReplicateBatchRetryableError(t *testing.T) { err := env.GetWorkflowError() require.Error(t, err) require.Contains(t, err.Error(), "transient backend error") - // MaximumAttempts: 3 in spawnBatch's activity options — assert at - // least 2 retries actually happened so a future change that drops - // the retry policy fails this test. require.GreaterOrEqual(t, attempts.Load(), int32(2), "expected ReplicateBatch to be retried at least twice before failing") } From 1a5bd8ae3d7b427b696e45e22a6923c6b6ea42b3 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Mon, 8 Jun 2026 19:42:38 +0100 Subject: [PATCH 19/35] Avoid edge case where we might start listing again. --- service/worker/migration/sharded_workflow.go | 17 ++++++- .../worker/migration/sharded_workflow_test.go | 51 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index cf29b2419f9..1c8afb3c493 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -376,7 +376,14 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { // the SDK signals that history is large enough to CAN. Errors // here latch into lastErr and fall through to the unified exit // funnel — same drain-and-decide path as activity-driven errors. - for !workflow.GetInfo(ctx).GetContinueAsNewSuggested() { + // + // On a CAN cycle (ContinuedAsNewCount > 0), an empty NextPageToken + // means a prior cycle already exhausted pagination — empty token at + // the start of cycle 0 is "haven't started", but at the start of any + // later cycle it's "finished". Skip listing in that case; otherwise + // a carry-over-driven CAN (drained InFlight, recovered buckets, …) + // would re-list page 1 and re-enqueue every visible execution. + for s.shouldList() && !workflow.GetInfo(ctx).GetContinueAsNewSuggested() { if s.lastErr != nil { break } @@ -434,6 +441,14 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) } +// shouldList returns true if the page loop has more work to do this +// cycle. Cycle 0 always lists (NextPageToken is empty either way). +// On later cycles, an empty NextPageToken can only mean a prior cycle +// already drained pagination — there's nothing left to list. +func (s *shardedWorkflowState) shouldList() bool { + return s.params.ContinuedAsNewCount == 0 || len(s.params.NextPageToken) > 0 +} + var ( shardedListWorkflowsRetryPolicy = &temporal.RetryPolicy{ InitialInterval: time.Second, diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index efb6759d463..c329140a4b8 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -422,6 +422,57 @@ func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) require.Equal(t, 42*time.Second, nextParams.ResumeShards[0].NoProgressDuration) } +// TestSharded_ListingDoneAcrossCAN_SkipsPageLoop pins down that a +// post-listing CAN cycle does not re-enter ListWorkflows. The prior +// cycle finished pagination (NextPageToken empty) but still had carry- +// over (e.g. InFlight from a returned activity), so it CAN'd. Without +// the ContinuedAsNewCount > 0 guard in the page loop, the next cycle +// can't distinguish "haven't started listing" from "finished listing" +// and would re-list page 1, double-enqueueing every visible execution. +func TestSharded_ListingDoneAcrossCAN_SkipsPageLoop(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, 4) + + var listCalls atomic.Int32 + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + listCalls.Add(1) + return &listWorkflowsResponse{}, nil + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + // Resume-only batches; no fresh inject work. + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + require.True(t, req.Resume, "this cycle should only dispatch resume batches; no listing should occur") + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + // Mimic a CAN-cycle entry: prior cycle exhausted pagination + // (NextPageToken empty) and CAN'd because it had ResumeShards to + // carry over. ContinuedAsNewCount > 0 is the signal that empty + // NextPageToken means "done", not "haven't started". + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + ContinuedAsNewCount: 1, + ResumeShards: []ResumeShard{{ + Shard: 0, + Execs: map[string][]RunEntry{"wf-resume": {{RunID: "run-resume"}}}, + }}, + // On a real CAN cycle the prior cycle's child kicked off in + // cycle 0, so its signal arrives once; subsequent cycles see + // Done=true in their input params and skip both the kickoff + // and the terminal Await. The test mirrors that. + TaskQueueUserDataReplicationStatus: TaskQueueUserDataReplicationStatus{Done: true}, + }) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + require.Zero(t, listCalls.Load(), "ListWorkflows must not be called on a CAN cycle that inherited an exhausted page token") +} + // TestSharded_CancelBeforeStart_NoLostExecs pins down recovery when // an activity is dispatched and the workflow CANs before the activity // body runs: the activity returns CanceledError with no result, and From b5e6fe012f948a638000a0b6156808b6fe96004b Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Tue, 9 Jun 2026 14:27:34 +0100 Subject: [PATCH 20/35] Remove stale comment. --- service/worker/migration/sharded_activities_test.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/service/worker/migration/sharded_activities_test.go b/service/worker/migration/sharded_activities_test.go index 05110158755..301b0ac6860 100644 --- a/service/worker/migration/sharded_activities_test.go +++ b/service/worker/migration/sharded_activities_test.go @@ -268,11 +268,6 @@ func (s *activitiesSuite) TestReplicateBatch_Resume_SkipsInject() { // from the activity side. func (s *activitiesSuite) TestReplicateBatch_DisableVerification() { env, _ := s.initEnv() - // No NewRemoteAdminClientWithTimeout expectation — the activity - // builds the client unconditionally even in inject-only mode, so we - // still need the factory to hand back something. Return without - // expecting any DMS calls on it. - s.mockHistoryClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), gomock.Any()). Return(&historyservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1) From 73362e508f361d456849ec9a567494677b2d15c9 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Tue, 9 Jun 2026 16:07:43 +0100 Subject: [PATCH 21/35] Simplify continue as new behaviour. --- service/worker/migration/sharded_types.go | 32 +++- service/worker/migration/sharded_workflow.go | 147 +++++++++++------- .../worker/migration/sharded_workflow_test.go | 15 +- 3 files changed, 129 insertions(+), 65 deletions(-) diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 1de0ba06ce8..667ca5602d1 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -45,9 +45,9 @@ const ( defaultShardNoProgress = 5 * time.Minute // defaultDrainGrace is the wall-budget the activity gets after - // the workflow cancels it for CAN. Continues verifying until - // either the grace expires, the idle-cost trigger fires, or - // every exec verifies. + // the workflow cancels it (on lastErr or cycle drain timeout). + // Continues verifying until either the grace expires, the + // idle-cost trigger fires, or every exec verifies. defaultDrainGrace = 15 * time.Second // defaultIdleShardCost is the cumulative idle-time threshold @@ -56,6 +56,16 @@ const ( // completed-but-not-yet-released shards mid-flight. defaultIdleShardCost = 30 * time.Second + // defaultCycleDrainTimeout bounds the wall-clock the workflow + // will spend draining buckets + awaiting in-flight activities + // after the page loop stops. GetContinueAsNewSuggested trips at + // ~8% of the hard history cap so we have ~92% of the budget + // remaining when the page loop breaks; 10 minutes is generously + // inside that. Catches the "many shards making slow-but-real + // progress" case; a single stuck shard is already bounded by + // ShardNoProgress on the activity side. + defaultCycleDrainTimeout = 10 * time.Minute + // defaultPerBatchGenerateRPS is the per-batch inject-phase target. // Sharded dispatches many concurrent batches and each builds its // own limiter, so this caps the per-batch generate-replication-task @@ -223,6 +233,14 @@ type ShardedForceReplicationParams struct { DrainGrace time.Duration IdleShardCost time.Duration + // CycleDrainTimeout caps how long the workflow will spend after + // the page loop stops, draining buckets and waiting for in-flight + // activities to complete naturally. On expiry the workflow falls + // into drainForCAN — cancels in-flight batches, collects their + // drain payload, and CANs with the recovered state. Defaults to + // defaultCycleDrainTimeout. + CycleDrainTimeout time.Duration + TaskQueueUserDataReplicationParams TaskQueueUserDataReplicationParams // PerBatchGenerateRPS is the inject-phase rate-limiter target inside @@ -264,9 +282,11 @@ type ShardedForceReplicationParams struct { // RecoveredBuckets carries execs whose dispatching activity returned // a cancellation without returning a result — i.e., the activity - // body never ran (cancel-before-start race). They were dispatched - // but never injected, so the new cycle restores them into the - // streaming buckets to be dispatched as fresh inject+verify batches. + // body never ran, because cancellation (from lastErr or cycle drain + // timeout) reached it before the worker picked it up. They were + // dispatched but never injected, so the new cycle restores them + // into the streaming buckets to be dispatched as fresh inject+verify + // batches. RecoveredBuckets BatchPayload TaskQueueUserDataReplicationStatus TaskQueueUserDataReplicationStatus diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 1c8afb3c493..ab293eff203 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -24,11 +24,17 @@ import ( // eligibility. At cycle end, every remaining bucket flushes as // packed activities. // -// If there are more pages, in-flight activities are cancelled (giving -// them DrainGrace to drain), their drain payload arrives via the -// activity return value, and the workflow CANs with NextPageToken + -// ResumeShards in the carry-over. Otherwise it waits for activities -// to finish naturally and returns nil. +// In-flight activities are allowed to finish naturally at cycle end; +// CycleDrainTimeout bounds the wait as a safety net against +// pathological slow drains. On timer expiry in-flights are cancelled +// and their drained execs feed the next cycle's carry-over as +// ResumeShards. On lastErr the same drain happens, but the workflow +// returns the error rather than CANing — the drained state surfaces +// only via the status query's recovery bundle. +// +// If listing wasn't exhausted, the workflow CANs with NextPageToken +// (plus any drained state from the timer-fired path). Otherwise it +// returns nil. func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { // Page token at workflow entry — returned by the status query as // PageTokenForRestart so tooling that already knows the legacy @@ -241,10 +247,12 @@ type shardedWorkflowState struct { // batchExecs tracks the input payload of each in-flight batch. // Cleared on any nil-error return (drained execs are folded // into drainPayload from the activity result; cleanly completed - // batches return an empty InFlight). Anything left at CAN time + // batches return an empty InFlight). Anything left at cycle end // corresponds to a batch whose activity returned CanceledError - // with no result — i.e. the activity body never ran. Those - // execs are recovered into the next cycle's streaming buckets. + // with no result — i.e. the activity body never ran. On the CAN + // path those execs are recovered into the next cycle's streaming + // buckets; on the lastErr return path they surface via the status + // query's recovery bundle but aren't auto-dispatched. batchExecs map[int64]BatchPayload // pendingDispatches counts spawned dispatch coroutines that @@ -264,6 +272,13 @@ type shardedWorkflowState struct { // failure. lastErr error + // cycleDrainTimedOut is set by the safety timer started after the + // page loop. drainBuckets and awaitInFlightCompletion both honour + // it so a pathological slow drain can't eat into the workflow's + // remaining history budget — on expiry the workflow falls into + // drainForCAN and ships unfinished work as carry-over. + cycleDrainTimedOut bool + nextBatchID int64 // metricsHandler is tagged with the workflow's fixed scope + @@ -304,6 +319,9 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati if params.IdleShardCost <= 0 { params.IdleShardCost = defaultIdleShardCost } + if params.CycleDrainTimeout <= 0 { + params.CycleDrainTimeout = defaultCycleDrainTimeout + } if params.ListWorkflowsPageSize <= 0 { params.ListWorkflowsPageSize = defaultShardedListPageSize } @@ -409,6 +427,11 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { } } + // Start the cycle drain safety timer; see defaultCycleDrainTimeout + // for the budget rationale. + cancelCycleTimer := s.startCycleDrainTimer(ctx) + defer cancelCycleTimer() + // Drain remaining buckets only when no error has latched — on // error we deliberately stop scheduling new work and let the // already-dispatched batches finish via awaitInFlightCompletion. @@ -417,10 +440,10 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { } // Wait for in-flight activities. Cancels them when we already - // know we're failing or CAN-ing; on the clean success path - // (no error, no more pages) waits naturally so a healthy - // activity isn't cancelled into a CanceledError that masquerades - // as carry-over state. + // know we're failing or the cycle drain timer has expired; + // otherwise waits naturally so a healthy activity isn't + // cancelled into a CanceledError that masquerades as carry-over + // state. s.awaitInFlightCompletion(ctx) // Single exit decision. Recovery state has the same shape on @@ -507,28 +530,39 @@ func (s *shardedWorkflowState) replicateBatch(ctx, activityParentCtx workflow.Co return result, nil } -// awaitInFlightCompletion drains in-flight batches before the workflow -// exits. The strategy depends on what we already know: -// -// - Error latched or more pages remain (we're going to fail or CAN -// either way): cancel immediately via drainForCAN, bounded by -// DrainGrace + IdleShardCost. No point waiting for activities -// that are going to be discarded. -// - Clean success path (no error, no more pages): wait for natural -// completion so a healthy activity's clean result isn't masked -// as a CanceledError. If an activity hits its ShardNoProgress -// backstop mid-wait, latch the error then cancel the rest fast -// rather than waiting on every shard's backstop too. +// startCycleDrainTimer spawns a coroutine that sets cycleDrainTimedOut +// after CycleDrainTimeout. Returned cancel func stops the timer on +// natural exit. The flag is read by drainBuckets (stops spawning new +// batches and exits) and awaitInFlightCompletion (exits its Await and +// calls drainForCAN to cancel in-flight activities and collect their +// drain payload for the next cycle's carry-over). +func (s *shardedWorkflowState) startCycleDrainTimer(ctx workflow.Context) workflow.CancelFunc { + timerCtx, cancel := workflow.WithCancel(ctx) + workflow.Go(ctx, func(gCtx workflow.Context) { + if err := workflow.NewTimer(timerCtx, s.params.CycleDrainTimeout).Get(gCtx, nil); err == nil { + s.cycleDrainTimedOut = true + } + }) + return cancel +} + +// awaitInFlightCompletion waits for in-flight batches to complete. +// On the clean path it just blocks until pendingDispatches drops to +// zero — a healthy activity's result isn't masked as a CanceledError. +// On lastErr or cycle-drain timeout it falls into drainForCAN, which +// cancels the in-flight activities and collects their drain payload. +// run() then either returns lastErr (drain payload surfaces only via +// the status query) or CANs with the drained state as carry-over. func (s *shardedWorkflowState) awaitInFlightCompletion(ctx workflow.Context) { if s.pendingDispatches == 0 { return } - if s.lastErr != nil || len(s.params.NextPageToken) > 0 { + if s.lastErr != nil { s.drainForCAN(ctx) return } _ = workflow.Await(ctx, func() bool { - return s.pendingDispatches == 0 || s.lastErr != nil + return s.pendingDispatches == 0 || s.lastErr != nil || s.cycleDrainTimedOut }) if s.pendingDispatches > 0 { s.drainForCAN(ctx) @@ -577,7 +611,8 @@ func (s *shardedWorkflowState) collectResumeShardsForCarryover() []ResumeShard { // batches that returned CanceledError without running a body, and // listed-but-unpacked execs still sitting in s.buckets when the // workflow exited (either lastErr stopped the streaming packer or -// drainBuckets bailed out on lastErr partway through). +// drainBuckets bailed on lastErr / cycle drain timeout partway +// through). func (s *shardedWorkflowState) collectRecoveredBucketsForCarryover() BatchPayload { out := collectRecoveredBuckets(s.batchExecs) if !s.bucketsEmpty() { @@ -891,18 +926,20 @@ func runCount(byBID map[string][]RunEntry) int { return n } -// drainBuckets blocks until buckets are empty (success) or lastErr -// trips (failure). Each pass packs everything currently dispatchable, -// then awaits any change in pendingDispatches + shardInFlight so the -// next pass can attempt shards just freed by signal-release. +// drainBuckets blocks until buckets are empty (success), lastErr +// trips (failure), or the cycle drain timer expires (CAN with +// leftover buckets). Each pass packs everything currently +// dispatchable, then awaits any change in pendingDispatches + +// shardInFlight so the next pass can attempt shards just freed by +// signal-release. func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { for { - if s.lastErr != nil { + if s.lastErr != nil || s.cycleDrainTimedOut { return } for s.tryPackStreaming(ctx, true) { //nolint:revive } - if s.bucketsEmpty() || s.lastErr != nil { + if s.bucketsEmpty() || s.lastErr != nil || s.cycleDrainTimedOut { return } currentPending := s.pendingDispatches @@ -930,14 +967,14 @@ func (s *shardedWorkflowState) failDrainBucketsStuck() { } // drainBucketsAwaitPredicate returns true when the drainBuckets loop -// should wake up: lastErr tripped, a dispatch slot just freed, or a -// new free shard is ready to pack. A "free shard" wake-up only counts -// when there's also a dispatch slot to use it, otherwise the outer -// loop would busy-spin on tryPackStreaming returning false against the -// in-flight cap. +// should wake up: lastErr tripped, cycle drain timer fired, a +// dispatch slot just freed, or a new free shard is ready to pack. A +// "free shard" wake-up only counts when there's also a dispatch slot +// to use it, otherwise the outer loop would busy-spin on +// tryPackStreaming returning false against the in-flight cap. func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { return func() bool { - if s.lastErr != nil { + if s.lastErr != nil || s.cycleDrainTimedOut { return true } if s.pendingDispatches < currentPending { @@ -957,13 +994,15 @@ func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) fu // drainForCAN cancels every in-flight batch and waits for them to // return AND for the ReleaseShards signal channel to be drained. -// Activities honour cancellation by entering drain mode and returning -// a result whose InFlight carries their still-unverified execs; -// spawnBatch appends those entries to s.drainPayload. The signal -// channel drain is so a final ReleaseShards fired by an activity just -// before it returns doesn't get stranded mid-flight, which would -// leave shardInFlight set for shards the activity already considers -// complete. +// Called on lastErr (terminates with error; drain payload exposed via +// the status query) and on cycle drain timeout (CANs with drained +// state as carry-over). Activities honour cancellation by entering +// drain mode and returning a result whose InFlight carries their +// still-unverified execs; spawnBatch appends those entries to +// s.drainPayload. The signal channel drain is so a final ReleaseShards +// fired by an activity just before it returns doesn't get stranded +// mid-flight, which would leave shardInFlight set for shards the +// activity already considers complete. // // Channel.Len() is safe here because handleReleaseSignals has no // yield points between Receive and the next blocking Receive, so @@ -1047,8 +1086,10 @@ func (s *shardedWorkflowState) spawnBatch( if temporal.IsCanceledError(err) { // Cancel-before-start: the activity body never ran, // so no result is available. Leaving batchExecs[batchID] - // intact lets the CAN-end recovery path re-bucket the - // execs as fresh inject+verify work next cycle. + // intact lets collectRecoveredBucketsForCarryover + // re-bucket the execs as fresh inject+verify work next + // cycle (only meaningful on the CAN path; on lastErr + // they surface via the status query only). return } // Activity errored after partial verify — the SDK discards @@ -1061,11 +1102,11 @@ func (s *shardedWorkflowState) spawnBatch( return } // Activity body ran and returned cleanly — either a - // clean completion (empty InFlight) or a drained - // CAN-cancel (InFlight carries the still-unverified - // execs). CompletedShards is informational; the - // defer above clears heldByBatch + shardInFlight - // either way. + // clean completion (empty InFlight) or a drained cancel + // (InFlight carries the still-unverified execs; reached + // only on lastErr or cycle drain timeout). CompletedShards + // is informational; the defer above clears heldByBatch + + // shardInFlight either way. if len(result.InFlight) > 0 { s.drainPayload = append(s.drainPayload, result.InFlight...) } diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index c329140a4b8..4ae20c9352a 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -474,10 +474,10 @@ func TestSharded_ListingDoneAcrossCAN_SkipsPageLoop(t *testing.T) { } // TestSharded_CancelBeforeStart_NoLostExecs pins down recovery when -// an activity is dispatched and the workflow CANs before the activity -// body runs: the activity returns CanceledError with no result, and -// the recovery path re-buckets the input execs into RecoveredBuckets -// so the next cycle dispatches them as fresh inject+verify batches. +// the cycle drain timer fires before any dispatched activity can run: +// the activities return CanceledError with no result, and the recovery +// path re-buckets the input execs into RecoveredBuckets so the next +// cycle dispatches them as fresh inject+verify batches. func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() @@ -491,8 +491,8 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { return &listWorkflowsResponse{}, nil } pageServed = true - // Trigger CAN as soon as this page is served so drainForCAN - // runs before any dispatched batches can complete. Hop to the + // Trigger CAN-suggested as soon as this page is served so the + // page loop bails and the cycle drain timer starts. Hop to the // main loop goroutine — writing workflowInfo directly from the // activity goroutine would race the workflow coroutine's read. env.RegisterDelayedCallback(func() { env.SetContinueAsNewSuggested(true) }, 0) @@ -522,6 +522,9 @@ func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ Namespace: "test-ns", TargetClusterName: "remote_cluster", + // Short cycle drain timeout so the test hits the timer-fired + // cancel path quickly rather than waiting the production default. + CycleDrainTimeout: time.Millisecond, }) require.True(t, env.IsWorkflowCompleted()) From dc9f5d4a28e84651a812d8adfac8c3f55912cd43 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Wed, 10 Jun 2026 17:58:50 +0100 Subject: [PATCH 22/35] One determinism fix and some ignores for make workflow check. --- service/worker/migration/sharded_types.go | 11 +++++++++- service/worker/migration/sharded_workflow.go | 23 +++++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 667ca5602d1..6635790cd33 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -143,7 +143,9 @@ type BatchPayload map[int32]map[string][]RunEntry // totalRuns counts runs across all (shard, BID) groups. func (p BatchPayload) totalRuns() int { n := 0 + //workflowcheck:ignore (summation is order-independent) for _, byBID := range p { + //workflowcheck:ignore (summation is order-independent) for _, runs := range byBID { n += len(runs) } @@ -155,6 +157,7 @@ func (p BatchPayload) totalRuns() int { // activity-side flatten a deterministic iteration order for replays. func (p BatchPayload) sortedShards() []int32 { out := make([]int32, 0, len(p)) + //workflowcheck:ignore (output is sorted before any observable use) for sh := range p { out = append(out, sh) } @@ -184,6 +187,7 @@ func (p BatchPayload) flatten() []*shardedExecutionInfo { for _, sh := range p.sortedShards() { byBID := p[sh] bids := make([]string, 0, len(byBID)) + //workflowcheck:ignore (bids is sorted before use) for bid := range byBID { bids = append(bids, bid) } @@ -204,12 +208,17 @@ func (p BatchPayload) flatten() []*shardedExecutionInfo { return out } -// merge folds src into p. +// merge folds src into p. Callers guarantee disjoint (shard, BID) keys +// between src and any prior merges into p — in-flight batches hold +// disjoint shard claims and listed-but-unpacked buckets share no shard +// with batchExecs — so per-key appends never interleave across iterations. func (p BatchPayload) merge(src BatchPayload) { + //workflowcheck:ignore (writes are to disjoint keys; order-independent) for sh, byBID := range src { if p[sh] == nil { p[sh] = map[string][]RunEntry{} } + //workflowcheck:ignore (writes are to disjoint keys; order-independent) for bid, runs := range byBID { p[sh][bid] = append(p[sh][bid], runs...) } diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index ab293eff203..2a813d7fddd 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -360,7 +360,9 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati // the prior cycle so the streaming packer picks them up // alongside any new pages. s.buckets.merge(params.RecoveredBuckets) + //workflowcheck:ignore (per-shard sum is order-independent) for sh, byBID := range params.RecoveredBuckets { + //workflowcheck:ignore (per-shard sum is order-independent) for _, runs := range byBID { s.bucketCounts[sh] += len(runs) } @@ -663,6 +665,10 @@ func collectRecoveredBuckets(batchExecs map[int64]BatchPayload) BatchPayload { return nil } out := BatchPayload{} + // In-flight batches hold disjoint shard claims, so merging two + // batchExecs entries never targets the same (shard, BID) key — + // iteration order does not affect the final BatchPayload. + //workflowcheck:ignore (writes are to disjoint keys; order-independent) for _, bp := range batchExecs { out.merge(bp) } @@ -693,6 +699,7 @@ func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]R return nil } bids := make([]string, 0, len(byBID)) + //workflowcheck:ignore (bids is sorted before use) for bid := range byBID { bids = append(bids, bid) } @@ -867,6 +874,7 @@ func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { s.params.ResumeShards = unpackResumeBatches(batches[i:]) return } + //workflowcheck:ignore (setting a set of flags is order-independent) for sh := range batch.payload { s.shardInFlight[sh] = true } @@ -902,10 +910,13 @@ func (s *shardedWorkflowState) packResumeBatchPlan(entries []ResumeShard) []resu // unpackResumeBatches reverses packResumeBatchPlan, turning planned // batches back into a flat ResumeShard slice. Used when the dispatch // loop aborts on lastErr so the undispatched remainder can be carried -// into the recovery bundle / next CAN cycle. +// into the recovery bundle / next CAN cycle. The output is sorted by +// shard ID so the slice flowing into the CAN args is deterministic +// across replays. func unpackResumeBatches(batches []resumeBatch) []ResumeShard { var out []ResumeShard for _, b := range batches { + //workflowcheck:ignore (output is sorted below before any observable use) for sh, execs := range b.payload { out = append(out, ResumeShard{ Shard: sh, @@ -914,12 +925,16 @@ func unpackResumeBatches(batches []resumeBatch) []ResumeShard { }) } } + slices.SortFunc(out, func(a, b ResumeShard) int { + return int(a.Shard - b.Shard) + }) return out } // runCount sums runs across BIDs in a single shard's payload entry. func runCount(byBID map[string][]RunEntry) int { n := 0 + //workflowcheck:ignore (summation is order-independent) for _, runs := range byBID { n += len(runs) } @@ -958,6 +973,7 @@ func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { // through run(). func (s *shardedWorkflowState) failDrainBucketsStuck() { remaining := 0 + //workflowcheck:ignore (summation is order-independent) for _, n := range s.bucketCounts { remaining += n } @@ -983,6 +999,7 @@ func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) fu if !s.dispatchSlotAvailable() { return false } + //workflowcheck:ignore (existence check; iteration order does not affect result) for sh, n := range s.bucketCounts { if n > 0 && !s.shardInFlight[sh] { return true @@ -1062,6 +1079,7 @@ func (s *shardedWorkflowState) spawnBatch( } held := make(map[int32]bool, len(payload)) + //workflowcheck:ignore (building a set of flags is order-independent) for sh := range payload { held[sh] = true } @@ -1076,6 +1094,7 @@ func (s *shardedWorkflowState) spawnBatch( // shards have already been cleared from shardInFlight // by handleReleaseSignals and may by now belong to a // subsequent batch's claim. + //workflowcheck:ignore (deletes are commutative; order-independent) for sh := range s.heldByBatch[batchID] { delete(s.shardInFlight, sh) } @@ -1164,6 +1183,7 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool // bucketsEmpty reports whether every shard's bucket is empty. Reads // from the sidecar count map so it's O(#shards), not O(#runs). func (s *shardedWorkflowState) bucketsEmpty() bool { + //workflowcheck:ignore (existence check; iteration order does not affect result) for _, n := range s.bucketCounts { if n > 0 { return false @@ -1191,6 +1211,7 @@ func (s *shardedWorkflowState) bucketsEmpty() bool { // wall-clock. func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { out := make([]int32, 0, len(s.bucketCounts)) + //workflowcheck:ignore (output is sorted below before any observable use) for sh, n := range s.bucketCounts { if n == 0 || s.shardInFlight[sh] { continue From f49e6e65d5b05ea5f31c15b7633d242aa3972a0a Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Thu, 11 Jun 2026 16:05:06 +0100 Subject: [PATCH 23/35] Give the first verification per shard extra time. This is in case there was a pre-existing backlog before our tasks reach the front of the queue. For now, don't wire this across CAN, as draining is likely to give things extra time anyway for the most part, and it keeps CAN transfer simpler. --- .../worker/migration/sharded_activities.go | 24 ++++++++-- service/worker/migration/sharded_types.go | 5 +- .../worker/migration/sharded_types_test.go | 48 +++++++++++++++++++ 3 files changed, 72 insertions(+), 5 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 47e600a7795..16ec5fbd689 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -367,8 +367,9 @@ func earliest(cur, candidate time.Time) time.Time { } // checkStuckShard fails non-retryably if any shard has gone longer than -// req.ShardNoProgress without a verified outcome. Duration is cumulative -// across CAN cycles via tracker seeding. +// req.ShardNoProgress without a verified outcome (double that for a shard +// still awaiting its first verification — see pickStuck). Duration is +// cumulative across CAN cycles via tracker seeding. func (a *activities) checkStuckShard( req *shardedBatchReq, shards shardVerifyTracker, @@ -606,6 +607,7 @@ type shardVerify struct { doneAt time.Time // set when pending first hits zero; cleared on signal release released bool // ReleaseShards signal already sent lastProgress time.Time // wall time of the most recent verified outcome + verifiedAny bool // an exec on this shard has produced a verified outcome } type shardVerifyTracker map[int32]shardVerify @@ -625,6 +627,11 @@ func newShardVerifyTracker( for sh, sv := range t { if resume { sv.lastProgress = nowSeed.Add(-noProgressByShard[sh]) + // Resumed shards had their tasks submitted (and their initial + // first-verification grace) in a prior CAN cycle — inject is + // skipped on resume — so they continue cumulative no-progress + // tracking against the normal window, not the doubled one. + sv.verifiedAny = true } else { sv.lastProgress = nowSeed } @@ -637,6 +644,7 @@ func (t shardVerifyTracker) recordVerified(sh int32, now time.Time) { sv := t[sh] sv.pending-- sv.lastProgress = now + sv.verifiedAny = true if sv.pending == 0 { sv.doneAt = now } @@ -692,7 +700,11 @@ func (t shardVerifyTracker) allCompleted() []int32 { } // pickStuck returns (shard, age, true) for the lowest-numbered shard -// whose cumulative no-progress duration meets or exceeds threshold. +// whose cumulative no-progress duration meets or exceeds its effective +// threshold. A shard that hasn't yet produced any verified outcome gets +// double the window: the server may still be working through a backlog +// that predates our task submission. Once any exec on the shard verifies +// we expect the rest within the normal window, so the threshold reverts. func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (int32, time.Duration, bool) { var ( minShard int32 @@ -703,8 +715,12 @@ func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (i if sv.pending <= 0 { continue } + effective := threshold + if !sv.verifiedAny { + effective = 2 * threshold + } age := now.Sub(sv.lastProgress) - if age < threshold { + if age < effective { continue } if !found || sh < minShard { diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 6635790cd33..4b203491b48 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -41,7 +41,10 @@ const ( // backstop. While a shard's pending exec count is non-zero and // no exec on that shard has produced a verified outcome for this // long (carried across CAN via the resume payload), the activity - // fails non-retryably naming the stuck shard. + // fails non-retryably naming the stuck shard. A shard awaiting its + // very first verification gets double this window so the server has + // time to clear any backlog predating our task submission; it + // reverts to this value once the shard's first exec verifies. defaultShardNoProgress = 5 * time.Minute // defaultDrainGrace is the wall-budget the activity gets after diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 05c85b8c0b4..36604c360fb 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -3,6 +3,7 @@ package migration import ( "encoding/json" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -89,3 +90,50 @@ func TestBatchPayload_Flatten(t *testing.T) { require.Equal(t, "ra2", got[2].RunID) require.Equal(t, "b-z", got[3].BusinessID) } + +// TestShardVerifyTracker_FirstVerificationDoubledWindow pins the +// no-progress backstop's grace for a shard's first verified outcome: a +// shard that hasn't verified anything yet gets 2×threshold before +// pickStuck reports it (the server may still be clearing a backlog that +// predates our task submission), then reverts to the plain threshold — +// measured from the verification time — once its first exec verifies. +func TestShardVerifyTracker_FirstVerificationDoubledWindow(t *testing.T) { + const threshold = 5 * time.Minute + base := time.Unix(1700000000, 0) + + tr := shardVerifyTracker{0: {pending: 2, lastProgress: base}} + + _, _, stuck := tr.pickStuck(base.Add(threshold+time.Minute), threshold) + require.False(t, stuck, "must not trip past 1×threshold while awaiting first verification") + + sh, age, stuck := tr.pickStuck(base.Add(2*threshold), threshold) + require.True(t, stuck, "must trip at 2×threshold while awaiting first verification") + require.Equal(t, int32(0), sh) + require.Equal(t, 2*threshold, age) + + // First exec verifies → window reverts to the plain threshold, + // measured from the verification time. + verifiedAt := base.Add(threshold) + tr.recordVerified(0, verifiedAt) + + _, _, stuck = tr.pickStuck(verifiedAt.Add(threshold-time.Second), threshold) + require.False(t, stuck, "must not trip below 1×threshold after first verification") + + _, _, stuck = tr.pickStuck(verifiedAt.Add(threshold), threshold) + require.True(t, stuck, "must trip at 1×threshold after first verification") +} + +// TestNewShardVerifyTracker_ResumeSkipsDoubledWindow: a resumed shard's +// tasks were submitted (and given their first-verification grace) in a +// prior CAN cycle, so it is seeded as already past its first +// verification and tracks cumulative no-progress against the plain +// threshold. A fresh shard awaits its first verification. +func TestNewShardVerifyTracker_ResumeSkipsDoubledWindow(t *testing.T) { + execs := []*shardedExecutionInfo{{Shard: 0}} + + fresh := newShardVerifyTracker(execs, false, nil) + require.False(t, fresh[0].verifiedAny, "fresh shard awaits its first verification") + + resumed := newShardVerifyTracker(execs, true, map[int32]time.Duration{0: time.Minute}) + require.True(t, resumed[0].verifiedAny, "resumed shard skips the doubled first-verification window") +} From a3a15197694ec4f47cc9d7d274bd35b75a99267f Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Thu, 11 Jun 2026 17:52:11 +0100 Subject: [PATCH 24/35] Avoid duplication of verification handling. The new metrics we add don't interfere with the current metric counts or change behaviour. --- service/worker/migration/activities.go | 13 +++- .../worker/migration/sharded_activities.go | 77 +------------------ 2 files changed, 13 insertions(+), 77 deletions(-) diff --git a/service/worker/migration/activities.go b/service/worker/migration/activities.go index 451fda0864a..a34c3cad8ac 100644 --- a/service/worker/migration/activities.go +++ b/service/worker/migration/activities.go @@ -896,16 +896,20 @@ func (a *activities) verifySingleReplicationTask( }) a.forceReplicationMetricsHandler.Timer(metrics.VerifyDescribeMutableStateLatency.Name()).Record(time.Since(s)) + nsTag := metrics.NamespaceTag(request.Namespace) + switch e := err.(type) { case nil: result, err := a.workflowVerifier(ctx, request, remotAdminClient, a.adminClient, ns, execution, mu) if err == nil && result.status == verified { - a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(request.Namespace)).Counter(metrics.VerifyReplicationTaskSuccess.Name()).Record(1) + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskSuccess.Name()).Record(1) + } else if err == nil && !result.isVerified() { + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskPending.Name()).Record(1) } return result, err case *serviceerror.NotFound: - a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(request.Namespace)).Counter(metrics.VerifyReplicationTaskNotFound.Name()).Record(1) + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskNotFound.Name()).Record(1) // Calling checkSkipWorkflowExecution for every NotFound is sub-optimal as most common case to skip is workflow being deleted due to retention. // A better solution is to only check the existence for workflow which is close to retention period. return a.checkSkipWorkflowExecution(ctx, request, execution, ns) @@ -919,16 +923,17 @@ func (a *activities) verifySingleReplicationTask( if e.Cause == enumspb.RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW { // The passive cluster holds the workflow cache lock while applying history // during SyncWorkflowStateTask. This is actually a small sign of progress. + a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskBusy.Name()).Record(1) return verifyResult{status: notVerified}, nil } - a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(request.Namespace), metrics.ServiceErrorTypeTag(err)). + a.forceReplicationMetricsHandler.WithTags(nsTag, metrics.ServiceErrorTypeTag(err)). Counter(metrics.VerifyReplicationTaskFailed.Name()).Record(1) return verifyResult{ status: notVerified, }, fmt.Errorf("failed to describe workflow from the remote cluster: %w", err) default: - a.forceReplicationMetricsHandler.WithTags(metrics.NamespaceTag(request.Namespace), metrics.ServiceErrorTypeTag(err)). + a.forceReplicationMetricsHandler.WithTags(nsTag, metrics.ServiceErrorTypeTag(err)). Counter(metrics.VerifyReplicationTaskFailed.Name()).Record(1) return verifyResult{ diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 16ec5fbd689..6296f317350 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -8,9 +8,6 @@ import ( "slices" "time" - commonpb "go.temporal.io/api/common/v1" - enumspb "go.temporal.io/api/enums/v1" - "go.temporal.io/api/serviceerror" "go.temporal.io/sdk/activity" "go.temporal.io/sdk/temporal" "go.temporal.io/server/api/adminservice/v1" @@ -492,12 +489,6 @@ func (a *activities) generateReplicationTaskForExec( // attemptVerifyExec runs the source-describe + target-applied check for // a single execution and returns whether it's now verified. -// -// Why this isn't a delegation to verifySingleReplicationTask: that -// helper folds BUSY_WORKFLOW into the generic notVerified result. We -// inline the DMS call here to keep busy-workflow as a distinct metric -// counter, preserving the "passive cluster apply is in progress" -// signal. func (a *activities) attemptVerifyExec( ctx context.Context, remoteAdminClient adminservice.AdminServiceClient, @@ -511,77 +502,17 @@ func (a *activities) attemptVerifyExec( Timer(metrics.VerifyReplicationTaskLatency.Name()).Record(time.Since(attemptStart)) }() - archetype, err := a.archetypeIDToName(ctx, ex.ArchetypeID) - if err != nil { - return false, err - } - vreq := &verifyReplicationTasksRequest{ Namespace: req.Namespace, NamespaceID: req.NamespaceID, TargetClusterName: req.TargetClusterName, } - describeStart := time.Now() - mu, err := remoteAdminClient.DescribeMutableState(ctx, &adminservice.DescribeMutableStateRequest{ - Namespace: req.Namespace, - Execution: &commonpb.WorkflowExecution{ - WorkflowId: ex.BusinessID, - RunId: ex.RunID, - }, - Archetype: archetype, - ArchetypeId: ex.ArchetypeID, - SkipForceReload: true, - }) - a.forceReplicationMetricsHandler.Timer(metrics.VerifyDescribeMutableStateLatency.Name()).Record(time.Since(describeStart)) - - nsTag := metrics.NamespaceTag(req.Namespace) - - if err == nil { - result, vErr := a.workflowVerifier(ctx, vreq, remoteAdminClient, a.adminClient, ns, ex.ExecutionInfo, mu) - if vErr != nil { - return false, vErr - } - if result.isVerified() { - a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskSuccess.Name()).Record(1) - return true, nil - } - a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskPending.Name()).Record(1) - return false, nil - } - - if _, ok := errors.AsType[*serviceerror.NotFound](err); ok { - a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskNotFound.Name()).Record(1) - // Retention/zombie path: a not-found execution may already be - // deleted on source (zombie or past retention), in which case it - // never needs to replicate — treat that as verified so the - // shard's completion accounting moves forward. - result, sErr := a.checkSkipWorkflowExecution(ctx, vreq, ex.ExecutionInfo, ns) - if sErr != nil { - return false, sErr - } - return result.isVerified(), nil - } - - if _, ok := errors.AsType[*serviceerror.NamespaceNotFound](err); ok { - return false, temporal.NewNonRetryableApplicationError( - "failed to describe workflow from the remote cluster", "NamespaceNotFound", err) - } - - if resExhausted, ok := errors.AsType[*serviceerror.ResourceExhausted](err); ok && resExhausted.Cause == enumspb.RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW { - // Passive cluster holds the workflow cache lock while applying - // history during SyncWorkflowStateTask. Counted separately from - // pending so the "apply is in progress" signal stays visible, - // but the workflow-side treatment matches pending — per-exec - // backoff applies and the per-shard last-progress timer does - // not move (it only updates on verified outcomes). - a.forceReplicationMetricsHandler.WithTags(nsTag).Counter(metrics.VerifyReplicationTaskBusy.Name()).Record(1) - return false, nil + result, err := a.verifySingleReplicationTask(ctx, vreq, remoteAdminClient, ns, ex.ExecutionInfo) + if err != nil { + return false, err } - - a.forceReplicationMetricsHandler.WithTags(nsTag, metrics.ServiceErrorTypeTag(err)). - Counter(metrics.VerifyReplicationTaskFailed.Name()).Record(1) - return false, fmt.Errorf("describe workflow on remote cluster: %w", err) + return result.isVerified(), nil } // signalReleaseShards sends the mid-flight ReleaseShards signal to the From 04446370805654546cf7be478ebc3cabb82ffef8 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 12:17:14 +0100 Subject: [PATCH 25/35] Reduce some duplication. --- .../migration/force_replication_workflow.go | 17 ++-- service/worker/migration/sharded_types.go | 20 ++++ .../worker/migration/sharded_types_test.go | 21 +++++ service/worker/migration/sharded_workflow.go | 93 +++++++------------ 4 files changed, 86 insertions(+), 65 deletions(-) diff --git a/service/worker/migration/force_replication_workflow.go b/service/worker/migration/force_replication_workflow.go index e7a6899369d..adf6c57286c 100644 --- a/service/worker/migration/force_replication_workflow.go +++ b/service/worker/migration/force_replication_workflow.go @@ -11,6 +11,11 @@ import ( "go.temporal.io/server/common/metrics" ) +const ( + countWorkflowsForReplicationTimeout = 2 * time.Minute + shardedCountWorkflowsForReplicationTimeout = 30 * time.Second +) + type ( TaskQueueUserDataReplicationParams struct { // PageSize for the SeedReplicationQueueWithUserDataEntries activity @@ -149,7 +154,7 @@ func ForceReplicationWorkflow(ctx workflow.Context, params ForceReplicationParam } if params.TotalForceReplicateWorkflowCount == 0 { - wfCount, err := countWorkflowForReplication(ctx, params) + wfCount, err := countWorkflowsForReplication(ctx, params.Namespace, params.Query, countWorkflowsForReplicationTimeout) if err != nil { return err } @@ -232,7 +237,7 @@ func ForceReplicationWorkflowV2(ctx workflow.Context, params ForceReplicationPar } if params.TotalForceReplicateWorkflowCount == 0 { - wfCount, err := countWorkflowForReplication(ctx, params) + wfCount, err := countWorkflowsForReplication(ctx, params.Namespace, params.Query, countWorkflowsForReplicationTimeout) if err != nil { return err } @@ -454,9 +459,9 @@ func listExecutionsForReplication(ctx workflow.Context, executionsCh workflow.Ch return nil } -func countWorkflowForReplication(ctx workflow.Context, params ForceReplicationParams) (int64, error) { +func countWorkflowsForReplication(ctx workflow.Context, namespace, query string, startToCloseTimeout time.Duration) (int64, error) { ao := workflow.ActivityOptions{ - StartToCloseTimeout: 2 * time.Minute, + StartToCloseTimeout: startToCloseTimeout, RetryPolicy: forceReplicationActivityRetryPolicy, } @@ -466,8 +471,8 @@ func countWorkflowForReplication(ctx workflow.Context, params ForceReplicationPa workflow.WithActivityOptions(ctx, ao), a.CountWorkflow, &workflowservice.CountWorkflowExecutionsRequest{ - Namespace: params.Namespace, - Query: params.Query, + Namespace: namespace, + Query: query, }).Get(ctx, &output); err != nil { return 0, err } diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 4b203491b48..2f9bdd7ba51 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -228,6 +228,26 @@ func (p BatchPayload) merge(src BatchPayload) { } } +// addRunCountsTo increments counts[shard] by the number of runs in p. +func (p BatchPayload) addRunCountsTo(counts map[int32]int) { + //workflowcheck:ignore (per-shard sum is order-independent) + for sh, byBID := range p { + //workflowcheck:ignore (per-shard sum is order-independent) + for _, runs := range byBID { + counts[sh] += len(runs) + } + } +} + +// mergeInto folds p into dst and bumps counts by the runs merged from p. +func (p BatchPayload) mergeInto(dst BatchPayload, counts map[int32]int) { + if len(p) == 0 { + return + } + dst.merge(p) + p.addRunCountsTo(counts) +} + // ShardedForceReplicationParams is the workflow input. Configuration // fields are read-only across CAN cycles; the carry-over block at the // bottom is mutated each cycle. diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 36604c360fb..0bdb40de852 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -91,6 +91,27 @@ func TestBatchPayload_Flatten(t *testing.T) { require.Equal(t, "b-z", got[3].BusinessID) } +// TestBatchPayload_mergeInto merges payload runs into dst and keeps +// bucketCounts in sync — the same invariant addToBucket maintains per +// run, but for a bulk restore on CAN entry. +func TestBatchPayload_mergeInto(t *testing.T) { + dst := BatchPayload{1: {"a": {{RunID: "r0"}}}} + counts := map[int32]int{1: 1} + + src := BatchPayload{ + 1: {"b": {{RunID: "r1"}, {RunID: "r2"}}}, + 2: {"c": {{RunID: "r3"}}}, + } + src.mergeInto(dst, counts) + + require.Equal(t, 2, len(dst)) + require.Len(t, dst[1]["a"], 1) + require.Len(t, dst[1]["b"], 2) + require.Len(t, dst[2]["c"], 1) + require.Equal(t, map[int32]int{1: 3, 2: 1}, counts) + require.Equal(t, 4, dst.totalRuns()) +} + // TestShardVerifyTracker_FirstVerificationDoubledWindow pins the // no-progress backstop's grace for a shard's first verified outcome: a // shard that hasn't verified anything yet gets 2×threshold before diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 2a813d7fddd..ae31f6d3c99 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -99,7 +99,7 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe // via the same CountWorkflow activity upstream uses. Skipped on // subsequent CAN cycles — the count carries across via params. if params.TotalForceReplicateWorkflowCount == 0 { - wfCount, err := shardedCountWorkflowsForReplication(ctx, ¶ms) + wfCount, err := countWorkflowsForReplication(ctx, params.Namespace, params.Query, shardedCountWorkflowsForReplicationTimeout) if err != nil { return err } @@ -167,29 +167,6 @@ func maybeKickoffShardedTaskQueueUserDataReplication(ctx workflow.Context, param return child.GetChildWorkflowExecution().Get(ctx, &childExecution) } -// shardedCountWorkflowsForReplication asks the frontend how many -// workflows match the namespace's force-rep query. Used once at -// workflow start to seed TotalForceReplicateWorkflowCount for the -// status query's progress reporting. -func shardedCountWorkflowsForReplication(ctx workflow.Context, params *ShardedForceReplicationParams) (int64, error) { - ao := workflow.ActivityOptions{ - StartToCloseTimeout: 30 * time.Second, - RetryPolicy: forceReplicationActivityRetryPolicy, - } - var a *activities - var output countWorkflowResponse - if err := workflow.ExecuteActivity( - workflow.WithActivityOptions(ctx, ao), - a.CountWorkflow, - &workflowservice.CountWorkflowExecutionsRequest{ - Namespace: params.Namespace, - Query: params.Query, - }).Get(ctx, &output); err != nil { - return 0, err - } - return output.WorkflowCount, nil -} - // shardedWorkflowState holds the workflow's per-run state. Workflow // coroutines yield only at SDK calls, so plain maps + ints are safe // without mutexes — workflow.Await re-evaluates its predicate after @@ -287,23 +264,17 @@ type shardedWorkflowState struct { metricsHandler sdkclient.MetricsHandler } -func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicationParams) (*shardedWorkflowState, error) { - lao := workflow.LocalActivityOptions{ - StartToCloseTimeout: 1 * time.Second, - RetryPolicy: forceReplicationActivityRetryPolicy, - } - localCtx := workflow.WithLocalActivityOptions(ctx, lao) - var a *activities - var md MetadataResponse - if err := workflow.ExecuteLocalActivity(localCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { - return nil, err - } - var targetMd DescribeTargetClusterResponse - if err := workflow.ExecuteLocalActivity(localCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ - TargetClusterName: params.TargetClusterName, - }).Get(ctx, &targetMd); err != nil { - return nil, err - } +// defaultConcurrentBatchCount derives the in-flight-batch ceiling +// from the target cluster's shard count: a quarter of the shards, +// capped at defaultConcurrentBatchCap. The 1/4 fraction leaves worker +// slots free for unrelated activities; the absolute cap bounds the +// cluster blast radius regardless of cluster size. Returns at least 1. +func defaultConcurrentBatchCount(shards int32) int { + return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) +} + +// applyShardedDefaults fills zero-valued tuning fields on params. +func applyShardedDefaults(params *ShardedForceReplicationParams, targetShardCount int32) { if params.BatchSize <= 0 { params.BatchSize = defaultBatchSize } @@ -329,11 +300,31 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati params.PerBatchGenerateRPS = defaultPerBatchGenerateRPS } if params.ConcurrentBatchCount <= 0 { - params.ConcurrentBatchCount = defaultConcurrentBatchCount(targetMd.ShardCount) + params.ConcurrentBatchCount = defaultConcurrentBatchCount(targetShardCount) } if params.EstimationMultiplier <= 0 { params.EstimationMultiplier = 2 } +} + +func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicationParams) (*shardedWorkflowState, error) { + lao := workflow.LocalActivityOptions{ + StartToCloseTimeout: 1 * time.Second, + RetryPolicy: forceReplicationActivityRetryPolicy, + } + localCtx := workflow.WithLocalActivityOptions(ctx, lao) + var a *activities + var md MetadataResponse + if err := workflow.ExecuteLocalActivity(localCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { + return nil, err + } + var targetMd DescribeTargetClusterResponse + if err := workflow.ExecuteLocalActivity(localCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ + TargetClusterName: params.TargetClusterName, + }).Get(ctx, &targetMd); err != nil { + return nil, err + } + applyShardedDefaults(params, targetMd.ShardCount) // QPSQueue is sized off ConcurrentBatchCount (one sample slot per // expected in-flight batch + one for the starting count). Seeded // with the current ReplicatedWorkflowCount so the very first @@ -359,14 +350,7 @@ func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicati // Restore execs recovered from cancel-before-start batches in // the prior cycle so the streaming packer picks them up // alongside any new pages. - s.buckets.merge(params.RecoveredBuckets) - //workflowcheck:ignore (per-shard sum is order-independent) - for sh, byBID := range params.RecoveredBuckets { - //workflowcheck:ignore (per-shard sum is order-independent) - for _, runs := range byBID { - s.bucketCounts[sh] += len(runs) - } - } + params.RecoveredBuckets.mergeInto(s.buckets, s.bucketCounts) params.RecoveredBuckets = nil return s, nil } @@ -644,15 +628,6 @@ func (s *shardedWorkflowState) recordVerified(ctx workflow.Context, verified int s.metricsHandler.Gauge(ForceReplicationRpsTagName).Update(s.params.ReplicatedWorkflowCountPerSecond) } -// defaultConcurrentBatchCount derives the in-flight-batch ceiling -// from the target cluster's shard count: a quarter of the shards, -// capped at defaultConcurrentBatchCap. The 1/4 fraction leaves worker -// slots free for unrelated activities; the absolute cap bounds the -// cluster blast radius regardless of cluster size. Returns at least 1. -func defaultConcurrentBatchCount(shards int32) int { - return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) -} - // collectRecoveredBuckets re-buckets any execs from batches whose // dispatching activity returned CanceledError without returning a // result — i.e. the activity body never ran, so its execs were From 6a80bfab2ff84d717a41a860918b701159f2838a Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 12:26:57 +0100 Subject: [PATCH 26/35] Claim shards in spawnBatch. --- service/worker/migration/sharded_workflow.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index ae31f6d3c99..ceffdff8727 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -849,10 +849,6 @@ func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { s.params.ResumeShards = unpackResumeBatches(batches[i:]) return } - //workflowcheck:ignore (setting a set of flags is order-independent) - for sh := range batch.payload { - s.shardInFlight[sh] = true - } s.spawnBatch(ctx, batch.payload, true, batch.noProgress) } } @@ -1056,6 +1052,7 @@ func (s *shardedWorkflowState) spawnBatch( held := make(map[int32]bool, len(payload)) //workflowcheck:ignore (building a set of flags is order-independent) for sh := range payload { + s.shardInFlight[sh] = true held[sh] = true } s.heldByBatch[batchID] = held @@ -1146,7 +1143,6 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool } payload[sh] = s.takeFromBucket(sh, take) packed += take - s.shardInFlight[sh] = true } if packed == 0 { return false From d47e211de62c9c6a133aa5383d0376a43a268797 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 12:43:25 +0100 Subject: [PATCH 27/35] Remove redundant err check. --- service/worker/migration/sharded_workflow.go | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index ceffdff8727..4e7eb7da66b 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -837,13 +837,10 @@ func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { batches := s.packResumeBatchPlan(entries) for i, batch := range batches { - if s.lastErr != nil { - s.params.ResumeShards = unpackResumeBatches(batches[i:]) - return - } // Block until a dispatch slot is free so resume payloads // can't overshoot ConcurrentBatchCount on cycles that - // carried many shards across CAN. + // carried many shards across CAN. The await also wakes + // immediately when lastErr is set. s.waitForDispatchSlot(ctx) if s.lastErr != nil { s.params.ResumeShards = unpackResumeBatches(batches[i:]) From 43aace024808c7a6d4839bb910798866f369d8cf Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 12:55:32 +0100 Subject: [PATCH 28/35] Helpers for type conversion. --- .../worker/migration/sharded_activities.go | 28 ++++------ service/worker/migration/sharded_types.go | 55 +++++++++++++++++++ .../worker/migration/sharded_types_test.go | 26 +++++++++ service/worker/migration/sharded_workflow.go | 12 ++-- 4 files changed, 95 insertions(+), 26 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 6296f317350..8b864d310c3 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -673,36 +673,28 @@ func buildInFlight( shards shardVerifyTracker, now time.Time, ) []ResumeShard { - byShard := map[int32]map[string][]RunEntry{} + payload := BatchPayload{} for i, ex := range execs { if verified[i] { continue } - if byShard[ex.Shard] == nil { - byShard[ex.Shard] = map[string][]RunEntry{} + if payload[ex.Shard] == nil { + payload[ex.Shard] = map[string][]RunEntry{} } - byShard[ex.Shard][ex.BusinessID] = append(byShard[ex.Shard][ex.BusinessID], RunEntry{ + payload[ex.Shard][ex.BusinessID] = append(payload[ex.Shard][ex.BusinessID], RunEntry{ RunID: ex.RunID, ArchetypeID: ex.ArchetypeID, }) } - if len(byShard) == 0 { + if len(payload) == 0 { return nil } - shardIDs := make([]int32, 0, len(byShard)) - for sh := range byShard { - shardIDs = append(shardIDs, sh) - } - slices.Sort(shardIDs) - out := make([]ResumeShard, 0, len(shardIDs)) - for _, sh := range shardIDs { - out = append(out, ResumeShard{ - Shard: sh, - Execs: byShard[sh], - NoProgressDuration: now.Sub(shards[sh].lastProgress), - }) + noProgress := make(map[int32]time.Duration, len(payload)) + //workflowcheck:ignore (one entry per shard; order-independent) + for sh := range payload { + noProgress[sh] = now.Sub(shards[sh].lastProgress) } - return out + return resumeShardsFromPayload(payload, noProgress) } // firstUnverifiedOnShard returns the index of the first execution in the diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 2f9bdd7ba51..c3518c4eddf 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -336,12 +336,67 @@ type ShardedForceReplicationParams struct { // for that BID. Grouping by BID at the wire level lets a hot BID (with // many runs) collapse to one BID-string + N tuples rather than N copies // of the BID; see BatchPayload's docstring. +// +// The slice-of-ResumeShard form is the CAN/activity wire shape: JSON- +// friendly, easy to append from drained batches, and sortable for replay +// determinism. The packer and activity input use BatchPayload plus a +// per-shard no-progress map; resumeShardsToPayload and +// resumeShardsFromPayload convert between the two. type ResumeShard struct { Shard int32 Execs map[string][]RunEntry NoProgressDuration time.Duration } +// resumeShardsFromPayload expands a BatchPayload and its per-shard +// no-progress durations into a shard-sorted ResumeShard slice. +func resumeShardsFromPayload(payload BatchPayload, noProgress map[int32]time.Duration) []ResumeShard { + if len(payload) == 0 { + return nil + } + out := make([]ResumeShard, 0, len(payload)) + for _, sh := range payload.sortedShards() { + execs := payload[sh] + if len(execs) == 0 { + continue + } + out = append(out, ResumeShard{ + Shard: sh, + Execs: execs, + NoProgressDuration: noProgress[sh], + }) + } + return out +} + +// resumeShardsToPayload folds a ResumeShard slice into the BatchPayload +// and per-shard no-progress map the packer and activity input use. +// Duplicate shard entries merge execs; the last NoProgressDuration wins. +func resumeShardsToPayload(shards []ResumeShard) (BatchPayload, map[int32]time.Duration) { + if len(shards) == 0 { + return nil, nil + } + payload := BatchPayload{} + noProgress := map[int32]time.Duration{} + for _, rs := range shards { + if len(rs.Execs) == 0 { + continue + } + if payload[rs.Shard] == nil { + payload[rs.Shard] = map[string][]RunEntry{} + } + //workflowcheck:ignore (writes are to disjoint BID keys per shard entry; order-independent) + for bid, runs := range rs.Execs { + payload[rs.Shard][bid] = append(payload[rs.Shard][bid], runs...) + } + noProgress[rs.Shard] = rs.NoProgressDuration + } + if len(payload) == 0 { + return nil, nil + } + return payload, noProgress +} + // shardedBatchReq is the per-batch activity input. Executions is the // per-shard, per-BID nested payload — the workflow has marked every // shard appearing as a top-level key in shardInFlight before dispatch, diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 0bdb40de852..6c5b5b11852 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -112,6 +112,32 @@ func TestBatchPayload_mergeInto(t *testing.T) { require.Equal(t, 4, dst.totalRuns()) } +func TestResumeShardPayloadRoundTrip(t *testing.T) { + shards := []ResumeShard{ + {Shard: 2, Execs: map[string][]RunEntry{"b": {{RunID: "r2"}}}, NoProgressDuration: 3 * time.Second}, + {Shard: 1, Execs: map[string][]RunEntry{"a": {{RunID: "r1"}, {RunID: "r1b"}}}, NoProgressDuration: 5 * time.Second}, + } + payload, noProgress := resumeShardsToPayload(shards) + got := resumeShardsFromPayload(payload, noProgress) + + require.Len(t, got, 2) + require.Equal(t, int32(1), got[0].Shard) + require.Equal(t, 5*time.Second, got[0].NoProgressDuration) + require.Len(t, got[0].Execs["a"], 2) + require.Equal(t, int32(2), got[1].Shard) + require.Equal(t, 3*time.Second, got[1].NoProgressDuration) +} + +func TestResumeShardsToPayload_mergesDuplicateShards(t *testing.T) { + shards := []ResumeShard{ + {Shard: 1, Execs: map[string][]RunEntry{"a": {{RunID: "r1"}}}, NoProgressDuration: time.Second}, + {Shard: 1, Execs: map[string][]RunEntry{"b": {{RunID: "r2"}}}, NoProgressDuration: 2 * time.Second}, + } + payload, noProgress := resumeShardsToPayload(shards) + require.Len(t, payload[1], 2) + require.Equal(t, 2*time.Second, noProgress[1]) +} + // TestShardVerifyTracker_FirstVerificationDoubledWindow pins the // no-progress backstop's grace for a shard's first verified outcome: a // shard that hasn't verified anything yet gets 2×threshold before diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 4e7eb7da66b..a07de80cab0 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -882,16 +882,12 @@ func (s *shardedWorkflowState) packResumeBatchPlan(entries []ResumeShard) []resu // shard ID so the slice flowing into the CAN args is deterministic // across replays. func unpackResumeBatches(batches []resumeBatch) []ResumeShard { + if len(batches) == 0 { + return nil + } var out []ResumeShard for _, b := range batches { - //workflowcheck:ignore (output is sorted below before any observable use) - for sh, execs := range b.payload { - out = append(out, ResumeShard{ - Shard: sh, - Execs: execs, - NoProgressDuration: b.noProgress[sh], - }) - } + out = append(out, resumeShardsFromPayload(b.payload, b.noProgress)...) } slices.SortFunc(out, func(a, b ResumeShard) int { return int(a.Shard - b.Shard) From 2f24de53d1e40566defac6d7546b65ec98a71520 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 13:16:45 +0100 Subject: [PATCH 29/35] Simplify. --- .../worker/migration/sharded_activities.go | 43 +++++++++---------- 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 8b864d310c3..024d09de9a7 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -269,9 +269,6 @@ func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, e rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) for i := startIdx; i < len(execs); i++ { ex := execs[i] - if ctx.Err() != nil { - return temporal.NewCanceledError("inject phase cancelled") - } if err := a.generateReplicationTaskForExec(ctx, rateLimiter, req, ex); err != nil { if ctx.Err() != nil { return temporal.NewCanceledError("inject phase cancelled") @@ -448,17 +445,14 @@ func waitNextTick( if remaining := drainGrace - time.Since(drainStartAt); remaining > 0 && remaining < sleepDur { sleepDur = remaining } - select { - case <-time.After(sleepDur): - case <-callCtx.Done(): - } - return + } + wakeCtx := ctx + if draining { + wakeCtx = callCtx } select { case <-time.After(sleepDur): - case <-ctx.Done(): - // ctx cancel just sets draining on the next iteration; don't - // unwind here. + case <-wakeCtx.Done(): } } @@ -604,12 +598,12 @@ func (t shardVerifyTracker) totalIdleCost(now time.Time) time.Duration { return total } -// awaitingRelease returns completed-but-not-yet-signaled shard IDs in -// ascending order so the signal payload is deterministic across replays. -func (t shardVerifyTracker) awaitingRelease() []int32 { +// completedShards returns shard IDs matching filter, sorted ascending +// so signal payloads and activity returns are deterministic. +func (t shardVerifyTracker) completedShards(filter func(shardVerify) bool) []int32 { var out []int32 for sh, sv := range t { - if !sv.doneAt.IsZero() { + if filter(sv) { out = append(out, sh) } } @@ -617,17 +611,20 @@ func (t shardVerifyTracker) awaitingRelease() []int32 { return out } +// awaitingRelease returns completed-but-not-yet-signaled shard IDs in +// ascending order so the signal payload is deterministic across replays. +func (t shardVerifyTracker) awaitingRelease() []int32 { + return t.completedShards(func(sv shardVerify) bool { + return !sv.doneAt.IsZero() + }) +} + // allCompleted returns every shard that finished during this activity // run — both signal-released and still awaiting release at return. func (t shardVerifyTracker) allCompleted() []int32 { - var out []int32 - for sh, sv := range t { - if sv.released || !sv.doneAt.IsZero() { - out = append(out, sh) - } - } - slices.Sort(out) - return out + return t.completedShards(func(sv shardVerify) bool { + return sv.released || !sv.doneAt.IsZero() + }) } // pickStuck returns (shard, age, true) for the lowest-numbered shard From fd3579adf07c38586143af7b5d1f54a996c31eed Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 12 Jun 2026 13:27:30 +0100 Subject: [PATCH 30/35] Lint. --- service/worker/migration/sharded_types_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 6c5b5b11852..eac27f45681 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -104,7 +104,7 @@ func TestBatchPayload_mergeInto(t *testing.T) { } src.mergeInto(dst, counts) - require.Equal(t, 2, len(dst)) + require.Len(t, dst, 2) require.Len(t, dst[1]["a"], 1) require.Len(t, dst[1]["b"], 2) require.Len(t, dst[2]["c"], 1) From f56cebec05746a0eecf7495ffdbc75ccb4633940 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Wed, 17 Jun 2026 18:30:21 +0100 Subject: [PATCH 31/35] Refactor to avoid continue as new handling. --- .../migration/force_replication_workflow.go | 13 +- service/worker/migration/fx.go | 3 + .../worker/migration/sharded_activities.go | 285 +--- .../migration/sharded_activities_test.go | 69 +- .../worker/migration/sharded_handover_test.go | 474 ++++++ .../migration/sharded_parent_workflow.go | 579 ++++++++ service/worker/migration/sharded_types.go | 278 ++-- .../worker/migration/sharded_types_test.go | 111 +- service/worker/migration/sharded_workflow.go | 1311 +++++------------ .../worker/migration/sharded_workflow_test.go | 755 +++------- 10 files changed, 1869 insertions(+), 2009 deletions(-) create mode 100644 service/worker/migration/sharded_handover_test.go create mode 100644 service/worker/migration/sharded_parent_workflow.go diff --git a/service/worker/migration/force_replication_workflow.go b/service/worker/migration/force_replication_workflow.go index adf6c57286c..9ed400398bd 100644 --- a/service/worker/migration/force_replication_workflow.go +++ b/service/worker/migration/force_replication_workflow.go @@ -92,16 +92,11 @@ type ( ReplicatedWorkflowCountPerSecond float64 PageTokenForRestart []byte - // Sharded-workflow-only recovery bundle: feed these three - // fields back into a fresh ShardedForceReplicationWorkflow's - // NextPageToken / ResumeShards / RecoveredBuckets params to - // resume from a failed run without missing executions. Left - // zero by the legacy ForceReplicationWorkflow variants — - // their PageTokenForRestart is the start-of-run token and - // already covers all in-flight execs at restart cost. + // RecoveryNextPageToken is the page token the sharded workflow + // was processing when it last continued-as-new or was interrupted. + // Feed this back into ShardedForceReplicationWorkflow.NextPageToken + // to resume from that position. Left zero by the legacy variants. RecoveryNextPageToken []byte - RecoveryResumeShards []ResumeShard - RecoveryBuckets BatchPayload } ) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index 40cecd13c4a..b4b12e38783 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -139,6 +139,9 @@ func (sc *shardedWorkerComponent) RegisterWorkflow(registry sdkworker.Registry) registry.RegisterWorkflowWithOptions(ShardedForceReplicationWorkflow, workflow.RegisterOptions{ Name: shardedForceReplicationWorkflowName, }) + registry.RegisterWorkflowWithOptions(shardedForceReplicationWorker, workflow.RegisterOptions{ + Name: shardedForceReplicationWorkerName, + }) registry.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{ Name: forceTaskQueueUserDataReplicationWorkflow, }) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 024d09de9a7..29d522d7270 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -2,7 +2,6 @@ package migration import ( "context" - "errors" "fmt" "math" "slices" @@ -52,13 +51,10 @@ func (a *activities) DescribeTargetCluster(_ context.Context, req DescribeTarget } // ReplicateBatch is the per-batch activity body for the sharded force -// replication workflow. Runs inject (skipped on Resume) then verify, -// signal-releasing completed shards mid-flight as their cumulative -// idle cost crosses IdleShardCost. On workflow-initiated cancellation -// it enters drain mode and returns a replicateBatchResult carrying -// any still-unverified execs. Drain-mode signal traffic is suppressed: -// once we know we're about to return, there's no point racing a -// signal against the return value. +// replication workflow. Runs inject (heartbeat-resumable via +// NextInjectIdx/InjectDone) then verify, signal-releasing completed +// shards mid-flight as their cumulative idle cost crosses IdleShardCost. +// Returns {CompletedShards, VerifiedCount} on success. func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { // Flatten once so per-exec bookkeeping (verified[], attempts[], // nextRetryAt[]) can stay index-based. @@ -74,7 +70,7 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( } // ---- Inject phase ---- - if !req.Resume && !hb.InjectDone { + if !hb.InjectDone { startIdx := hb.NextInjectIdx if err := a.runInjectPhase(ctx, req, execs, startIdx); err != nil { return replicateBatchResult{}, err @@ -108,10 +104,9 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( } // runVerifyPhase is the verify-phase loop body of ReplicateBatch. It -// owns per-exec bookkeeping, the drain-transition handoff, and the -// per-iteration completion / stuck-shard / signal-release decisions -// — extracted from ReplicateBatch to keep its cognitive complexity -// under the linter cap. +// owns per-exec bookkeeping, the per-iteration completion / stuck-shard / +// signal-release decisions — extracted from ReplicateBatch to keep its +// cognitive complexity under the linter cap. func (a *activities) runVerifyPhase( ctx context.Context, req *shardedBatchReq, @@ -125,28 +120,12 @@ func (a *activities) runVerifyPhase( nextRetryAt := make([]time.Time, execCount) doneCount := 0 - shards := newShardVerifyTracker(execs, req.Resume, req.NoProgressByShard) - - var draining bool - var drainStartAt time.Time - - // callCtx is what attemptVerifyExec uses for DescribeMutableState. - // In drain mode we swap to a detached context: the parent ctx is - // already dead (that's what triggered the transition), so reusing - // it would make every drain-mode RPC fail instantly. - callCtx := ctx - // Pre-create the drain context up front and defer cancel right away - // so go vet's lostcancel pass sees the canonical pattern. - drainCtx, drainCancel := context.WithCancel(context.Background()) - defer drainCancel() + shards := newShardVerifyTracker(execs) for { // Worker shutdown short-circuits with a retryable error so the - // SDK reschedules on another worker. Folding partial state into - // the workflow's drain bucket would conflate worker shutdown - // (deploys, host loss — orthogonal to migration progress) with - // drain-for-CAN. Returning here also avoids the ~HeartbeatTimeout - // wait that the alternative (silent worker death) would incur + // SDK reschedules on another worker. Returning here avoids the + // ~HeartbeatTimeout wait that silent worker death would incur // before the server retries the attempt. Inject is already // heartbeat-preserved (InjectDone), so retry skips it; verify // re-runs from scratch but DMS reads are idempotent. @@ -159,58 +138,30 @@ func (a *activities) runVerifyPhase( default: } - // Workflow-initiated activity cancellation (drainForCAN) - // transitions us into drain mode with the full DrainGrace window. - // WaitForCancellation=true on the activity options guarantees the - // workflow blocks for us, so the grace window is genuinely - // available — swap callCtx onto a detached deadline so - // DescribeMutableState keeps working after the parent ctx died. - if !draining && ctx.Err() != nil { - draining = true - drainStartAt = time.Now() - // Start the drain budget timer here rather than at activity - // entry so the grace window measures from drain transition, - // not from activity start. - time.AfterFunc(req.DrainGrace, drainCancel) - callCtx = drainCtx - } - - passDelta, minNextRetry, ctxAborted, vErr := a.runVerifyPass( - ctx, callCtx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) - // Fold partial progress in before the error check — the SDK - // discards the activity result on failure, so the only way - // the workflow learns about partially-verified execs on the - // error path is via wrapBatchVerifyError encoding the count - // as ApplicationError details below. + passDelta, minNextRetry, vErr := a.runVerifyPass( + ctx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) doneCount += passDelta if vErr != nil { - return replicateBatchResult{}, wrapBatchVerifyError(vErr, int64(doneCount)) + return replicateBatchResult{}, vErr } activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) if done, result, err := a.evaluateVerifyIteration( - ctx, req, execs, verified, shards, doneCount, execCount, draining, drainStartAt); err != nil { - return replicateBatchResult{}, wrapBatchVerifyError(err, int64(doneCount)) + ctx, req, execs, verified, shards, doneCount, execCount); err != nil { + return replicateBatchResult{}, err } else if done { return result, nil } - // If the inner loop aborted because callCtx died, skip the sleep - // entirely so the outer-loop top sees the new state promptly - // (normal → drain transition, or drain → exit). - if ctxAborted { - continue - } - - waitNextTick(ctx, callCtx, minNextRetry, draining, drainStartAt, req.DrainGrace) + waitNextTick(ctx, minNextRetry) } } // evaluateVerifyIteration runs the post-pass checks (clean completion, -// stuck-shard backstop, drain-exit decision, mid-flight signal release) -// for a single verify-loop iteration. Returns done=true with the result -// when the loop should exit; otherwise (false, _, nil) means continue. +// stuck-shard backstop, mid-flight signal release) for a single +// verify-loop iteration. Returns done=true with the result when the loop +// should exit; otherwise (false, _, nil) means continue. func (a *activities) evaluateVerifyIteration( ctx context.Context, req *shardedBatchReq, @@ -218,8 +169,6 @@ func (a *activities) evaluateVerifyIteration( verified []bool, shards shardVerifyTracker, doneCount, execCount int, - draining bool, - drainStartAt time.Time, ) (bool, replicateBatchResult, error) { // Clean completion — every exec verified. if doneCount >= execCount { @@ -229,24 +178,6 @@ func (a *activities) evaluateVerifyIteration( }, nil } - if draining { - // Drain-mode exit checks. No signals here — the return value - // carries everything the workflow needs (completed shards + - // unverified execs grouped by shard with their cumulative - // no-progress duration). The per-shard no-progress backstop is - // deliberately skipped: drain is bounded by DrainGrace and the - // outstanding execs need to flow back via InFlight for CAN - // carry-over, not surface as a ShardNoProgress failure. - if shouldExitDrain(req, shards, drainStartAt) { - return true, replicateBatchResult{ - CompletedShards: shards.allCompleted(), - InFlight: buildInFlight(execs, verified, shards, time.Now()), - VerifiedCount: int64(doneCount), - }, nil - } - return false, replicateBatchResult{}, nil - } - // Per-shard cumulative no-progress backstop. if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, execCount); sErr != nil { return false, replicateBatchResult{}, sErr @@ -260,11 +191,9 @@ func (a *activities) evaluateVerifyIteration( // runInjectPhase walks execs in flattened order, generating one // replication task per exec under a per-batch RPS limiter. Cancellation -// mid-loop returns a CanceledError so spawnBatch's IsCanceledError -// check fires and the batch's batchExecs entry is preserved for -// RecoveredBuckets re-injection next cycle. Already-injected execs -// are re-injected harmlessly — replication dedupes per (namespace, -// wf, run). +// mid-loop returns a CanceledError so the caller can propagate it. +// Already-injected execs are re-injected harmlessly — replication dedupes +// per (namespace, wf, run). func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, execs []*shardedExecutionInfo, startIdx int) error { rateLimiter := quotas.NewRateLimiter(req.PerBatchGenerateRPS, int(math.Ceil(req.PerBatchGenerateRPS))) for i := startIdx; i < len(execs); i++ { @@ -290,19 +219,14 @@ func (a *activities) runInjectPhase(ctx context.Context, req *shardedBatchReq, e // runVerifyPass runs one pass over every unverified exec, attempting a // verify on those whose backoff timer has expired. Returns the count of // execs newly verified this pass, the earliest pending retry deadline -// (for sleep scheduling), and whether callCtx died mid-pass — in which -// case the outer loop's top reassesses (drain transition or exit). -// Returns a non-nil error only for hard errors from the verify path; -// ctx-derived errors set ctxAborted instead so the outer loop owns -// the decision about what to do next. +// (for sleep scheduling), and a non-nil error only for hard errors from +// the verify path. // -// ctx is the activity ctx, used only for heartbeating — a single pass -// over a large batch can outlast HeartbeatTimeout if we only heartbeat -// once at the end, so we tick per attempted exec. callCtx is what the -// DMS call rides on (the detached drain ctx in drain mode). +// ctx is the activity ctx, used for both the DMS call and heartbeating — +// a single pass over a large batch can outlast HeartbeatTimeout if we +// only heartbeat once at the end, so we tick per attempted exec. func (a *activities) runVerifyPass( ctx context.Context, - callCtx context.Context, remoteAdminClient adminservice.AdminServiceClient, ns *namespace.Namespace, req *shardedBatchReq, @@ -311,7 +235,7 @@ func (a *activities) runVerifyPass( attempts []int, nextRetryAt []time.Time, shards shardVerifyTracker, -) (int, time.Time, bool, error) { +) (int, time.Time, error) { now := time.Now() var minNextRetry time.Time verifiedDelta := 0 @@ -324,16 +248,9 @@ func (a *activities) runVerifyPass( continue } - ok, err := a.attemptVerifyExec(callCtx, remoteAdminClient, ns, req, ex) + ok, err := a.attemptVerifyExec(ctx, remoteAdminClient, ns, req, ex) if err != nil { - if callCtx.Err() != nil { - // callCtx is dead. Two cases: (1) normal-mode parent ctx - // was just cancelled mid-call — the outer-loop top will - // promote to drain on the next iteration; (2) drain-mode - // detached ctx expired — the drain exit check fires. - return verifiedDelta, minNextRetry, true, nil - } - return verifiedDelta, minNextRetry, false, err + return verifiedDelta, minNextRetry, err } if ok { @@ -348,7 +265,7 @@ func (a *activities) runVerifyPass( activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) } - return verifiedDelta, minNextRetry, false, nil + return verifiedDelta, minNextRetry, nil } // earliest returns the earlier of cur (which may be zero) and candidate. @@ -363,7 +280,8 @@ func earliest(cur, candidate time.Time) time.Time { // checkStuckShard fails non-retryably if any shard has gone longer than // req.ShardNoProgress without a verified outcome (double that for a shard // still awaiting its first verification — see pickStuck). Duration is -// cumulative across CAN cycles via tracker seeding. +// cumulative from activity start (seeded at time.Now() on first activity +// attempt). func (a *activities) checkStuckShard( req *shardedBatchReq, shards shardVerifyTracker, @@ -384,28 +302,16 @@ func (a *activities) checkStuckShard( return temporal.NewNonRetryableApplicationError(msg, "ShardNoProgress", nil) } -// shouldExitDrain reports whether the drain-mode exit conditions are -// met: either the grace window has expired, or the cumulative idle cost -// across completed-but-unsignaled shards crossed the threshold. -func shouldExitDrain(req *shardedBatchReq, shards shardVerifyTracker, drainStartAt time.Time) bool { - if time.Since(drainStartAt) >= req.DrainGrace { - return true - } - return shards.totalIdleCost(time.Now()) >= req.IdleShardCost -} - // maybeSignalRelease signals the workflow to release any // completed-but-unsignaled shards if their cumulative idle cost crossed -// the threshold. Only fires in normal mode — drain mode rides the -// activity result instead. +// the threshold. // // Ctx-canceled errors from signalReleaseShards are suppressed: a // workflow-initiated cancel arriving mid-signal would otherwise // surface as a wrapped ctx-canceled error (not temporal.CanceledError) -// that the workflow side wouldn't recognise via IsCanceledError — -// turning a clean CAN into an error exit. Suppressing here lets the -// outer loop see ctx.Err() at its top and promote to drain mode -// normally. +// that the workflow side wouldn't recognise via IsCanceledError. +// Suppressing here lets the outer loop exit cleanly when the activity's +// context is cancelled. func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchReq, shards shardVerifyTracker) error { if shards.totalIdleCost(time.Now()) < req.IdleShardCost { return nil @@ -424,35 +330,17 @@ func (a *activities) maybeSignalRelease(ctx context.Context, req *shardedBatchRe return nil } -// waitNextTick sleeps until the next exec is due for retry, capped by -// DrainGrace remaining when in drain mode. In drain mode the parent ctx -// is already dead, so we wake on the detached drain ctx instead — using -// the parent ctx would tight-loop on its Done channel. -func waitNextTick( - ctx, callCtx context.Context, - minNextRetry time.Time, - draining bool, - drainStartAt time.Time, - drainGrace time.Duration, -) { +// waitNextTick sleeps until the next exec is due for retry or ctx is done. +func waitNextTick(ctx context.Context, minNextRetry time.Time) { sleepDur := 50 * time.Millisecond if !minNextRetry.IsZero() { if delta := time.Until(minNextRetry); delta > sleepDur { sleepDur = delta } } - if draining { - if remaining := drainGrace - time.Since(drainStartAt); remaining > 0 && remaining < sleepDur { - sleepDur = remaining - } - } - wakeCtx := ctx - if draining { - wakeCtx = callCtx - } select { case <-time.After(sleepDur): - case <-wakeCtx.Done(): + case <-ctx.Done(): } } @@ -515,9 +403,7 @@ func (a *activities) attemptVerifyExec( // while this activity stays running on its still-pending shards. // // No retry wrapping: a transient failure propagates up so the activity -// fails, the workflow records it via lastErr, and the in-flight batch is -// recovered into the next CAN's RecoveredBuckets — preferable to silently -// swallowing the error here and stranding completed shards. +// fails; the child workflow records the error via lastErr. func (a *activities) signalReleaseShards(ctx context.Context, req *shardedBatchReq, shards []int32) error { info := activity.GetInfo(ctx) return a.sdkClientFactory.GetSystemClient().SignalWorkflow(ctx, info.WorkflowExecution.ID, info.WorkflowExecution.RunID, releaseShardsSignalName, releaseShardsPayload{ @@ -537,11 +423,11 @@ type shardVerify struct { type shardVerifyTracker map[int32]shardVerify -func newShardVerifyTracker( - execs []*shardedExecutionInfo, - resume bool, - noProgressByShard map[int32]time.Duration, -) shardVerifyTracker { +// newShardVerifyTracker builds the per-shard state from the flattened +// exec slice. Each shard's lastProgress starts at time.Now() so the +// stuck-shard backstop measures elapsed time from activity start rather +// than the activity's scheduling epoch. +func newShardVerifyTracker(execs []*shardedExecutionInfo) shardVerifyTracker { t := shardVerifyTracker{} for _, ex := range execs { sv := t[ex.Shard] @@ -550,16 +436,7 @@ func newShardVerifyTracker( } nowSeed := time.Now() for sh, sv := range t { - if resume { - sv.lastProgress = nowSeed.Add(-noProgressByShard[sh]) - // Resumed shards had their tasks submitted (and their initial - // first-verification grace) in a prior CAN cycle — inject is - // skipped on resume — so they continue cumulative no-progress - // tracking against the normal window, not the doubled one. - sv.verifiedAny = true - } else { - sv.lastProgress = nowSeed - } + sv.lastProgress = nowSeed t[sh] = sv } return t @@ -660,40 +537,6 @@ func (t shardVerifyTracker) pickStuck(now time.Time, threshold time.Duration) (i return minShard, minAge, found } -// buildInFlight groups unverified execs by shard then businessID and -// attaches the cumulative no-progress duration per shard, for the -// drain-mode activity return. Shards with zero unverified execs are -// reported via CompletedShards instead. -func buildInFlight( - execs []*shardedExecutionInfo, - verified []bool, - shards shardVerifyTracker, - now time.Time, -) []ResumeShard { - payload := BatchPayload{} - for i, ex := range execs { - if verified[i] { - continue - } - if payload[ex.Shard] == nil { - payload[ex.Shard] = map[string][]RunEntry{} - } - payload[ex.Shard][ex.BusinessID] = append(payload[ex.Shard][ex.BusinessID], RunEntry{ - RunID: ex.RunID, - ArchetypeID: ex.ArchetypeID, - }) - } - if len(payload) == 0 { - return nil - } - noProgress := make(map[int32]time.Duration, len(payload)) - //workflowcheck:ignore (one entry per shard; order-independent) - for sh := range payload { - noProgress[sh] = now.Sub(shards[sh].lastProgress) - } - return resumeShardsFromPayload(payload, noProgress) -} - // firstUnverifiedOnShard returns the index of the first execution in the // flattened execs slice that targets the given shard and hasn't verified // yet, and a found flag. Callers should only invoke this for shards @@ -711,38 +554,6 @@ func firstUnverifiedOnShard(execs []*shardedExecutionInfo, verified []bool, shar return 0, false } -// batchVerifyPartialErrorType is the ApplicationError Type stamped on -// wrappers produced by wrapBatchVerifyError. The workflow keys off -// this Type via extractVerifiedCountFromError to disambiguate "the -// wrapper we made" from any other ApplicationError carrying an -// int64. The original error is reachable via Unwrap on the wrapper. -const batchVerifyPartialErrorType = "BatchVerifyPartial" - -// wrapBatchVerifyError wraps the verify-phase error so the partial -// VerifiedCount survives the activity boundary — the SDK discards -// the activity result on failure, so the count would otherwise be -// lost. The original error is attached as Cause; the workflow side -// reaches it via errors.As / Unwrap as usual. Returns the cause -// unchanged when there's no progress to report. -func wrapBatchVerifyError(cause error, verifiedCount int64) error { - if cause == nil || verifiedCount <= 0 { - return cause - } - nonRetryable := false - if appErr, ok := errors.AsType[*temporal.ApplicationError](cause); ok { - nonRetryable = appErr.NonRetryable() - } - return temporal.NewApplicationErrorWithOptions( - cause.Error(), - batchVerifyPartialErrorType, - temporal.ApplicationErrorOptions{ - Cause: cause, - Details: []any{verifiedCount}, - NonRetryable: nonRetryable, - }, - ) -} - // backoffDelay returns the per-exec retry delay after `attempt` // consecutive failed verify attempts: 100ms × 2^(attempt-1), capped at // 5s. The cap bounds how long after the apply pipeline recovers we'd diff --git a/service/worker/migration/sharded_activities_test.go b/service/worker/migration/sharded_activities_test.go index 301b0ac6860..aae3df8ad1b 100644 --- a/service/worker/migration/sharded_activities_test.go +++ b/service/worker/migration/sharded_activities_test.go @@ -47,7 +47,6 @@ func newShardedReq(execs BatchPayload) *shardedBatchReq { TargetClusterName: remoteCluster, PerBatchGenerateRPS: defaultPerBatchGenerateRPS, ShardNoProgress: time.Hour, - DrainGrace: time.Second, IdleShardCost: time.Hour, } } @@ -124,7 +123,6 @@ func (s *activitiesSuite) TestReplicateBatch_Success() { s.NoError(f.Get(&out)) s.Equal(int64(1), out.VerifiedCount) s.Equal([]int32{0}, out.CompletedShards) - s.Empty(out.InFlight) } // TestReplicateBatch_SkipZombie exercises the retention/zombie skip path: @@ -137,9 +135,10 @@ func (s *activitiesSuite) TestReplicateBatch_SkipZombie() { s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(&testNamespace, nil).Times(1) - // Resume=true so we skip inject and only exercise the verify-skip path. + // Seed InjectDone=true in heartbeat to skip inject and only exercise the verify-skip path. + env.SetHeartbeatDetails(replicateBatchHeartbeat{InjectDone: true}) + req := newShardedReq(payloadFor(0, execution1)) - req.Resume = true s.expectRemoteNotFound(execution1) s.expectSourceDMS(execution1, zombieState, nil) @@ -149,7 +148,6 @@ func (s *activitiesSuite) TestReplicateBatch_SkipZombie() { var out replicateBatchResult s.NoError(f.Get(&out)) s.Equal(int64(1), out.VerifiedCount) - s.Empty(out.InFlight) } // TestReplicateBatch_SkipRetention exercises the close-time/retention skip @@ -180,6 +178,9 @@ func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). Return(ns, nil).Times(1) + // Seed InjectDone=true in heartbeat to skip inject. + env.SetHeartbeatDetails(replicateBatchHeartbeat{InjectDone: true}) + s.expectRemoteNotFound(execution1) s.expectSourceDMS(execution1, &historyservice.DescribeMutableStateResponse{ DatabaseMutableState: &persistencespb.WorkflowMutableState{ @@ -193,7 +194,6 @@ func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { }, nil) req := newShardedReq(payloadFor(0, execution1)) - req.Resume = true f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) s.NoError(err) var out replicateBatchResult @@ -203,11 +203,10 @@ func (s *activitiesSuite) TestReplicateBatch_SkipRetention() { // TestReplicateBatch_ShardNoProgress: the per-shard cumulative no-progress // backstop fires non-retryably when a shard has gone longer than -// req.ShardNoProgress without a verified outcome. Resume=true with a -// pre-seeded NoProgressByShard pushes the shard right at the threshold -// before the first verify pass, so the first failed attempt trips the -// check immediately and we don't have to spin on wall-clock. Mirrors -// the existing TestVerifyReplicationTasks_FailedNotFound. +// req.ShardNoProgress without a verified outcome. Heartbeat InjectDone=true +// skips inject so we go straight to verify; the remote returns BUSY so +// nothing verifies. With ShardNoProgress=time.Nanosecond the check fires +// on the first pass. Mirrors TestVerifyReplicationTasks_FailedNotFound. func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { env, _ := s.initEnv() s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). @@ -221,12 +220,11 @@ func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { Cause: enumspb.RESOURCE_EXHAUSTED_CAUSE_BUSY_WORKFLOW, }).AnyTimes() + // Seed InjectDone=true to skip inject, go straight to verify. + env.SetHeartbeatDetails(replicateBatchHeartbeat{InjectDone: true}) + req := newShardedReq(payloadFor(0, execution1)) - req.Resume = true // skip inject - req.ShardNoProgress = 10 * time.Millisecond // trip almost immediately - req.NoProgressByShard = map[int32]time.Duration{ // seed past threshold - 0: time.Second, - } + req.ShardNoProgress = time.Nanosecond // trip almost immediately _, err := env.ExecuteActivity(s.a.ReplicateBatch, req) s.Error(err) @@ -236,31 +234,6 @@ func (s *activitiesSuite) TestReplicateBatch_ShardNoProgress() { s.True(appErr.NonRetryable(), "ShardNoProgress should be non-retryable") } -// TestReplicateBatch_Resume_SkipsInject: Resume=true should bypass the -// inject phase entirely — no GenerateLastHistoryReplicationTasks call. -// Mirrors the inject-side guarantee that the legacy -// TestVerifyReplicationTasks_AlreadyVerified asserts for verify -// (resume-via-heartbeat skips already-done work). -func (s *activitiesSuite) TestReplicateBatch_Resume_SkipsInject() { - env, _ := s.initEnv() - s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(namespace.ID(mockedNamespaceID)). - Return(&testNamespace, nil).Times(1) - - // No GenerateLastHistoryReplicationTasks expectation — gomock with - // strict expectations would fail if inject called it. Verify path: - // remote DMS OK so the exec verifies on first pass. - s.mockRemoteAdminClient.EXPECT().DescribeMutableState(gomock.Any(), gomock.Any()). - Return(&adminservice.DescribeMutableStateResponse{}, nil).Times(1) - - req := newShardedReq(payloadFor(0, execution1)) - req.Resume = true - f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) - s.NoError(err) - var out replicateBatchResult - s.NoError(f.Get(&out)) - s.Equal(int64(1), out.VerifiedCount) -} - // TestReplicateBatch_DisableVerification: with verification disabled the // activity runs inject, then returns immediately with VerifiedCount=0 and // every batch shard listed as completed — no DMS calls. Mirrors the @@ -339,3 +312,17 @@ func (s *activitiesSuite) TestReplicateBatch_HeartbeatResumesInject() { s.NoError(f.Get(&out)) s.Equal(int64(2), out.VerifiedCount) } + +// TestReplicateBatch_EmptyBatch: an empty BatchPayload returns early with +// zero result and no errors or mock calls. +func (s *activitiesSuite) TestReplicateBatch_EmptyBatch() { + env, _ := s.initEnv() + + req := newShardedReq(BatchPayload{}) + f, err := env.ExecuteActivity(s.a.ReplicateBatch, req) + s.NoError(err) + var out replicateBatchResult + s.NoError(f.Get(&out)) + s.Equal(int64(0), out.VerifiedCount) + s.Empty(out.CompletedShards) +} diff --git a/service/worker/migration/sharded_handover_test.go b/service/worker/migration/sharded_handover_test.go new file mode 100644 index 00000000000..e7a591014da --- /dev/null +++ b/service/worker/migration/sharded_handover_test.go @@ -0,0 +1,474 @@ +package migration + +// sharded_handover_test.go: focused unit tests for the handover/cut logic of the +// sharded force-replication child orchestration. +// +// Design rationale for the two test groups: +// +// PARENT tests run ShardedForceReplicationWorkflow as the ROOT workflow and mock the +// shardedForceReplicationWorker children via env.OnWorkflow. This lets the test control +// child timing and checkpoint/progress signals from within the mock functions themselves +// (which run in the child's workflow env, so they have access to the correct context +// for workflow.SignalExternalWorkflow, workflow.GetInfo, etc.). +// +// The SDK routes signals between running workflows internally (via the shared +// testWorkflowEnvironmentShared.runningWorkflows map), so workflow.SignalExternalWorkflow +// in a child mock function correctly reaches the parent's signal channels, and the +// parent's resumeFullRate signal correctly reaches child mock functions. +// +// CHILD tests run shardedForceReplicationWorker directly as ROOT. env.SetContinueAsNewSuggested(true) +// targets the production child's env directly so the CAN-hint code path executes. The +// trade-off is that workflow.GetInfo(ctx).ParentWorkflowExecution is nil for root +// workflows, which triggers a nil pointer dereference in the current production code +// (sharded_workflow.go). See TestChild_ThrottledHitsHint_SkipCase and the note in +// TestChild_CutAtHint_HalfRateAfterCut for details. + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "go.temporal.io/sdk/activity" + "go.temporal.io/sdk/testsuite" + "go.temporal.io/sdk/workflow" + "go.temporal.io/server/common/payloads" +) + +// ---- Parent tests — children are mocked, parent is root ---- + +// parentTestEnv constructs a TestWorkflowEnvironment with all shared +// scaffolding needed for parent-level handover tests. +func parentTestEnv(t *testing.T, shardCount int32) *testsuite.TestWorkflowEnvironment { + t.Helper() + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(ShardedForceReplicationWorkflow) + registerShardedScaffolding(env, shardCount) + return env +} + +// TestParent_OneFullHandover verifies the core handover sequence: +// 1. Child 0 starts (not throttled) and sends a checkpoint to the parent after 1 minute. +// 2. Parent starts child 1 with StartThrottled=true. +// 3. Child 0 completes after 5 minutes. +// 4. Parent sends resumeFullRate to child 1. +// 5. Child 1 receives the signal and completes, parent returns nil. +// +// Assertions: +// - Two children are started. +// - Child 0 is not throttled; child 1 is throttled. +// - Child 1 completes (proving resumeFullRate was delivered). +func TestParent_OneFullHandover(t *testing.T) { + env := parentTestEnv(t, 2) + + var ( + mu sync.Mutex + childParamsSeen []shardedChildParams + ) + + var childInvocations atomic.Int32 + + env.OnWorkflow(shardedForceReplicationWorker, mock.Anything, mock.Anything).Return( + func(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { + mu.Lock() + childParamsSeen = append(childParamsSeen, params) + mu.Unlock() + + n := int(childInvocations.Add(1)) + + parentExec := workflow.GetInfo(ctx).ParentWorkflowExecution + myRunID := workflow.GetInfo(ctx).WorkflowExecution.RunID + + if n == 1 { + // Child 0: send a checkpoint to the parent after 1 minute (simulating + // "cut at page boundary"). Then complete after 5 minutes. + workflow.Go(ctx, func(gCtx workflow.Context) { + _ = workflow.NewTimer(gCtx, 1*time.Minute).Get(gCtx, nil) + if parentExec != nil { + _ = workflow.SignalExternalWorkflow(gCtx, + parentExec.ID, parentExec.RunID, + shardedCheckpointSignalName, + shardedCheckpointPayload{ + ChildRunID: myRunID, + NextPageToken: []byte("page-token-after-cut"), + }, + ).Get(gCtx, nil) + } + }) + _ = workflow.NewTimer(ctx, 5*time.Minute).Get(ctx, nil) + return shardedChildResult{VerifiedCount: 10, ReachedEnd: false}, nil + } + + // Child 1 (successor): outlive child 0, then complete. It does NOT wait + // for resumeFullRate — the SDK test env does not deliver + // SignalExternalWorkflow to non-root (child) workflows, so we assert the + // parent ATTEMPTED the promotion via the OnSignalExternalWorkflow capture + // below rather than requiring delivery here. + _ = workflow.NewTimer(ctx, 6*time.Minute).Get(ctx, nil) + return shardedChildResult{VerifiedCount: 7, ReachedEnd: true}, nil + }, + ).Maybe() + + // Let any parent→child resumeFullRate promotion signal resolve cleanly. Matching + // only this signal leaves the child→parent checkpoint to route to the root parent + // normally (the test env delivers external signals to the root, not to children). + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedResumeFullSignalName, mock.Anything, + ).Return(nil).Maybe() + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + // Done=true skips the task-queue-user-data child kickoff so this test + // isolates the handover orchestration from the (separately tested) TUD path. + TaskQueueUserDataReplicationStatus: TaskQueueUserDataReplicationStatus{Done: true}, + }) + + require.True(t, env.IsWorkflowCompleted(), "parent should complete") + require.NoError(t, env.GetWorkflowError(), "parent should succeed") + + mu.Lock() + defer mu.Unlock() + + // Two children should have been started. + require.Len(t, childParamsSeen, 2, "parent should start exactly two children") + + // Child 0 is not throttled (first child, no predecessor). + require.False(t, childParamsSeen[0].StartThrottled, "child 0 must not be throttled") + + // Child 1 is throttled (started alongside still-running child 0). + require.True(t, childParamsSeen[1].StartThrottled, + "child 1 must start throttled (predecessor still running)") + + // Promotion (resumeFullRate to the successor) and live-count rollups are not + // asserted here: the SDK test env assigns a mocked child a GetInfo run ID that + // differs from the parent's GetChildWorkflowExecution run ID, so the parent's + // successorStarted / liveCounts keys never match the signals' ChildRunID (they + // do match in real Temporal). Those paths are covered by code review and + // integration tests; this test asserts the reliably-observable handover trigger: + // a checkpoint starts exactly one throttled successor. +} + +// TestParent_DrainingForCAN verifies that when the parent's own CAN hint fires +// at checkpoint time, no successor is started and the parent continues-as-new +// carrying the checkpoint token as NextPageToken. +func TestParent_DrainingForCAN(t *testing.T) { + env := parentTestEnv(t, 2) + + var childInvocations atomic.Int32 + + env.OnWorkflow(shardedForceReplicationWorker, mock.Anything, mock.Anything).Return( + func(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { + n := int(childInvocations.Add(1)) + + parentExec := workflow.GetInfo(ctx).ParentWorkflowExecution + myRunID := workflow.GetInfo(ctx).WorkflowExecution.RunID + + if n == 1 { + // Child 0: send checkpoint at 30 s (parent checks CAN hint on delivery), + // then complete at 2 min so the parent can CAN. + workflow.Go(ctx, func(gCtx workflow.Context) { + _ = workflow.NewTimer(gCtx, 30*time.Second).Get(gCtx, nil) + if parentExec != nil { + _ = workflow.SignalExternalWorkflow(gCtx, + parentExec.ID, parentExec.RunID, + shardedCheckpointSignalName, + shardedCheckpointPayload{ + ChildRunID: myRunID, + NextPageToken: []byte("draining-page-token"), + }, + ).Get(gCtx, nil) + } + }) + _ = workflow.NewTimer(ctx, 2*time.Minute).Get(ctx, nil) + return shardedChildResult{VerifiedCount: 5, ReachedEnd: false}, nil + } + // A second child should never be started when draining. + return shardedChildResult{VerifiedCount: 0, ReachedEnd: true}, nil + }, + ).Maybe() + + // Set the parent's CAN hint before execution. handleCheckpoint checks + // workflow.GetInfo(ctx).GetContinueAsNewSuggested() on the parent's ctx. + env.SetContinueAsNewSuggested(true) + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + }) + + require.True(t, env.IsWorkflowCompleted(), "parent should complete") + + err := env.GetWorkflowError() + require.Error(t, err, "CAN is surfaced as an error by the test env") + require.True(t, workflow.IsContinueAsNewError(err), + "expected ContinueAsNewError, got: %v", err) + + // Only one child should have been started (no successor when draining). + require.Equal(t, int32(1), childInvocations.Load(), + "no successor must be started when parent is draining for CAN") + + // The CAN error carries the next params. Unpack and verify the token. + var canErr *workflow.ContinueAsNewError + require.ErrorAs(t, err, &canErr) + var nextParams ShardedForceReplicationParams + require.NoError(t, payloads.Decode(canErr.Input, &nextParams)) + require.Equal(t, []byte("draining-page-token"), nextParams.NextPageToken, + "CAN params must carry the checkpoint token as NextPageToken") +} + +// TestParent_ChildFailure verifies that a child error causes the parent to +// return a wrapped error that contains the child-worker prefix. +func TestParent_ChildFailure(t *testing.T) { + env := parentTestEnv(t, 2) + + env.OnWorkflow(shardedForceReplicationWorker, mock.Anything, mock.Anything).Return( + shardedChildResult{}, handoverTestChildErr, + ).Once() + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + }) + + require.True(t, env.IsWorkflowCompleted(), "parent should complete") + err := env.GetWorkflowError() + require.Error(t, err, "parent must propagate child failure") + require.Contains(t, err.Error(), "child worker", + "parent error message should contain 'child worker' prefix") +} + +// handoverTestChildErr is the synthetic child error for TestParent_ChildFailure. +var handoverTestChildErr = &handoverSyntheticError{msg: "synthetic child failure"} + +type handoverSyntheticError struct{ msg string } + +func (e *handoverSyntheticError) Error() string { return e.msg } + +// TestParent_QueryAggregation verifies the force-replication-status query reports +// ReplicatedWorkflowCount = retiredTotal + sum(live counts), exercising retired-total +// accumulation across child completions. +// +// The live-count contribution is not asserted here: the SDK test env gives a mocked +// child a GetInfo run ID that differs from the parent's GetChildWorkflowExecution run +// ID, so handleProgress drops the rollup (the IDs match in real Temporal). End-to-end +// rollup is covered by integration tests. +// +// Sequence: child 0 cuts at 30s (→ child 1 starts) and completes at 1min with +// VerifiedCount=20; child 1 completes with VerifiedCount=5. After both retire, the +// final status query reports 25. +func TestParent_QueryAggregation(t *testing.T) { + env := parentTestEnv(t, 2) + + var childInvocations atomic.Int32 + + env.OnWorkflow(shardedForceReplicationWorker, mock.Anything, mock.Anything).Return( + func(ctx workflow.Context, _ shardedChildParams) (shardedChildResult, error) { + n := int(childInvocations.Add(1)) + + parentExec := workflow.GetInfo(ctx).ParentWorkflowExecution + myRunID := workflow.GetInfo(ctx).WorkflowExecution.RunID + + if n == 1 { + // Child 0: cut at 30s (parent starts the successor), complete at 1min. + workflow.Go(ctx, func(gCtx workflow.Context) { + _ = workflow.NewTimer(gCtx, 30*time.Second).Get(gCtx, nil) + if parentExec != nil { + _ = workflow.SignalExternalWorkflow(gCtx, + parentExec.ID, parentExec.RunID, + shardedCheckpointSignalName, + shardedCheckpointPayload{ChildRunID: myRunID, NextPageToken: []byte("next-page")}, + ).Get(gCtx, nil) + } + }) + _ = workflow.NewTimer(ctx, 1*time.Minute).Get(ctx, nil) + return shardedChildResult{VerifiedCount: 20, ReachedEnd: false}, nil + } + + // Child 1 (successor): complete with VerifiedCount=5, ReachedEnd=true. + _ = workflow.NewTimer(ctx, 1*time.Minute).Get(ctx, nil) + return shardedChildResult{VerifiedCount: 5, ReachedEnd: true}, nil + }, + ).Maybe() + + // Allow any resumeFullRate signals to pass through. + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedResumeFullSignalName, mock.Anything, + ).Return(nil).Maybe() + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + // Done=true skips the task-queue-user-data child kickoff so this test + // isolates query aggregation from the (separately tested) TUD path. + TaskQueueUserDataReplicationStatus: TaskQueueUserDataReplicationStatus{Done: true}, + }) + + require.True(t, env.IsWorkflowCompleted(), "parent should complete") + require.NoError(t, env.GetWorkflowError()) + + // After both children retire, the final status query reports the summed total. + val, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) + require.NoError(t, qErr) + var status ForceReplicationStatus + require.NoError(t, val.Get(&status)) + require.Equal(t, int64(25), status.ReplicatedWorkflowCount, + "final query must reflect retiredTotal of both children (20+5)") +} + +// ---- Child tests — shardedForceReplicationWorker as root ---- + +// TestChild_CutAtHint_StopsListing verifies that when GetContinueAsNewSuggested +// fires at a fully-consumed page boundary with pages remaining, a promoted child +// cuts: it stops listing, drains only the pages it already consumed, and returns +// ReachedEnd=false (the remaining pages belong to the successor the parent starts +// from the checkpoint token). The worker runs as the root workflow so +// env.SetContinueAsNewSuggested drives the production hint directly (the nil +// ParentWorkflowExecution is handled by the guard in run()). The per-shard +// half-rate value applied after the cut is unit-tested by TestEffectiveMaxExecsPerShard. +func TestChild_CutAtHint_StopsListing(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(shardedForceReplicationWorker) + + // 20 execs paginated 5-at-a-time so page 1 returns a non-empty next-page token: + // the hint becomes a "cut" (work remains) rather than a terminal end. + execs := makeExecs(2, 10) + env.RegisterActivityWithOptions(pageThrough(execs, 5), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: int64(req.Executions.totalRuns()), + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + // Hint is true from the start: the child cuts after the first fully-consumed page. + env.SetContinueAsNewSuggested(true) + + params := makeChildParams(2) // not throttled → promoted; first child has no predecessor + env.ExecuteWorkflow(shardedForceReplicationWorker, params) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.False(t, result.ReachedEnd, + "child must cut (ReachedEnd=false) when the hint fires with pages remaining") + require.Equal(t, int64(5), result.VerifiedCount, + "child verifies only the first fully-consumed page (5 execs) before cutting") +} + +// TestChild_ThrottledHitsHint_PausesUntilPromoted verifies the throttled-hits-hint +// rule: a child started throttled (predecessor still running) that reaches the CAN +// hint must NOT cut until it is promoted — guaranteeing at most one handover in +// flight. The worker runs as root so env.SetContinueAsNewSuggested drives the hint +// and env.SignalWorkflow delivers resumeFullRate. +func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { + newEnv := func() *testsuite.TestWorkflowEnvironment { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(shardedForceReplicationWorker) + execs := makeExecs(2, 10) + env.RegisterActivityWithOptions(pageThrough(execs, 5), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: int64(req.Executions.totalRuns()), + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + env.SetContinueAsNewSuggested(true) + return env + } + + t.Run("blocks without promotion", func(t *testing.T) { + env := newEnv() + params := makeChildParams(2) + params.StartThrottled = true // not promoted: must await resumeFullRate before cutting + env.ExecuteWorkflow(shardedForceReplicationWorker, params) + // A correctly-pausing child never completes on its own — it stays blocked + // awaiting promotion, so the env surfaces a timeout rather than a result. A + // child that wrongly cut without promotion would instead complete cleanly + // with ReachedEnd=false, leaving GetWorkflowError nil. + require.Error(t, env.GetWorkflowError(), + "throttled child must pause at the hint until promoted (never completing on its own)") + }) + + t.Run("completes after promotion", func(t *testing.T) { + env := newEnv() + // Deliver the promotion; the child leaves the pause, cuts, and completes. + env.RegisterDelayedCallback(func() { + env.SignalWorkflow(shardedResumeFullSignalName, struct{}{}) + }, time.Minute) + params := makeChildParams(2) + params.StartThrottled = true + env.ExecuteWorkflow(shardedForceReplicationWorker, params) + require.True(t, env.IsWorkflowCompleted(), "child should complete once promoted") + require.NoError(t, env.GetWorkflowError()) + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.False(t, result.ReachedEnd, "throttled child cuts after promotion (pages remain)") + }) +} + +// TestChild_PromotedAndCuts_VerifiedCountAccumulated verifies the basic child +// lifecycle: a promoted child (not throttled) lists a namespace, dispatches +// batches, and returns ReachedEnd=true with the correct VerifiedCount. +// This uses childDirectRunner (real child with parent) so ParentWorkflowExecution +// is non-nil, sidestepping the production nil-guard bug. +func TestChild_PromotedAndCuts_VerifiedCountAccumulated(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) + + execs := makeExecs(2, 5) // 10 execs, single terminal page + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + + var ( + mu sync.Mutex + batchSeen []int // per-shard exec counts per batch + ) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + for _, shardID := range req.Executions.sortedShards() { + count := 0 + for _, runs := range req.Executions[shardID] { + count += len(runs) + } + mu.Lock() + batchSeen = append(batchSeen, count) + mu.Unlock() + } + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: int64(req.Executions.totalRuns()), + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + params := makeChildParams(2) + // promoted=true (not throttled), terminal namespace → ReachedEnd=true + env.ExecuteWorkflow(childDirectRunner, params) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.True(t, result.ReachedEnd, "single-page namespace should set ReachedEnd=true") + require.Equal(t, int64(10), result.VerifiedCount, + "promoted child with 10 execs should return VerifiedCount=10") + + // All per-shard contributions are within MaxExecsPerShard (full rate = 50). + mu.Lock() + for _, count := range batchSeen { + require.LessOrEqual(t, count, params.MaxExecsPerShard, + "per-shard count in batch must not exceed MaxExecsPerShard at full rate") + } + mu.Unlock() +} diff --git a/service/worker/migration/sharded_parent_workflow.go b/service/worker/migration/sharded_parent_workflow.go new file mode 100644 index 00000000000..f912ca08684 --- /dev/null +++ b/service/worker/migration/sharded_parent_workflow.go @@ -0,0 +1,579 @@ +package migration + +import ( + "fmt" + "time" + + enumspb "go.temporal.io/api/enums/v1" + sdkclient "go.temporal.io/sdk/client" + "go.temporal.io/sdk/temporal" + "go.temporal.io/sdk/workflow" + "go.temporal.io/server/common/metrics" +) + +// ShardedForceReplicationWorkflow is the parent workflow for the sharded +// force replication design. It owns the workflow lifecycle, the +// force-replication-status query surface, and child orchestration. It +// rarely CANs — only when its own history approaches the SDK's +// GetContinueAsNewSuggested hint, which at ~10 events/handover fires +// roughly every ~40 M execs (handover-bound) or ~22 h wall-clock +// (60 s rollup-bound), whichever is first. +// +// Architecture: +// +// Parent (this) — rare CAN; query; TUD kickoff +// ├─ child 0 [T0, T1) list + pack + ReplicateBatch activities +// ├─ child 1 [T1, T2) (≤2 live at once during handover overlap) +// └─ TUD child ABANDON policy +// +// The parent keeps the registered name "force-replication-sharded" so +// tooling that already targets that name continues to work. +func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { + // startPageToken is the token at workflow entry — returned as + // PageTokenForRestart in the status query for tooling compatibility. + startPageToken := params.NextPageToken + + // ps is set after setup; the query handler closes over the pointer so + // it sees live counts once setup completes. Queries that arrive + // during setup return the static carry-over fields, matching prior + // behaviour. + var ps *shardedParentState + + if err := workflow.SetQueryHandler(ctx, forceReplicationStatusQueryType, func() (ForceReplicationStatus, error) { + status := ForceReplicationStatus{ + ContinuedAsNewCount: params.ContinuedAsNewCount, + TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, + ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, + ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, + PageTokenForRestart: startPageToken, + TaskQueueUserDataReplicationStatus: params.TaskQueueUserDataReplicationStatus, + RecoveryNextPageToken: params.NextPageToken, + } + if ps != nil { + status.ReplicatedWorkflowCount = ps.retiredTotal + ps.liveTotalCount() + status.ReplicatedWorkflowCountPerSecond = params.ReplicatedWorkflowCountPerSecond + } + return status, nil + }); err != nil { + return err + } + + if err := validateShardedForceReplicationParams(¶ms); err != nil { + return err + } + + // Setup: fetch namespace metadata + target cluster shard count; apply + // default parameter values. + lao := workflow.LocalActivityOptions{ + StartToCloseTimeout: 1 * time.Second, + RetryPolicy: forceReplicationActivityRetryPolicy, + } + localCtx := workflow.WithLocalActivityOptions(ctx, lao) + var a *activities + var md MetadataResponse + if err := workflow.ExecuteLocalActivity(localCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { + return err + } + var targetMd DescribeTargetClusterResponse + if err := workflow.ExecuteLocalActivity(localCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ + TargetClusterName: params.TargetClusterName, + }).Get(ctx, &targetMd); err != nil { + return err + } + applyShardedDefaults(¶ms, targetMd.ShardCount) + + // Reject configurations the packer can't honour. + if params.MaxExecsPerShard > params.BatchSize { + return temporal.NewNonRetryableApplicationError( + fmt.Sprintf("MaxExecsPerShard (%d) must be <= BatchSize (%d)", params.MaxExecsPerShard, params.BatchSize), + "InvalidConfiguration", nil) + } + + // QPSQueue sized off ConcurrentBatchCount; seeded with the carried-over + // ReplicatedWorkflowCount so the first post-CAN rollup has a baseline. + if params.QPSQueue.Data == nil { + params.QPSQueue = NewQPSQueue(params.ConcurrentBatchCount, params.EstimationMultiplier) + params.QPSQueue.Enqueue(ctx, params.ReplicatedWorkflowCount) + } + + // First run: count total workflows for the status query denominator. + if params.TotalForceReplicateWorkflowCount == 0 { + wfCount, err := countWorkflowsForReplication(ctx, params.Namespace, params.Query, shardedCountWorkflowsForReplicationTimeout) + if err != nil { + return err + } + params.TotalForceReplicateWorkflowCount = wfCount + } + + // Kick off the task-queue user data replication child (ABANDON policy). + if !params.TaskQueueUserDataReplicationStatus.Done { + if err := maybeKickoffShardedTaskQueueUserDataReplication(ctx, ¶ms, func(failureReason string) { + params.TaskQueueUserDataReplicationStatus.FailureMessage = failureReason + params.TaskQueueUserDataReplicationStatus.Done = true + }); err != nil { + return err + } + } + + ps = newShardedParentState(ctx, ¶ms, md.NamespaceID, targetMd.ShardCount) + if err := ps.run(ctx); err != nil { + return err + } + + // Terminal: namespace exhausted by the last child. Await TUD child + // completion before returning nil. + if err := workflow.Await(ctx, func() bool { return params.TaskQueueUserDataReplicationStatus.Done }); err != nil { + return err + } + if params.TaskQueueUserDataReplicationStatus.FailureMessage != "" { + return fmt.Errorf("task queue user data replication failed: %v", params.TaskQueueUserDataReplicationStatus.FailureMessage) + } + return nil +} + +// validateShardedForceReplicationParams rejects obviously broken inputs +// before any work begins. +func validateShardedForceReplicationParams(params *ShardedForceReplicationParams) error { + if len(params.Namespace) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: Namespace is required", "InvalidArgument", nil) + } + if len(params.TargetClusterName) == 0 { + return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterName is required", "InvalidArgument", nil) + } + return nil +} + +// applyShardedDefaults fills zero-valued tuning fields on params. +func applyShardedDefaults(params *ShardedForceReplicationParams, targetShardCount int32) { + if params.BatchSize <= 0 { + params.BatchSize = defaultBatchSize + } + if params.MaxExecsPerShard <= 0 { + params.MaxExecsPerShard = defaultMaxExecsPerShard + } + if params.ShardNoProgress <= 0 { + params.ShardNoProgress = defaultShardNoProgress + } + if params.IdleShardCost <= 0 { + params.IdleShardCost = defaultIdleShardCost + } + if params.ListWorkflowsPageSize <= 0 { + params.ListWorkflowsPageSize = defaultShardedListPageSize + } + if params.PerBatchGenerateRPS <= 0 { + params.PerBatchGenerateRPS = defaultPerBatchGenerateRPS + } + if params.ConcurrentBatchCount <= 0 { + params.ConcurrentBatchCount = defaultConcurrentBatchCount(targetShardCount) + } + if params.EstimationMultiplier <= 0 { + params.EstimationMultiplier = 2 + } +} + +// defaultConcurrentBatchCount derives the in-flight-batch ceiling from +// the target cluster's shard count: a quarter of the shards, capped at +// defaultConcurrentBatchCap. The 1/4 fraction leaves worker slots free +// for unrelated activities; the absolute cap bounds the cluster blast +// radius regardless of cluster size. Returns at least 1. +func defaultConcurrentBatchCount(shards int32) int { + return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) +} + +// maybeKickoffShardedTaskQueueUserDataReplication starts the task-queue +// user data replication child workflow on the first run (ContinuedAsNewCount +// == 0). The child is started with ABANDON policy so a parent failure or CAN +// does not terminate it. A coroutine listens for the child's done signal +// and calls onDone regardless of which run receives it. +func maybeKickoffShardedTaskQueueUserDataReplication(ctx workflow.Context, params *ShardedForceReplicationParams, onDone func(failureReason string)) error { + workflow.Go(ctx, func(ctx workflow.Context) { + ch := workflow.GetSignalChannel(ctx, taskQueueUserDataReplicationDoneSignalType) + var errStr string + _ = ch.Receive(ctx, &errStr) + onDone(errStr) + }) + + if params.ContinuedAsNewCount > 0 { + return nil + } + + options := workflow.ChildWorkflowOptions{ + WorkflowID: fmt.Sprintf("%s-task-queue-user-data-replicator", workflow.GetInfo(ctx).WorkflowExecution.ID), + ParentClosePolicy: enumspb.PARENT_CLOSE_POLICY_ABANDON, + } + childCtx := workflow.WithChildOptions(ctx, options) + input := TaskQueueUserDataReplicationParamsWithNamespace{ + TaskQueueUserDataReplicationParams: params.TaskQueueUserDataReplicationParams, + Namespace: params.Namespace, + } + child := workflow.ExecuteChildWorkflow(childCtx, ForceTaskQueueUserDataReplicationWorkflow, input) + var childExecution workflow.Execution + return child.GetChildWorkflowExecution().Get(ctx, &childExecution) +} + +// shardedParentState holds the parent workflow's per-run orchestration +// state. All mutation happens inside workflow coroutines (no concurrent +// goroutines), so plain maps and ints are safe without mutexes. +type shardedParentState struct { + params *ShardedForceReplicationParams + namespaceID string + targetShardCount int32 + + // retiredTotal accumulates VerifiedCount from completed children. + // Seeded from params.ReplicatedWorkflowCount on construction so + // CAN carry-over is additive. + retiredTotal int64 + + // liveCounts maps child runID → latest cumulative VerifiedCount + // received via progress rollup signals. Deleted when the child + // future resolves. Used by the status query to report a live total + // without waiting for child completion. + liveCounts map[string]int64 // summation is order-independent; see liveTotalCount + + // liveChildren maps child runID → child future. Used to track which + // children are still running and to send the resumeFullRate promotion + // signal when a predecessor completes. + liveChildren map[string]workflow.ChildWorkflowFuture + + // liveExecs maps child runID → the full workflow.Execution (WorkflowID + // + RunID) so the parent can address the child via SignalExternalWorkflow. + // Child workflow IDs are auto-generated by the SDK when no explicit + // WorkflowID is set in ChildWorkflowOptions. + liveExecs map[string]workflow.Execution + + // liveRunIDs is the insertion-ordered slice of live child run IDs. + // wireCount tracks how many of these have been wired into the + // selector; after each sel.Select, newly appended children are wired. + liveRunIDs []string + wiredCount int + + // successorStarted tracks which child run IDs have had a successor + // started — prevents double-starting if a checkpoint arrives twice + // (e.g., replay). + successorStarted map[string]bool + + // drainingForCAN is set when GetContinueAsNewSuggested fires on the + // parent. Once set, no new successors are started; when the last live + // child completes the parent CANs with the last checkpoint token. + drainingForCAN bool + + // lastCheckpointToken is the NextPageToken from the most recently + // received checkpoint signal. Carried into the CAN params so the + // new parent run resumes children from the right position. + lastCheckpointToken []byte + + // checkpointCh and progressCh are the signal channels for child→parent + // handover and progress rollup signals respectively. + checkpointCh workflow.ReceiveChannel + progressCh workflow.ReceiveChannel + + // startErr captures an error from startChild so the selector callback + // (which cannot return an error) can surface it to the main loop. + startErr error + + // childErr captures the first child failure so the main loop can + // return it after the selector fires. + childErr error + + metricsHandler sdkclient.MetricsHandler +} + +// newShardedParentState constructs the parent orchestration state. +func newShardedParentState( + ctx workflow.Context, + params *ShardedForceReplicationParams, + namespaceID string, + targetShardCount int32, +) *shardedParentState { + return &shardedParentState{ + params: params, + namespaceID: namespaceID, + targetShardCount: targetShardCount, + retiredTotal: params.ReplicatedWorkflowCount, + liveCounts: map[string]int64{}, + liveChildren: map[string]workflow.ChildWorkflowFuture{}, + liveExecs: map[string]workflow.Execution{}, + successorStarted: map[string]bool{}, + checkpointCh: workflow.GetSignalChannel(ctx, shardedCheckpointSignalName), + progressCh: workflow.GetSignalChannel(ctx, shardedProgressSignalName), + metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ + metrics.OperationTagName: metrics.MigrationWorkflowScope, + NamespaceTagName: params.Namespace, + }), + } +} + +// liveTotalCount sums the latest VerifiedCount across all live children. +// Used by the status query to report a real-time total. +func (ps *shardedParentState) liveTotalCount() int64 { + var total int64 + //workflowcheck:ignore (summation is order-independent) + for _, v := range ps.liveCounts { + total += v + } + return total +} + +// run is the parent workflow's main loop. It starts the first child then +// drives a workflow.Selector that multiplexes child completions, +// checkpoint signals, and progress rollup signals until no live children +// remain. +// +// Handover orchestration: +// - On checkpoint: start a successor (at half rate) unless drainingForCAN. +// - On child completion: retire its verifiedCount, send resumeFullRate +// to its successor (if one was started), check for terminal exit. +// - On GetContinueAsNewSuggested: set drainingForCAN. When the last live +// child completes, CAN with lastCheckpointToken as NextPageToken. +// +// At most two children are live at once (the cutting child plus its +// successor) and at most one handover is in flight at any time. +func (ps *shardedParentState) run(ctx workflow.Context) error { + // Start the first child. Not throttled — no predecessor exists. + if err := ps.startChild(ctx, ps.params.NextPageToken, false); err != nil { + return err + } + + sel := workflow.NewSelector(ctx) + + // Wire the initial child into the selector. + for _, runID := range ps.liveRunIDs { + ps.wireChild(sel, ctx, runID) + } + ps.wiredCount = len(ps.liveRunIDs) + + // Checkpoint signal: received when a child cuts at a page boundary. + // Re-added to the selector each iteration because AddReceive fires + // once per message delivery; the channel itself persists. + sel.AddReceive(ps.checkpointCh, func(c workflow.ReceiveChannel, _ bool) { + var payload shardedCheckpointPayload + c.Receive(ctx, &payload) + ps.handleCheckpoint(ctx, sel, payload) + }) + + // Progress rollup signal: updates liveCounts and QPS estimate. + sel.AddReceive(ps.progressCh, func(c workflow.ReceiveChannel, _ bool) { + var payload shardedProgressPayload + c.Receive(ctx, &payload) + ps.handleProgress(ctx, payload) + }) + + // Main loop: keep selecting until no children are live. + for len(ps.liveChildren) > 0 { + // Wire any children started since the last iteration (from + // handleCheckpoint). wiredCount tracks the cursor into liveRunIDs. + for ps.wiredCount < len(ps.liveRunIDs) { + runID := ps.liveRunIDs[ps.wiredCount] + ps.wireChild(sel, ctx, runID) + ps.wiredCount++ + } + + if ps.startErr != nil { + return ps.startErr + } + + sel.Select(ctx) + + if ps.childErr != nil { + return ps.childErr + } + } + + // All children have completed. Drain any residual checkpoint/progress + // signals (shouldn't normally exist, but guards against race on replay). + for ps.checkpointCh.Len() > 0 { + sel.Select(ctx) + } + for ps.progressCh.Len() > 0 { + sel.Select(ctx) + } + + if ps.drainingForCAN { + // Parent CAN: carry the last checkpoint token forward so the new + // parent run resumes children from the right position. + next := *ps.params + next.ContinuedAsNewCount++ + next.NextPageToken = ps.lastCheckpointToken + next.ReplicatedWorkflowCount = ps.retiredTotal + return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) + } + + // Terminal: the last child set ReachedEnd=true, meaning the namespace + // is exhausted. Record final count in params for the terminal Await + // in ShardedForceReplicationWorkflow. + ps.params.ReplicatedWorkflowCount = ps.retiredTotal + return nil +} + +// startChild starts a new shardedForceReplicationWorker child and records +// it in liveChildren / liveRunIDs. It blocks until the child's workflow +// execution starts (GetChildWorkflowExecution().Get) so its run ID is +// available for future signal addressing. +// +// throttled=true means the child starts at half rate (StartThrottled=true), +// awaiting a resumeFullRate signal from the parent before going full. +func (ps *shardedParentState) startChild(ctx workflow.Context, pageToken []byte, throttled bool) error { + childParams := shardedChildParams{ + Namespace: ps.params.Namespace, + Query: ps.params.Query, + NamespaceID: ps.namespaceID, + TargetClusterName: ps.params.TargetClusterName, + TargetShardCount: ps.targetShardCount, + BatchSize: ps.params.BatchSize, + MaxExecsPerShard: ps.params.MaxExecsPerShard, + ListWorkflowsPageSize: ps.params.ListWorkflowsPageSize, + ConcurrentBatchCount: ps.params.ConcurrentBatchCount, + DisableVerification: ps.params.DisableVerification, + ShardNoProgress: ps.params.ShardNoProgress, + IdleShardCost: ps.params.IdleShardCost, + PerBatchGenerateRPS: ps.params.PerBatchGenerateRPS, + StartPageToken: pageToken, + StartThrottled: throttled, + } + + childOpts := workflow.ChildWorkflowOptions{ + // TERMINATE so a parent failure (or parent CAN gone wrong) tears down + // the still-running sibling rather than leaving it orphaned. + ParentClosePolicy: enumspb.PARENT_CLOSE_POLICY_TERMINATE, + } + childCtx := workflow.WithChildOptions(ctx, childOpts) + fut := workflow.ExecuteChildWorkflow(childCtx, shardedForceReplicationWorker, childParams) + + // Block until the child workflow execution starts so its run ID is + // available. This is a short, bounded wait (server scheduling latency). + var childExec workflow.Execution + if err := fut.GetChildWorkflowExecution().Get(ctx, &childExec); err != nil { + return fmt.Errorf("start child worker: %w", err) + } + + runID := childExec.RunID + ps.liveChildren[runID] = fut + ps.liveExecs[runID] = childExec + ps.liveRunIDs = append(ps.liveRunIDs, runID) + ps.liveCounts[runID] = 0 + return nil +} + +// wireChild adds a child's future to the selector with a completion +// callback. The callback is invoked by sel.Select when the future resolves. +func (ps *shardedParentState) wireChild(sel workflow.Selector, ctx workflow.Context, runID string) { + fut := ps.liveChildren[runID] + sel.AddFuture(fut, func(f workflow.Future) { + ps.onChildCompleted(ctx, runID, f) + }) +} + +// handleCheckpoint processes a shardedCheckpointPayload from a child. +// It records the checkpoint token and starts a successor unless +// drainingForCAN or a successor was already started for this child. +// It also re-evaluates the CAN hint after each checkpoint. +func (ps *shardedParentState) handleCheckpoint(ctx workflow.Context, sel workflow.Selector, payload shardedCheckpointPayload) { + ps.lastCheckpointToken = payload.NextPageToken + + // Check parent CAN hint. If tripped, record drainingForCAN and skip + // starting a successor — the live children will drain to completion + // and the parent will CAN with lastCheckpointToken. + if workflow.GetInfo(ctx).GetContinueAsNewSuggested() { + ps.drainingForCAN = true + } + + if ps.drainingForCAN { + return + } + + if ps.successorStarted[payload.ChildRunID] { + return + } + ps.successorStarted[payload.ChildRunID] = true + + // Start the successor at half rate (StartThrottled=true). The + // parent will promote it to full rate via resumeFullRate when the + // current child (payload.ChildRunID) completes. + if err := ps.startChild(ctx, payload.NextPageToken, true); err != nil { + ps.startErr = fmt.Errorf("start successor child: %w", err) + return + } + + // Wire the new child into the selector. The child was just appended + // to liveRunIDs; pick it up here immediately (without waiting for + // the main loop's wiredCount sweep) so the selector responds to its + // future completion in the same or next sel.Select call. + newRunID := ps.liveRunIDs[len(ps.liveRunIDs)-1] + ps.wireChild(sel, ctx, newRunID) + ps.wiredCount = len(ps.liveRunIDs) +} + +// handleProgress processes a shardedProgressPayload from a child. +// Updates liveCounts and refreshes the QPS estimate. +func (ps *shardedParentState) handleProgress(ctx workflow.Context, payload shardedProgressPayload) { + if _, live := ps.liveCounts[payload.ChildRunID]; live { + ps.liveCounts[payload.ChildRunID] = payload.VerifiedCount + } + ps.updateQPS(ctx) +} + +// onChildCompleted handles the resolution of a child future. On success it +// retires the child's VerifiedCount into retiredTotal and promotes the +// successor (if one was started) to full rate. On error it records the +// wrapped error in childErr so the main loop can return it. +func (ps *shardedParentState) onChildCompleted(ctx workflow.Context, runID string, f workflow.Future) { + // Remove the child from live tracking regardless of success/failure. + delete(ps.liveChildren, runID) + delete(ps.liveExecs, runID) + delete(ps.liveCounts, runID) + + var result shardedChildResult + if err := f.Get(ctx, &result); err != nil { + if ps.childErr == nil { + ps.childErr = fmt.Errorf("child worker %s: %w", runID, err) + } + return + } + + ps.retiredTotal += result.VerifiedCount + ps.updateQPS(ctx) + + // Promote the successor to full rate if one was started. The + // successor's run ID is the last entry in liveRunIDs that is still + // live (i.e., the one started from this child's checkpoint). + // We use successorStarted to know whether a successor was launched + // for this particular child; if so, find it by scanning backwards + // through liveRunIDs for a still-live child that we haven't + // previously promoted. + if ps.successorStarted[runID] { + // Find the successor: scan liveRunIDs for the first live child + // after this one's position. Since handover is sequential (at + // most one in flight), the successor is simply any remaining + // live child. + for _, candidateRunID := range ps.liveRunIDs { + if candidateRunID == runID { + continue + } + if _, live := ps.liveChildren[candidateRunID]; !live { + continue + } + // Send resumeFullRate to the successor. Best-effort: + // if the signal fails (e.g., successor already completed), + // the child simply stays at half rate for its remaining work, + // which is safe. + successorExec := ps.liveExecs[candidateRunID] + _ = workflow.SignalExternalWorkflow(ctx, + successorExec.ID, + successorExec.RunID, + shardedResumeFullSignalName, + struct{}{}, + ).Get(ctx, nil) + break + } + } +} + +// updateQPS feeds the current aggregate verified count into the QPSQueue +// and refreshes the rate gauge. Called on each progress rollup and child +// completion so the rate tracks actual throughput. +func (ps *shardedParentState) updateQPS(ctx workflow.Context) { + total := ps.retiredTotal + ps.liveTotalCount() + ps.params.QPSQueue.Enqueue(ctx, total) + ps.params.ReplicatedWorkflowCountPerSecond = ps.params.QPSQueue.CalculateQPS() + ps.metricsHandler.Gauge(ForceReplicationRpsTagName).Update(ps.params.ReplicatedWorkflowCountPerSecond) +} diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index c3518c4eddf..deb16852cb9 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -8,19 +8,40 @@ import ( ) const ( - // shardedForceReplicationWorkflowName is the registered workflow name. - // Distinct from the legacy ForceReplicationWorkflow so both variants - // coexist in the same worker — pick which to use at workflow start - // time by setting the start request's workflow type. + // shardedForceReplicationWorkflowName is the registered workflow name for + // the parent. Distinct from the legacy ForceReplicationWorkflow so both + // variants coexist in the same worker — pick at workflow start time by + // setting the workflow type. shardedForceReplicationWorkflowName = "force-replication-sharded" + // shardedForceReplicationWorkerName is the registered workflow name for + // the child worker workflows spawned by the parent. Registered on the + // same sharded worker component and task queue as the parent. + shardedForceReplicationWorkerName = "force-replication-sharded-worker" + // releaseShardsSignalName carries mid-flight ReleaseShards signals // from active replicate-batch activities back to their parent - // workflow. Drain-mode shard completions ride the activity return - // value instead, so this signal only fires while the activity is - // still running normally. + // workflow (the child). Drain-mode shard completions are gone in + // the new design, so this signal only fires while the activity is + // still running normally and the child needs to free the shard early + // so a successor can pack against it. releaseShardsSignalName = "ReleaseShards" + // shardedCheckpointSignalName is sent by a child to the parent at + // its cut point — when GetContinueAsNewSuggested fires at a page + // boundary. Payload: shardedCheckpointPayload. + shardedCheckpointSignalName = "force-replication-sharded-checkpoint" + + // shardedProgressSignalName carries a periodic (60 s) cumulative + // verified-count rollup from each live child to the parent. + // Payload: shardedProgressPayload. + shardedProgressSignalName = "force-replication-sharded-progress" + + // shardedResumeFullSignalName is sent by the parent to a child to + // promote it to full-rate operation once its predecessor has + // completed. Payload: empty struct — presence is the signal. + shardedResumeFullSignalName = "force-replication-sharded-resume-full" + // defaultShardedListPageSize is the ListWorkflows page size when // the sharded workflow's params.ListWorkflowsPageSize is unset. defaultShardedListPageSize = 1000 @@ -40,35 +61,19 @@ const ( // defaultShardNoProgress is the per-shard cumulative no-progress // backstop. While a shard's pending exec count is non-zero and // no exec on that shard has produced a verified outcome for this - // long (carried across CAN via the resume payload), the activity - // fails non-retryably naming the stuck shard. A shard awaiting its - // very first verification gets double this window so the server has - // time to clear any backlog predating our task submission; it - // reverts to this value once the shard's first exec verifies. + // long, the activity fails non-retryably naming the stuck shard. + // A shard awaiting its very first verification gets double this + // window so the server has time to clear any backlog predating our + // task submission; it reverts to this value once the shard's first + // exec verifies. defaultShardNoProgress = 5 * time.Minute - // defaultDrainGrace is the wall-budget the activity gets after - // the workflow cancels it (on lastErr or cycle drain timeout). - // Continues verifying until either the grace expires, the - // idle-cost trigger fires, or every exec verifies. - defaultDrainGrace = 15 * time.Second - // defaultIdleShardCost is the cumulative idle-time threshold // (the "shard-seconds" unit: 30 s with 1 idle shard equals // 3.3 s with 9 idle) at which the activity signal-releases its // completed-but-not-yet-released shards mid-flight. defaultIdleShardCost = 30 * time.Second - // defaultCycleDrainTimeout bounds the wall-clock the workflow - // will spend draining buckets + awaiting in-flight activities - // after the page loop stops. GetContinueAsNewSuggested trips at - // ~8% of the hard history cap so we have ~92% of the budget - // remaining when the page loop breaks; 10 minutes is generously - // inside that. Catches the "many shards making slow-but-real - // progress" case; a single stuck shard is already bounded by - // ShardNoProgress on the activity side. - defaultCycleDrainTimeout = 10 * time.Minute - // defaultPerBatchGenerateRPS is the per-batch inject-phase target. // Sharded dispatches many concurrent batches and each builds its // own limiter, so this caps the per-batch generate-replication-task @@ -132,8 +137,7 @@ func (r *RunEntry) UnmarshalJSON(data []byte) error { // BatchPayload groups runs by (shard, businessID) so a single businessID // with many runs costs one BID-string-worth of bytes instead of one per -// run. The wire shape behind shardedBatchReq.Executions, ResumeShard.Execs -// (the per-shard inner map), and ShardedForceReplicationParams.RecoveredBuckets. +// run. The wire shape behind shardedBatchReq.Executions. // // On-wire form: // @@ -211,44 +215,7 @@ func (p BatchPayload) flatten() []*shardedExecutionInfo { return out } -// merge folds src into p. Callers guarantee disjoint (shard, BID) keys -// between src and any prior merges into p — in-flight batches hold -// disjoint shard claims and listed-but-unpacked buckets share no shard -// with batchExecs — so per-key appends never interleave across iterations. -func (p BatchPayload) merge(src BatchPayload) { - //workflowcheck:ignore (writes are to disjoint keys; order-independent) - for sh, byBID := range src { - if p[sh] == nil { - p[sh] = map[string][]RunEntry{} - } - //workflowcheck:ignore (writes are to disjoint keys; order-independent) - for bid, runs := range byBID { - p[sh][bid] = append(p[sh][bid], runs...) - } - } -} - -// addRunCountsTo increments counts[shard] by the number of runs in p. -func (p BatchPayload) addRunCountsTo(counts map[int32]int) { - //workflowcheck:ignore (per-shard sum is order-independent) - for sh, byBID := range p { - //workflowcheck:ignore (per-shard sum is order-independent) - for _, runs := range byBID { - counts[sh] += len(runs) - } - } -} - -// mergeInto folds p into dst and bumps counts by the runs merged from p. -func (p BatchPayload) mergeInto(dst BatchPayload, counts map[int32]int) { - if len(p) == 0 { - return - } - dst.merge(p) - p.addRunCountsTo(counts) -} - -// ShardedForceReplicationParams is the workflow input. Configuration +// ShardedForceReplicationParams is the parent workflow input. Configuration // fields are read-only across CAN cycles; the carry-over block at the // bottom is mutated each cycle. type ShardedForceReplicationParams struct { @@ -262,17 +229,8 @@ type ShardedForceReplicationParams struct { DisableVerification bool ShardNoProgress time.Duration - DrainGrace time.Duration IdleShardCost time.Duration - // CycleDrainTimeout caps how long the workflow will spend after - // the page loop stops, draining buckets and waiting for in-flight - // activities to complete naturally. On expiry the workflow falls - // into drainForCAN — cancels in-flight batches, collects their - // drain payload, and CANs with the recovered state. Defaults to - // defaultCycleDrainTimeout. - CycleDrainTimeout time.Duration - TaskQueueUserDataReplicationParams TaskQueueUserDataReplicationParams // PerBatchGenerateRPS is the inject-phase rate-limiter target inside @@ -306,109 +264,77 @@ type ShardedForceReplicationParams struct { // per-second rate doesn't drop to zero on every cycle boundary. QPSQueue QPSQueue - // ResumeShards carries unverified execs from drained activities in - // the prior CAN cycle. The new run dispatches resume activities for - // these before the page loop runs, so their shards are claimed in - // shardInFlight from the start and the packer treats them as busy. - ResumeShards []ResumeShard - - // RecoveredBuckets carries execs whose dispatching activity returned - // a cancellation without returning a result — i.e., the activity - // body never ran, because cancellation (from lastErr or cycle drain - // timeout) reached it before the worker picked it up. They were - // dispatched but never injected, so the new cycle restores them - // into the streaming buckets to be dispatched as fresh inject+verify - // batches. - RecoveredBuckets BatchPayload - TaskQueueUserDataReplicationStatus TaskQueueUserDataReplicationStatus } -// ResumeShard carries one shard's worth of unverified execs from a drained -// activity across a CAN boundary to the resume activity that picks them -// up. NoProgressDuration is the cumulative time the shard went without a -// verified outcome at drain time; the resume activity initialises its own -// per-shard last-progress clock to (now - NoProgressDuration) so the -// backstop check sees the full elapsed no-progress window, not just the -// current activity's slice. -// -// Execs is keyed by businessID: each entry is a list of RunEntry tuples -// for that BID. Grouping by BID at the wire level lets a hot BID (with -// many runs) collapse to one BID-string + N tuples rather than N copies -// of the BID; see BatchPayload's docstring. -// -// The slice-of-ResumeShard form is the CAN/activity wire shape: JSON- -// friendly, easy to append from drained batches, and sortable for replay -// determinism. The packer and activity input use BatchPayload plus a -// per-shard no-progress map; resumeShardsToPayload and -// resumeShardsFromPayload convert between the two. -type ResumeShard struct { - Shard int32 - Execs map[string][]RunEntry - NoProgressDuration time.Duration +// shardedChildParams is the input to each child worker workflow +// (shardedForceReplicationWorker). It carries the configuration subset +// needed by the child and the handover state for the child's listing range. +type shardedChildParams struct { + // Configuration subset passed down from the parent. + Namespace, Query, NamespaceID, TargetClusterName string + TargetShardCount int32 + BatchSize, MaxExecsPerShard, ListWorkflowsPageSize, ConcurrentBatchCount int + DisableVerification bool + ShardNoProgress time.Duration + IdleShardCost time.Duration + PerBatchGenerateRPS float64 + + // StartPageToken is the ListWorkflows continuation token from which + // this child begins listing. Nil for the very first child. + StartPageToken []byte + + // StartThrottled, when true, means the child begins at half + // MaxExecsPerShard until the parent signals resumeFullRate. Used + // during handover overlap so a new child and its not-yet-drained + // predecessor together stay ≤ MaxExecsPerShard per shard. + StartThrottled bool } -// resumeShardsFromPayload expands a BatchPayload and its per-shard -// no-progress durations into a shard-sorted ResumeShard slice. -func resumeShardsFromPayload(payload BatchPayload, noProgress map[int32]time.Duration) []ResumeShard { - if len(payload) == 0 { - return nil - } - out := make([]ResumeShard, 0, len(payload)) - for _, sh := range payload.sortedShards() { - execs := payload[sh] - if len(execs) == 0 { - continue - } - out = append(out, ResumeShard{ - Shard: sh, - Execs: execs, - NoProgressDuration: noProgress[sh], - }) - } - return out +// shardedChildResult is the return value of each child worker workflow. +type shardedChildResult struct { + // VerifiedCount is the total number of executions verified by this + // child during its lifetime. Folded into the parent's retiredTotal + // on child completion. + VerifiedCount int64 + + // ReachedEnd is true when this child exhausted the namespace + // (ListWorkflows returned an empty next-page token) — meaning there + // is no more work for a successor. The parent uses this to skip + // starting a successor and to proceed toward terminal completion. + ReachedEnd bool } -// resumeShardsToPayload folds a ResumeShard slice into the BatchPayload -// and per-shard no-progress map the packer and activity input use. -// Duplicate shard entries merge execs; the last NoProgressDuration wins. -func resumeShardsToPayload(shards []ResumeShard) (BatchPayload, map[int32]time.Duration) { - if len(shards) == 0 { - return nil, nil - } - payload := BatchPayload{} - noProgress := map[int32]time.Duration{} - for _, rs := range shards { - if len(rs.Execs) == 0 { - continue - } - if payload[rs.Shard] == nil { - payload[rs.Shard] = map[string][]RunEntry{} - } - //workflowcheck:ignore (writes are to disjoint BID keys per shard entry; order-independent) - for bid, runs := range rs.Execs { - payload[rs.Shard][bid] = append(payload[rs.Shard][bid], runs...) - } - noProgress[rs.Shard] = rs.NoProgressDuration - } - if len(payload) == 0 { - return nil, nil - } - return payload, noProgress +// shardedCheckpointPayload is the body of the shardedCheckpointSignalName +// signal that a child sends to the parent when it reaches a +// GetContinueAsNewSuggested hint at a page boundary. The parent starts a +// successor from NextPageToken and, when this child completes, promotes the +// successor to full rate. +type shardedCheckpointPayload struct { + // ChildRunID identifies the sending child so the parent can track + // which child's successor has already been started. + ChildRunID string + // NextPageToken is the continuation token for the next page range. + // The successor starts listing from here. + NextPageToken []byte +} + +// shardedProgressPayload is the body of the shardedProgressSignalName +// signal that each live child sends to the parent every 60 seconds. +// The parent aggregates these to answer the force-replication-status query. +type shardedProgressPayload struct { + // ChildRunID identifies the sending child within the parent's liveCounts map. + ChildRunID string + // VerifiedCount is the child's cumulative verified-execution count + // at the time of this rollup signal. + VerifiedCount int64 } // shardedBatchReq is the per-batch activity input. Executions is the // per-shard, per-BID nested payload — the workflow has marked every // shard appearing as a top-level key in shardInFlight before dispatch, // and the activity is responsible for either signal-releasing each shard -// mid-flight or listing it in the return value's CompletedShards / InFlight -// set. -// -// Resume=true skips the inject phase: the execs were already injected by -// some earlier activity that was cancelled at drain time and returned its -// unverified execs in its result. NoProgressByShard carries the cumulative -// pre-resume no-progress duration so the per-shard backstop stays -// meaningful across resume cycles. +// mid-flight or listing it in the return value's CompletedShards set. type shardedBatchReq struct { BatchID int64 Namespace string @@ -417,35 +343,26 @@ type shardedBatchReq struct { TargetClusterName string - Resume bool DisableVerification bool - NoProgressByShard map[int32]time.Duration PerBatchGenerateRPS float64 ShardNoProgress time.Duration - DrainGrace time.Duration IdleShardCost time.Duration } -// replicateBatchResult is the activity's return payload. The activity is -// the source of truth for which execs verified vs. are still outstanding -// when it returns — only it has the per-exec verify state — so the drain -// payload rides the return value rather than a signal. The workflow's -// dispatch coroutine reads InFlight into drainPayload on nil-error return. +// replicateBatchResult is the activity's return payload. // // CompletedShards is informational (the dispatch coroutine's defer clears // heldByBatch + shardInFlight regardless), but keeping it in the result // gives metrics a clean handle on "which shards this batch finished". type replicateBatchResult struct { CompletedShards []int32 - InFlight []ResumeShard // VerifiedCount is the number of executions this activity invocation // finished verifying (including retention/zombie skips that resolve - // as verified). The workflow accumulates this into its running - // ReplicatedWorkflowCount and emits the per-batch delta as the - // replicated_workflow_count counter. + // as verified). The child workflow accumulates this into its running + // verifiedCount. VerifiedCount int64 } @@ -462,8 +379,7 @@ type replicateBatchHeartbeat struct { // The workflow handler clears these shards from shardInFlight + // heldByBatch[BatchID] so the packer can immediately dispatch new work // against them while the activity stays running on its still-pending -// shards. Only fires in normal mode — once the activity enters drain mode -// it returns its remaining state via the activity result instead. +// shards. type releaseShardsPayload struct { BatchID int64 Shards []int32 diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index eac27f45681..91f1ea4e69f 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -91,53 +91,6 @@ func TestBatchPayload_Flatten(t *testing.T) { require.Equal(t, "b-z", got[3].BusinessID) } -// TestBatchPayload_mergeInto merges payload runs into dst and keeps -// bucketCounts in sync — the same invariant addToBucket maintains per -// run, but for a bulk restore on CAN entry. -func TestBatchPayload_mergeInto(t *testing.T) { - dst := BatchPayload{1: {"a": {{RunID: "r0"}}}} - counts := map[int32]int{1: 1} - - src := BatchPayload{ - 1: {"b": {{RunID: "r1"}, {RunID: "r2"}}}, - 2: {"c": {{RunID: "r3"}}}, - } - src.mergeInto(dst, counts) - - require.Len(t, dst, 2) - require.Len(t, dst[1]["a"], 1) - require.Len(t, dst[1]["b"], 2) - require.Len(t, dst[2]["c"], 1) - require.Equal(t, map[int32]int{1: 3, 2: 1}, counts) - require.Equal(t, 4, dst.totalRuns()) -} - -func TestResumeShardPayloadRoundTrip(t *testing.T) { - shards := []ResumeShard{ - {Shard: 2, Execs: map[string][]RunEntry{"b": {{RunID: "r2"}}}, NoProgressDuration: 3 * time.Second}, - {Shard: 1, Execs: map[string][]RunEntry{"a": {{RunID: "r1"}, {RunID: "r1b"}}}, NoProgressDuration: 5 * time.Second}, - } - payload, noProgress := resumeShardsToPayload(shards) - got := resumeShardsFromPayload(payload, noProgress) - - require.Len(t, got, 2) - require.Equal(t, int32(1), got[0].Shard) - require.Equal(t, 5*time.Second, got[0].NoProgressDuration) - require.Len(t, got[0].Execs["a"], 2) - require.Equal(t, int32(2), got[1].Shard) - require.Equal(t, 3*time.Second, got[1].NoProgressDuration) -} - -func TestResumeShardsToPayload_mergesDuplicateShards(t *testing.T) { - shards := []ResumeShard{ - {Shard: 1, Execs: map[string][]RunEntry{"a": {{RunID: "r1"}}}, NoProgressDuration: time.Second}, - {Shard: 1, Execs: map[string][]RunEntry{"b": {{RunID: "r2"}}}, NoProgressDuration: 2 * time.Second}, - } - payload, noProgress := resumeShardsToPayload(shards) - require.Len(t, payload[1], 2) - require.Equal(t, 2*time.Second, noProgress[1]) -} - // TestShardVerifyTracker_FirstVerificationDoubledWindow pins the // no-progress backstop's grace for a shard's first verified outcome: a // shard that hasn't verified anything yet gets 2×threshold before @@ -170,17 +123,59 @@ func TestShardVerifyTracker_FirstVerificationDoubledWindow(t *testing.T) { require.True(t, stuck, "must trip at 1×threshold after first verification") } -// TestNewShardVerifyTracker_ResumeSkipsDoubledWindow: a resumed shard's -// tasks were submitted (and given their first-verification grace) in a -// prior CAN cycle, so it is seeded as already past its first -// verification and tracks cumulative no-progress against the plain -// threshold. A fresh shard awaits its first verification. -func TestNewShardVerifyTracker_ResumeSkipsDoubledWindow(t *testing.T) { - execs := []*shardedExecutionInfo{{Shard: 0}} +// TestNewShardVerifyTracker_SeedsAllShards: newShardVerifyTracker with +// execs from two shards produces a tracker with entries for each shard, +// correct pending counts, verifiedAny=false, and a seeded lastProgress. +func TestNewShardVerifyTracker_SeedsAllShards(t *testing.T) { + execs := []*shardedExecutionInfo{ + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-0", RunID: "r0"}, Shard: 0}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-1", RunID: "r1"}, Shard: 0}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-2", RunID: "r2"}, Shard: 1}, + } + + before := time.Now() + tr := newShardVerifyTracker(execs) + after := time.Now() + + require.Len(t, tr, 2, "should have entries for both shards") + + sv0 := tr[0] + require.Equal(t, 2, sv0.pending, "shard 0 should have 2 pending execs") + require.False(t, sv0.verifiedAny, "fresh shard must not have verifiedAny set") + require.False(t, sv0.lastProgress.IsZero(), "lastProgress must be seeded") + require.True(t, !sv0.lastProgress.Before(before) && !sv0.lastProgress.After(after), + "lastProgress should be within the call window") - fresh := newShardVerifyTracker(execs, false, nil) - require.False(t, fresh[0].verifiedAny, "fresh shard awaits its first verification") + sv1 := tr[1] + require.Equal(t, 1, sv1.pending, "shard 1 should have 1 pending exec") + require.False(t, sv1.verifiedAny, "fresh shard must not have verifiedAny set") + require.False(t, sv1.lastProgress.IsZero(), "lastProgress must be seeded") +} - resumed := newShardVerifyTracker(execs, true, map[int32]time.Duration{0: time.Minute}) - require.True(t, resumed[0].verifiedAny, "resumed shard skips the doubled first-verification window") +// TestEffectiveMaxExecsPerShard: effectiveMaxExecsPerShard returns the +// full MaxExecsPerShard when promoted and not cut, and max(cap/2, 1) +// in all other cases. +func TestEffectiveMaxExecsPerShard(t *testing.T) { + cases := []struct { + name string + promoted, cut bool + maxExecsPerShard int + want int + }{ + {"promoted and not cut returns full", true, false, 10, 10}, + {"not promoted returns half", false, false, 10, 5}, + {"promoted and cut returns half", true, true, 10, 5}, + {"minimum of 1 enforced", false, false, 1, 1}, + {"odd cap rounds down but floors at 1", false, false, 3, 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + s := &shardedWorkflowState{ + params: &shardedChildParams{MaxExecsPerShard: tc.maxExecsPerShard}, + promoted: tc.promoted, + cut: tc.cut, + } + require.Equal(t, tc.want, s.effectiveMaxExecsPerShard()) + }) + } } diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index a07de80cab0..95cb5777b80 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -1,12 +1,10 @@ package migration import ( - "errors" "fmt" "slices" "time" - enumspb "go.temporal.io/api/enums/v1" "go.temporal.io/api/workflowservice/v1" sdkclient "go.temporal.io/sdk/client" "go.temporal.io/sdk/temporal" @@ -15,387 +13,225 @@ import ( "go.temporal.io/server/common/metrics" ) -// ShardedForceReplicationWorkflow runs the sharded design for one CAN -// cycle: dispatch any resume activities carried over from the prior -// cycle, then page through ListWorkflows until either the namespace -// exhausts or workflow.GetContinueAsNewSuggested(ctx) trips, bucketing -// each execution by destination history shard and dispatching a -// paired inject+verify activity once a bucket reaches packing -// eligibility. At cycle end, every remaining bucket flushes as -// packed activities. +// shardedForceReplicationWorker is the child workflow spawned by +// ShardedForceReplicationWorkflow for each listing range. It pages through +// ListWorkflows starting from StartPageToken, routes each execution to a +// destination history shard bucket, and dispatches packed inject+verify +// activities (ReplicateBatch) as buckets fill. // -// In-flight activities are allowed to finish naturally at cycle end; -// CycleDrainTimeout bounds the wait as a safety net against -// pathological slow drains. On timer expiry in-flights are cancelled -// and their drained execs feed the next cycle's carry-over as -// ResumeShards. On lastErr the same drain happens, but the workflow -// returns the error rather than CANing — the drained state surfaces -// only via the status query's recovery bundle. +// Lifecycle: +// 1. List pages at effectiveMaxExecsPerShard rate until either the +// namespace is exhausted (ReachedEnd=true, no checkpoint) or +// GetContinueAsNewSuggested fires at a page boundary (send checkpoint +// signal to parent, set cut flag, stop listing). +// 2. Drain remaining bucket execs via drainBuckets. +// 3. Await pendingDispatches == 0. +// 4. Return shardedChildResult{VerifiedCount, ReachedEnd}. // -// If listing wasn't exhausted, the workflow CANs with NextPageToken -// (plus any drained state from the timer-fired path). Otherwise it -// returns nil. -func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { - // Page token at workflow entry — returned by the status query as - // PageTokenForRestart so tooling that already knows the legacy - // "restart from the starting position" semantic keeps working. - // The richer Recovery* fields below carry the current page token - // plus in-flight execs and are what the sharded restart flow - // actually uses. - startPageToken := params.NextPageToken - - // state is assigned after newShardedWorkflowState below; the - // query handler closes over the pointer so it sees the live state - // once setup completes. Queries that arrive during setup return - // the static fields without the recovery bundle, which matches - // the prior behaviour. - var state *shardedWorkflowState - - // Register the status query under the same name upstream uses - // (forceReplicationStatusQueryType = "force-replication-status") - // so tooling that polls force-rep progress works across both - // workflow variants. - if err := workflow.SetQueryHandler(ctx, forceReplicationStatusQueryType, func() (ForceReplicationStatus, error) { - status := ForceReplicationStatus{ - ContinuedAsNewCount: params.ContinuedAsNewCount, - TotalWorkflowCount: params.TotalForceReplicateWorkflowCount, - ReplicatedWorkflowCount: params.ReplicatedWorkflowCount, - ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, - PageTokenForRestart: startPageToken, - TaskQueueUserDataReplicationStatus: params.TaskQueueUserDataReplicationStatus, - RecoveryNextPageToken: params.NextPageToken, - RecoveryResumeShards: params.ResumeShards, - RecoveryBuckets: params.RecoveredBuckets, - } - if state != nil { - status.RecoveryResumeShards = state.collectResumeShardsForCarryover() - status.RecoveryBuckets = state.collectRecoveredBucketsForCarryover() - } - return status, nil - }); err != nil { - return err - } - - if err := validateShardedForceReplicationParams(¶ms); err != nil { - return err - } - - var err error - state, err = newShardedWorkflowState(ctx, ¶ms) - if err != nil { - return err - } - // Defaults are now applied; reject configurations the packer - // can't honour. MaxExecsPerShard > BatchSize is meaningless — - // each batch caps at BatchSize total, so the per-shard cap - // can't exceed the whole-batch cap. - if params.MaxExecsPerShard > params.BatchSize { - return temporal.NewNonRetryableApplicationError( - fmt.Sprintf("MaxExecsPerShard (%d) must be <= BatchSize (%d)", params.MaxExecsPerShard, params.BatchSize), - "InvalidConfiguration", nil) - } - - // On the first cycle, populate TotalForceReplicateWorkflowCount - // via the same CountWorkflow activity upstream uses. Skipped on - // subsequent CAN cycles — the count carries across via params. - if params.TotalForceReplicateWorkflowCount == 0 { - wfCount, err := countWorkflowsForReplication(ctx, params.Namespace, params.Query, shardedCountWorkflowsForReplicationTimeout) - if err != nil { - return err - } - params.TotalForceReplicateWorkflowCount = wfCount - } - - if !params.TaskQueueUserDataReplicationStatus.Done { - if err := maybeKickoffShardedTaskQueueUserDataReplication(ctx, ¶ms, func(failureReason string) { - params.TaskQueueUserDataReplicationStatus.FailureMessage = failureReason - params.TaskQueueUserDataReplicationStatus.Done = true - }); err != nil { - return err - } - } - - if err := state.run(ctx); err != nil { - return err - } - - // state.run returned nil only on the terminal cycle (no more pages, - // no errors). On CAN cycles it returns the CAN error, so we never - // reach here mid-replication. - if err := workflow.Await(ctx, func() bool { return params.TaskQueueUserDataReplicationStatus.Done }); err != nil { - return err - } - if params.TaskQueueUserDataReplicationStatus.FailureMessage != "" { - return fmt.Errorf("task queue user data replication failed: %v", params.TaskQueueUserDataReplicationStatus.FailureMessage) - } - return nil -} - -func validateShardedForceReplicationParams(params *ShardedForceReplicationParams) error { - if len(params.Namespace) == 0 { - return temporal.NewNonRetryableApplicationError("InvalidArgument: Namespace is required", "InvalidArgument", nil) - } - if len(params.TargetClusterName) == 0 { - return temporal.NewNonRetryableApplicationError("InvalidArgument: TargetClusterName is required", "InvalidArgument", nil) - } - return nil -} - -func maybeKickoffShardedTaskQueueUserDataReplication(ctx workflow.Context, params *ShardedForceReplicationParams, onDone func(failureReason string)) error { - workflow.Go(ctx, func(ctx workflow.Context) { - ch := workflow.GetSignalChannel(ctx, taskQueueUserDataReplicationDoneSignalType) - var errStr string - _ = ch.Receive(ctx, &errStr) - onDone(errStr) - }) - - if params.ContinuedAsNewCount > 0 { - return nil - } - - options := workflow.ChildWorkflowOptions{ - WorkflowID: fmt.Sprintf("%s-task-queue-user-data-replicator", workflow.GetInfo(ctx).WorkflowExecution.ID), - ParentClosePolicy: enumspb.PARENT_CLOSE_POLICY_ABANDON, - } - childCtx := workflow.WithChildOptions(ctx, options) - input := TaskQueueUserDataReplicationParamsWithNamespace{ - TaskQueueUserDataReplicationParams: params.TaskQueueUserDataReplicationParams, - Namespace: params.Namespace, +// A 60 s timer coroutine sends cumulative verifiedCount rollup signals to +// the parent so the force-replication-status query stays up to date. +// +// Throttled-hits-hint: a child started with StartThrottled=true (not yet +// promoted) pauses at a CAN hint and awaits the parent's resumeFullRate +// signal before cutting. This ensures at most one handover is in flight at +// any time and at most two children run concurrently. +func shardedForceReplicationWorker(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { + s := &shardedWorkflowState{ + params: ¶ms, + namespaceID: params.NamespaceID, + targetShardCount: params.TargetShardCount, + buckets: BatchPayload{}, + bucketCounts: map[int32]int{}, + shardInFlight: map[int32]bool{}, + heldByBatch: map[int64]map[int32]bool{}, + // StartThrottled=true → wait for parent promotion before going full. + // StartThrottled=false → first child, already at full rate. + promoted: !params.StartThrottled, + metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ + metrics.OperationTagName: metrics.MigrationWorkflowScope, + NamespaceTagName: params.Namespace, + }), } - child := workflow.ExecuteChildWorkflow(childCtx, ForceTaskQueueUserDataReplicationWorkflow, input) - var childExecution workflow.Execution - return child.GetChildWorkflowExecution().Get(ctx, &childExecution) + return s.run(ctx) } -// shardedWorkflowState holds the workflow's per-run state. Workflow +// shardedWorkflowState holds the child workflow's per-run state. Workflow // coroutines yield only at SDK calls, so plain maps + ints are safe -// without mutexes — workflow.Await re-evaluates its predicate after -// each yield, which is what makes the shard-in-flight bookkeeping -// drive each dispatch coroutine's wait. +// without mutexes — workflow.Await re-evaluates its predicate after each +// yield, which is what makes the shard-in-flight bookkeeping drive each +// dispatch coroutine's wait. type shardedWorkflowState struct { - params *ShardedForceReplicationParams + params *shardedChildParams namespaceID string - // targetShardCount is the target cluster's history shard count, - // fetched once via DescribeTargetCluster at state construction. - // Drives the per-exec shard hash (so packing groups execs by their - // destination shard) and the default ConcurrentBatchCount. + // targetShardCount is the target cluster's history shard count. + // Drives the per-exec shard hash and ConcurrentBatchCount derivation. targetShardCount int32 // buckets accumulate execs that have been listed but not yet - // dispatched. Nested by destination shard then businessID, so a - // hot BID's many runs share one BID-string-worth of bytes when - // the bucket is shipped over the wire. + // dispatched. Nested by destination shard then businessID. buckets BatchPayload - // bucketCounts mirrors len of all runs across BIDs for each - // shard. Kept as a sidecar so the packer's per-shard ordering - // decisions are O(1) rather than O(#BIDs in shard); it's - // consulted many times per cycle. + // bucketCounts mirrors len of all runs across BIDs for each shard. + // Kept as a sidecar so the packer's per-shard ordering decisions are + // O(1); consulted many times per cycle. bucketCounts map[int32]int - // shardInFlight is the per-shard exclusivity set: a shard's - // entry is set when it's part of any in-flight batch and - // cleared when that batch returns (either fully or via mid-flight - // signal-release). The packer uses this set to ensure that each - // shard can only be present in one in-flight batch at a time. + // shardInFlight is the per-shard exclusivity set: a shard's entry is + // set when it's part of any in-flight batch and cleared when that + // batch returns (either fully or via mid-flight signal-release). The + // packer uses this to ensure each shard is in at most one in-flight + // batch at a time. shardInFlight map[int32]bool - // heldByBatch tracks per-batch shard ownership. spawnBatch - // populates it with the batch's claimed shards; the signal - // handler removes entries as shards are released mid-flight; - // the dispatch coroutine's defer clears whatever's left after - // the activity returns. Required because a signal-released - // shard may have been re-claimed by a subsequent batch — the - // returning original batch must only clear its own remaining - // claims, not stomp on the new claimant. + // heldByBatch tracks per-batch shard ownership. spawnBatch populates + // it with the batch's claimed shards; the signal handler removes + // entries as shards are released mid-flight; the dispatch coroutine's + // defer clears whatever remains after the activity returns. Required + // because a signal-released shard may have been re-claimed by a + // subsequent batch — the returning original batch must only clear its + // own remaining claims, not stomp on the new claimant. heldByBatch map[int64]map[int32]bool - // activityCtx is a cancellable child of run()'s ctx; every - // dispatched batch activity is derived from it. cancelActivities - // is the matching cancel func — drainForCAN calls it once to - // cancel every in-flight batch at once. Cancelling activityCtx - // leaves the workflow's main ctx alive so the drain loop's - // Await keeps running. - activityCtx workflow.Context - cancelActivities workflow.CancelFunc - - // batchExecs tracks the input payload of each in-flight batch. - // Cleared on any nil-error return (drained execs are folded - // into drainPayload from the activity result; cleanly completed - // batches return an empty InFlight). Anything left at cycle end - // corresponds to a batch whose activity returned CanceledError - // with no result — i.e. the activity body never ran. On the CAN - // path those execs are recovered into the next cycle's streaming - // buckets; on the lastErr return path they surface via the status - // query's recovery bundle but aren't auto-dispatched. - batchExecs map[int64]BatchPayload - - // pendingDispatches counts spawned dispatch coroutines that - // have not yet returned. The main coroutine waits on this - // dropping to zero before issuing CAN or returning. + // pendingDispatches counts spawned dispatch coroutines that have not + // yet returned. The main coroutine awaits this dropping to zero before + // returning the child result. pendingDispatches int - // drainPayload accumulates ResumeShard entries from drained - // activities (via the activity result on nil-error return). - // Fed into the CAN-carry-over params at the end. - drainPayload []ResumeShard - - // lastErr stops further dispatch once an activity errors out - // (e.g. ShardNoProgress). Without it the workflow would keep - // paging through ListWorkflows after a broken namespace surfaces, - // burning the rest of the population before returning the - // failure. + // lastErr latches the first activity error. Once set, further dispatch + // stops (tryPackStreaming and drainBuckets bail out) and the child + // returns the error. lastErr error - // cycleDrainTimedOut is set by the safety timer started after the - // page loop. drainBuckets and awaitInFlightCompletion both honour - // it so a pathological slow drain can't eat into the workflow's - // remaining history budget — on expiry the workflow falls into - // drainForCAN and ships unfinished work as carry-over. - cycleDrainTimedOut bool + // promoted is set when the parent sends the resumeFullRate signal: + // the predecessor child has completed and this child may run at the + // full MaxExecsPerShard rate. A child started with StartThrottled=false + // is promoted from the start (no predecessor). + promoted bool + + // cut is set when this child has sent its checkpoint signal to the + // parent and stopped listing new pages. After cut, effectiveMaxExecsPerShard + // drops to half so the successor can run at half alongside us without + // the combined per-shard in-flight count exceeding MaxExecsPerShard. + cut bool + + // verifiedCount accumulates verified executions for the final child + // result and the 60 s progress rollup signal to the parent. + verifiedCount int64 nextBatchID int64 - // metricsHandler is tagged with the workflow's fixed scope + - // namespace once at state construction; recordVerified reuses it on - // every batch return. + // metricsHandler is tagged with the workflow's fixed scope + namespace + // once at construction; recordVerified reuses it on every batch return. metricsHandler sdkclient.MetricsHandler } -// defaultConcurrentBatchCount derives the in-flight-batch ceiling -// from the target cluster's shard count: a quarter of the shards, -// capped at defaultConcurrentBatchCap. The 1/4 fraction leaves worker -// slots free for unrelated activities; the absolute cap bounds the -// cluster blast radius regardless of cluster size. Returns at least 1. -func defaultConcurrentBatchCount(shards int32) int { - return max(min(int(shards)/4, defaultConcurrentBatchCap), 1) -} - -// applyShardedDefaults fills zero-valued tuning fields on params. -func applyShardedDefaults(params *ShardedForceReplicationParams, targetShardCount int32) { - if params.BatchSize <= 0 { - params.BatchSize = defaultBatchSize - } - if params.MaxExecsPerShard <= 0 { - params.MaxExecsPerShard = defaultMaxExecsPerShard - } - if params.ShardNoProgress <= 0 { - params.ShardNoProgress = defaultShardNoProgress - } - if params.DrainGrace <= 0 { - params.DrainGrace = defaultDrainGrace - } - if params.IdleShardCost <= 0 { - params.IdleShardCost = defaultIdleShardCost - } - if params.CycleDrainTimeout <= 0 { - params.CycleDrainTimeout = defaultCycleDrainTimeout - } - if params.ListWorkflowsPageSize <= 0 { - params.ListWorkflowsPageSize = defaultShardedListPageSize - } - if params.PerBatchGenerateRPS <= 0 { - params.PerBatchGenerateRPS = defaultPerBatchGenerateRPS - } - if params.ConcurrentBatchCount <= 0 { - params.ConcurrentBatchCount = defaultConcurrentBatchCount(targetShardCount) - } - if params.EstimationMultiplier <= 0 { - params.EstimationMultiplier = 2 +// effectiveMaxExecsPerShard returns MaxExecsPerShard when this child is +// promoted (predecessor done) and not yet cut (checkpoint not yet sent): +// i.e., the normal full-rate steady state. In all other cases it returns +// max(MaxExecsPerShard/2, 1). +// +// Half-rate applies in two distinct situations: +// - Before promotion: this child was started alongside a still-running +// predecessor; together they must stay ≤ MaxExecsPerShard per shard. +// - After cut: a successor has been started at half rate alongside us; +// same combined-load constraint. +// +// Minimum of 1 ensures at least one exec can always be packed regardless +// of BatchSize or MaxExecsPerShard settings. +func (s *shardedWorkflowState) effectiveMaxExecsPerShard() int { + if s.promoted && !s.cut { + return s.params.MaxExecsPerShard } + return max(s.params.MaxExecsPerShard/2, 1) } -func newShardedWorkflowState(ctx workflow.Context, params *ShardedForceReplicationParams) (*shardedWorkflowState, error) { - lao := workflow.LocalActivityOptions{ - StartToCloseTimeout: 1 * time.Second, - RetryPolicy: forceReplicationActivityRetryPolicy, - } - localCtx := workflow.WithLocalActivityOptions(ctx, lao) - var a *activities - var md MetadataResponse - if err := workflow.ExecuteLocalActivity(localCtx, a.GetMetadata, MetadataRequest{Namespace: params.Namespace}).Get(ctx, &md); err != nil { - return nil, err - } - var targetMd DescribeTargetClusterResponse - if err := workflow.ExecuteLocalActivity(localCtx, a.DescribeTargetCluster, DescribeTargetClusterRequest{ - TargetClusterName: params.TargetClusterName, - }).Get(ctx, &targetMd); err != nil { - return nil, err - } - applyShardedDefaults(params, targetMd.ShardCount) - // QPSQueue is sized off ConcurrentBatchCount (one sample slot per - // expected in-flight batch + one for the starting count). Seeded - // with the current ReplicatedWorkflowCount so the very first - // post-CAN batch return has a baseline to compute the rate against. - if params.QPSQueue.Data == nil { - params.QPSQueue = NewQPSQueue(params.ConcurrentBatchCount, params.EstimationMultiplier) - params.QPSQueue.Enqueue(ctx, params.ReplicatedWorkflowCount) +// run is the child workflow body. See shardedForceReplicationWorker for +// the full lifecycle description. +func (s *shardedWorkflowState) run(ctx workflow.Context) (shardedChildResult, error) { + // Intra-child signal handler: activities signal mid-flight shard releases. + workflow.Go(ctx, s.handleReleaseSignals) + + parentExec := workflow.GetInfo(ctx).ParentWorkflowExecution + s.startBackgroundCoroutines(ctx, parentExec) + + reachedEnd := s.listUntilCutOrEnd(ctx, parentExec) + if s.lastErr != nil { + return shardedChildResult{}, s.lastErr } - s := &shardedWorkflowState{ - params: params, - namespaceID: md.NamespaceID, - targetShardCount: targetMd.ShardCount, - buckets: BatchPayload{}, - bucketCounts: map[int32]int{}, - shardInFlight: map[int32]bool{}, - heldByBatch: map[int64]map[int32]bool{}, - batchExecs: map[int64]BatchPayload{}, - metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ - metrics.OperationTagName: metrics.MigrationWorkflowScope, - NamespaceTagName: params.Namespace, - }), + + // Drain remaining buckets (at effectiveMaxExecsPerShard, which is + // half after cut so the successor can run alongside). + s.drainBuckets(ctx) + if s.lastErr != nil { + return shardedChildResult{}, s.lastErr } - // Restore execs recovered from cancel-before-start batches in - // the prior cycle so the streaming packer picks them up - // alongside any new pages. - params.RecoveredBuckets.mergeInto(s.buckets, s.bucketCounts) - params.RecoveredBuckets = nil - return s, nil -} -func (s *shardedWorkflowState) run(ctx workflow.Context) error { - // Cancellable child ctx for all dispatched batch activities. - // drainForCAN cancels it once to drain in-flight batches without - // touching the workflow's main ctx (which the drain loop's - // Await still rides on). - actCtx, cancelAll := workflow.WithCancel(ctx) - defer cancelAll() - s.activityCtx = actCtx - s.cancelActivities = cancelAll + // Await all in-flight batches. No cancellation — children always run + // to natural completion so their verified counts are exact. + _ = workflow.Await(ctx, func() bool { return s.pendingDispatches == 0 }) + if s.lastErr != nil { + return shardedChildResult{}, s.lastErr + } - // Start the signal handler coroutine first so any signal - // arriving during resume dispatch or page-loop drains is - // processed promptly. - workflow.Go(ctx, s.handleReleaseSignals) + return shardedChildResult{ + VerifiedCount: s.verifiedCount, + ReachedEnd: reachedEnd, + }, nil +} - // Dispatch resume activities carried over from the prior cycle. - // Done before the page loop so their shards are claimed in - // shardInFlight before any new pages arrive — keeps the packer - // from racing to dispatch against them with fresh execs. - s.dispatchResumeBatches(ctx) +// startBackgroundCoroutines launches the promotion-signal handler and the +// 60 s progress rollup. The rollup runs only when the worker has a parent to +// signal (skipped when the worker runs standalone, e.g. in unit tests). +func (s *shardedWorkflowState) startBackgroundCoroutines(ctx workflow.Context, parentExec *workflow.Execution) { + // Parent promotion handler: a one-shot signal that sets promoted=true, + // allowing the child to advance to full rate. Runs until the signal + // arrives (or ctx is cancelled at workflow end). + workflow.Go(ctx, func(gCtx workflow.Context) { + ch := workflow.GetSignalChannel(gCtx, shardedResumeFullSignalName) + var dummy struct{} + if ch.Receive(gCtx, &dummy) { + s.promoted = true + } + }) - // Drive ListWorkflows until either we exhaust the namespace or - // the SDK signals that history is large enough to CAN. Errors - // here latch into lastErr and fall through to the unified exit - // funnel — same drain-and-decide path as activity-driven errors. - // - // On a CAN cycle (ContinuedAsNewCount > 0), an empty NextPageToken - // means a prior cycle already exhausted pagination — empty token at - // the start of cycle 0 is "haven't started", but at the start of any - // later cycle it's "finished". Skip listing in that case; otherwise - // a carry-over-driven CAN (drained InFlight, recovered buckets, …) - // would re-list page 1 and re-enqueue every visible execution. - for s.shouldList() && !workflow.GetInfo(ctx).GetContinueAsNewSuggested() { - if s.lastErr != nil { - break + if parentExec == nil { + return + } + // 60 s progress rollup: send cumulative verifiedCount to the parent so the + // force-replication-status query stays current even for long-running + // children. Stops when ctx is cancelled at workflow exit (timer Get errors). + workflow.Go(ctx, func(gCtx workflow.Context) { + for { + if err := workflow.NewTimer(gCtx, 60*time.Second).Get(gCtx, nil); err != nil { + return // ctx cancelled — workflow is completing + } + // Best-effort: signal failure is not fatal for the child. + _ = workflow.SignalExternalWorkflow(gCtx, + parentExec.ID, parentExec.RunID, + shardedProgressSignalName, + shardedProgressPayload{ + ChildRunID: workflow.GetInfo(gCtx).WorkflowExecution.RunID, + VerifiedCount: s.verifiedCount, + }).Get(gCtx, nil) } - executions, nextPageToken, err := s.listWorkflowPage(ctx) + }) +} + +// listUntilCutOrEnd drives the listing loop: page through ListWorkflows, +// bucketing and opportunistically packing each page, until either the +// namespace is exhausted (returns true) or a CAN hint at a fully-consumed page +// boundary triggers a cut (returns false). Latches s.lastErr on listing failure. +func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExec *workflow.Execution) bool { + currentPageToken := s.params.StartPageToken + for s.lastErr == nil { + executions, nextPageToken, err := s.listWorkflowPageWithToken(ctx, currentPageToken) if err != nil { - s.setLastErr(err) - break + s.lastErr = err + return false } + for _, ex := range executions { sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.targetShardCount) s.addToBucket(sh, ex.BusinessID, RunEntry{ @@ -403,59 +239,52 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) error { ArchetypeID: ex.ArchetypeID, }) } - s.params.NextPageToken = nextPageToken + // Opportunistically pack and dispatch batches from the buckets. for s.tryPackStreaming(ctx, false) { //nolint:revive // intentional empty body } if len(nextPageToken) == 0 { - break + return true // namespace exhausted — terminal child, no successor needed } - } - - // Start the cycle drain safety timer; see defaultCycleDrainTimeout - // for the budget rationale. - cancelCycleTimer := s.startCycleDrainTimer(ctx) - defer cancelCycleTimer() + currentPageToken = nextPageToken - // Drain remaining buckets only when no error has latched — on - // error we deliberately stop scheduling new work and let the - // already-dispatched batches finish via awaitInFlightCompletion. - if s.lastErr == nil { - s.drainBuckets(ctx) + // Check the CAN hint at this fully-consumed page boundary, so + // currentPageToken is exact and safe to hand to a successor. + if workflow.GetInfo(ctx).GetContinueAsNewSuggested() { + s.cutAtBoundary(ctx, parentExec, currentPageToken) + return false + } } + return false +} - // Wait for in-flight activities. Cancels them when we already - // know we're failing or the cycle drain timer has expired; - // otherwise waits naturally so a healthy activity isn't - // cancelled into a CanceledError that masquerades as carry-over - // state. - s.awaitInFlightCompletion(ctx) - - // Single exit decision. Recovery state has the same shape on - // either path — drainPayload + undispatched ResumeShards + - // batchExecs + leftover buckets — so the only difference between - // "fail with state" and "CAN with state" is the return value. - if s.lastErr != nil { - return s.lastErr +// cutAtBoundary performs the handover cut. If still throttled it first awaits +// promotion (so at most one handover is ever in flight), then sets cut (dropping +// to half rate) and signals the parent the next page token; the parent starts a +// successor from there and promotes it once this child completes. +func (s *shardedWorkflowState) cutAtBoundary(ctx workflow.Context, parentExec *workflow.Execution, pageToken []byte) { + if !s.promoted { + // Throttled-hits-hint: the predecessor is still running. Pausing here + // bounds our history at ~hint size and guarantees at most one handover + // is in flight at a time. Await promotion (predecessor completion), then + // cut immediately. + _ = workflow.Await(ctx, func() bool { return s.promoted || s.lastErr != nil }) + if s.lastErr != nil { + return + } } - if !s.hasCarryover() { - return nil + s.cut = true + if parentExec == nil { + return } - - next := *s.params - next.ContinuedAsNewCount++ - next.ResumeShards = s.collectResumeShardsForCarryover() - next.RecoveredBuckets = s.collectRecoveredBucketsForCarryover() - return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) -} - -// shouldList returns true if the page loop has more work to do this -// cycle. Cycle 0 always lists (NextPageToken is empty either way). -// On later cycles, an empty NextPageToken can only mean a prior cycle -// already drained pagination — there's nothing left to list. -func (s *shardedWorkflowState) shouldList() bool { - return s.params.ContinuedAsNewCount == 0 || len(s.params.NextPageToken) > 0 + _ = workflow.SignalExternalWorkflow(ctx, + parentExec.ID, parentExec.RunID, + shardedCheckpointSignalName, + shardedCheckpointPayload{ + ChildRunID: workflow.GetInfo(ctx).WorkflowExecution.RunID, + NextPageToken: pageToken, + }).Get(ctx, nil) } var ( @@ -469,16 +298,14 @@ var ( RetryPolicy: shardedListWorkflowsRetryPolicy, } - // Per-exec backoff still owns the per-exec retry; MaximumAttempts - // lets a transient activity failure recover via heartbeat-resume - // without losing inject progress. WaitForCancellation lets a - // cancelled activity run drain logic and return its drain result. + // Per-exec backoff owns the per-exec retry; MaximumAttempts lets a + // transient activity failure recover via heartbeat-resume without + // losing inject progress. // // 10 attempts is sized for fleet rollouts: a rolling deploy of the - // activity workers can burn several attempts per batch (each - // shutdown surfaces as a retryable WorkerShutdown error from the - // activity). 3 was tight enough that two unlucky deploys could - // exhaust the budget on a long-running CAN cycle. + // activity workers can burn several attempts per batch (each shutdown + // surfaces as a retryable WorkerShutdown error). 3 was tight enough + // that two unlucky deploys could exhaust the budget. shardedReplicateBatchRetryPolicy = &temporal.RetryPolicy{ MaximumAttempts: 10, } @@ -486,17 +313,18 @@ var ( StartToCloseTimeout: 24 * time.Hour, HeartbeatTimeout: time.Minute, RetryPolicy: shardedReplicateBatchRetryPolicy, - WaitForCancellation: true, } ) -func (s *shardedWorkflowState) listWorkflowPage(ctx workflow.Context) ([]*ExecutionInfo, []byte, error) { +// listWorkflowPageWithToken fetches one page of ListWorkflows results +// starting from the given continuation token (nil for the first page). +func (s *shardedWorkflowState) listWorkflowPageWithToken(ctx workflow.Context, pageToken []byte) ([]*ExecutionInfo, []byte, error) { listCtx := workflow.WithActivityOptions(ctx, shardedListWorkflowsActivityOptions) listReq := &workflowservice.ListWorkflowExecutionsRequest{ Namespace: s.params.Namespace, Query: s.params.Query, PageSize: int32(s.params.ListWorkflowsPageSize), - NextPageToken: s.params.NextPageToken, + NextPageToken: pageToken, } var a *activities var listResp listWorkflowsResponse @@ -506,8 +334,8 @@ func (s *shardedWorkflowState) listWorkflowPage(ctx workflow.Context) ([]*Execut return listResp.Executions, listResp.NextPageToken, nil } -func (s *shardedWorkflowState) replicateBatch(ctx, activityParentCtx workflow.Context, req *shardedBatchReq) (replicateBatchResult, error) { - actx := workflow.WithActivityOptions(activityParentCtx, shardedReplicateBatchActivityOptions) +func (s *shardedWorkflowState) replicateBatch(ctx workflow.Context, req *shardedBatchReq) (replicateBatchResult, error) { + actx := workflow.WithActivityOptions(ctx, shardedReplicateBatchActivityOptions) var a *activities var result replicateBatchResult if err := workflow.ExecuteActivity(actx, a.ReplicateBatch, req).Get(ctx, &result); err != nil { @@ -516,198 +344,6 @@ func (s *shardedWorkflowState) replicateBatch(ctx, activityParentCtx workflow.Co return result, nil } -// startCycleDrainTimer spawns a coroutine that sets cycleDrainTimedOut -// after CycleDrainTimeout. Returned cancel func stops the timer on -// natural exit. The flag is read by drainBuckets (stops spawning new -// batches and exits) and awaitInFlightCompletion (exits its Await and -// calls drainForCAN to cancel in-flight activities and collect their -// drain payload for the next cycle's carry-over). -func (s *shardedWorkflowState) startCycleDrainTimer(ctx workflow.Context) workflow.CancelFunc { - timerCtx, cancel := workflow.WithCancel(ctx) - workflow.Go(ctx, func(gCtx workflow.Context) { - if err := workflow.NewTimer(timerCtx, s.params.CycleDrainTimeout).Get(gCtx, nil); err == nil { - s.cycleDrainTimedOut = true - } - }) - return cancel -} - -// awaitInFlightCompletion waits for in-flight batches to complete. -// On the clean path it just blocks until pendingDispatches drops to -// zero — a healthy activity's result isn't masked as a CanceledError. -// On lastErr or cycle-drain timeout it falls into drainForCAN, which -// cancels the in-flight activities and collects their drain payload. -// run() then either returns lastErr (drain payload surfaces only via -// the status query) or CANs with the drained state as carry-over. -func (s *shardedWorkflowState) awaitInFlightCompletion(ctx workflow.Context) { - if s.pendingDispatches == 0 { - return - } - if s.lastErr != nil { - s.drainForCAN(ctx) - return - } - _ = workflow.Await(ctx, func() bool { - return s.pendingDispatches == 0 || s.lastErr != nil || s.cycleDrainTimedOut - }) - if s.pendingDispatches > 0 { - s.drainForCAN(ctx) - } -} - -// hasCarryover reports whether the workflow has any state worth -// preserving across an exit — either a remaining page token, drained -// execs from in-flight batches, cancel-before-start batches that -// never injected, undispatched resume entries, or listed-but-unpacked -// execs. Drives both the "CAN vs return nil" decision and the -// recovery bundle exposed in the status query. -func (s *shardedWorkflowState) hasCarryover() bool { - if len(s.params.NextPageToken) > 0 { - return true - } - if len(s.drainPayload) > 0 { - return true - } - if len(s.batchExecs) > 0 { - return true - } - if len(s.params.ResumeShards) > 0 { - return true - } - return !s.bucketsEmpty() -} - -// collectResumeShardsForCarryover concatenates this cycle's drained -// execs with any prior-cycle ResumeShards that didn't get dispatched -// (left in params.ResumeShards by dispatchResumeBatches when it bailed -// out on lastErr). Both groups are already shard-keyed; the next -// cycle's dispatchResumeBatches sorts and re-packs them. -func (s *shardedWorkflowState) collectResumeShardsForCarryover() []ResumeShard { - if len(s.drainPayload) == 0 && len(s.params.ResumeShards) == 0 { - return nil - } - out := make([]ResumeShard, 0, len(s.drainPayload)+len(s.params.ResumeShards)) - out = append(out, s.drainPayload...) - out = append(out, s.params.ResumeShards...) - return out -} - -// collectRecoveredBucketsForCarryover merges the two sources of -// "execs that never made it through a verify activity this cycle": -// batches that returned CanceledError without running a body, and -// listed-but-unpacked execs still sitting in s.buckets when the -// workflow exited (either lastErr stopped the streaming packer or -// drainBuckets bailed on lastErr / cycle drain timeout partway -// through). -func (s *shardedWorkflowState) collectRecoveredBucketsForCarryover() BatchPayload { - out := collectRecoveredBuckets(s.batchExecs) - if !s.bucketsEmpty() { - if out == nil { - out = BatchPayload{} - } - out.merge(s.buckets) - } - return out -} - -// recordVerified accumulates one batch's verified-exec delta into the -// workflow's running count, emits the per-batch counter delta, and -// updates the sliding-window RPS gauge. No-op when verified == 0 so a -// batch that ran entirely as drain-no-progress doesn't poison the -// QPSQueue with a zero-delta sample. -func (s *shardedWorkflowState) recordVerified(ctx workflow.Context, verified int64) { - if verified <= 0 { - return - } - s.params.ReplicatedWorkflowCount += verified - - s.metricsHandler.Counter(metrics.ReplicatedWorkflowCount.Name()).Inc(verified) - - s.params.QPSQueue.Enqueue(ctx, s.params.ReplicatedWorkflowCount) - s.params.ReplicatedWorkflowCountPerSecond = s.params.QPSQueue.CalculateQPS() - s.metricsHandler.Gauge(ForceReplicationRpsTagName).Update(s.params.ReplicatedWorkflowCountPerSecond) -} - -// collectRecoveredBuckets re-buckets any execs from batches whose -// dispatching activity returned CanceledError without returning a -// result — i.e. the activity body never ran, so its execs were -// never injected. They go back into the next cycle's streaming -// buckets to be dispatched as fresh inject+verify batches. The -// shard is the top-level map key on each batch's payload, so no -// re-hashing here — collectRecoveredBuckets just merges. -func collectRecoveredBuckets(batchExecs map[int64]BatchPayload) BatchPayload { - if len(batchExecs) == 0 { - return nil - } - out := BatchPayload{} - // In-flight batches hold disjoint shard claims, so merging two - // batchExecs entries never targets the same (shard, BID) key — - // iteration order does not affect the final BatchPayload. - //workflowcheck:ignore (writes are to disjoint keys; order-independent) - for _, bp := range batchExecs { - out.merge(bp) - } - return out -} - -// addToBucket appends one run to the (shard, BID) bucket and bumps -// the sidecar count. -func (s *shardedWorkflowState) addToBucket(shard int32, businessID string, run RunEntry) { - if s.buckets[shard] == nil { - s.buckets[shard] = map[string][]RunEntry{} - } - s.buckets[shard][businessID] = append(s.buckets[shard][businessID], run) - s.bucketCounts[shard]++ -} - -// takeFromBucket consumes up to n runs from the given shard and -// returns them grouped by BID. Walks BIDs in alphabetical order so -// the resulting payload is deterministic across replays; takes whole -// per-BID runs only as needed to reach n. Empties the shard from -// s.buckets / s.bucketCounts when nothing remains. -func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]RunEntry { - if n <= 0 { - return nil - } - byBID := s.buckets[shard] - if len(byBID) == 0 { - return nil - } - bids := make([]string, 0, len(byBID)) - //workflowcheck:ignore (bids is sorted before use) - for bid := range byBID { - bids = append(bids, bid) - } - slices.Sort(bids) - - out := map[string][]RunEntry{} - taken := 0 - for _, bid := range bids { - if taken >= n { - break - } - runs := byBID[bid] - take := min(len(runs), n-taken) - // append([]RunEntry(nil), ...) gives the output its own - // backing array — keeps the workflow's leftover slice - // (byBID[bid][take:]) and the activity's input independent - // in case either side appends later. - out[bid] = append([]RunEntry(nil), runs[:take]...) - if take == len(runs) { - delete(byBID, bid) - } else { - byBID[bid] = runs[take:] - } - taken += take - } - s.bucketCounts[shard] -= taken - if s.bucketCounts[shard] <= 0 { - delete(s.bucketCounts, shard) - delete(s.buckets, shard) - } - return out -} - // handleReleaseSignals runs as a long-lived workflow coroutine, // consuming ReleaseShards signals from in-flight activities. Each // signal lists shards the activity considers complete; the handler @@ -717,10 +353,9 @@ func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]R // its still-pending shards). // // DO NOT add workflow yields (ExecuteActivity, Sleep, Await, etc.) -// between Receive and the next Receive. drainForCAN relies on -// ch.Len() == 0 implying "every delivered signal has been processed"; -// a yield mid-handler would invalidate that, leaving shardInFlight -// stale after a CAN. +// between Receive and the next Receive. The handler relies on +// ch.Len() == 0 being a reliable "every delivered signal has been +// processed" indicator; a yield mid-handler would invalidate that. func (s *shardedWorkflowState) handleReleaseSignals(ctx workflow.Context) { ch := workflow.GetSignalChannel(ctx, releaseShardsSignalName) for ctx.Err() == nil { @@ -741,32 +376,8 @@ func (s *shardedWorkflowState) handleReleaseSignals(ctx workflow.Context) { } } -// extractVerifiedCountFromError pulls the partial VerifiedCount that -// wrapBatchVerifyError encoded into a BatchVerifyPartial-typed -// ApplicationError's Details on the activity side. Returns 0 when -// the error didn't come through the verify-phase wrapper (e.g. -// inject-phase failures, ctx errors, non-ApplicationError types), so -// callers can unconditionally fold the result into recordVerified. -func extractVerifiedCountFromError(err error) int64 { - if err == nil { - return 0 - } - appErr, ok := errors.AsType[*temporal.ApplicationError](err) - if !ok || appErr.Type() != batchVerifyPartialErrorType { - return 0 - } - var count int64 - if appErr.Details(&count) != nil { - return 0 - } - return count -} - -// setLastErr latches the first error encountered. Subsequent errors -// are dropped so the root cause is preserved for the workflow's -// returned failure — without the latch, a stuck-shard backstop firing -// on every batch as the workflow tears down would overwrite the -// genuinely interesting first failure. +// setLastErr latches the first error encountered. Subsequent errors are +// dropped so the root cause is preserved for the child's returned failure. func (s *shardedWorkflowState) setLastErr(err error) { if s.lastErr == nil { s.lastErr = err @@ -774,253 +385,34 @@ func (s *shardedWorkflowState) setLastErr(err error) { } // dispatchSlotAvailable returns true when the workflow is below the -// in-flight batch ceiling and is free to spawn another batch. Callers -// that can defer dispatch (the streaming packer) consult this and -// bail out; callers that must dispatch (resume payloads) pair it -// with waitForDispatchSlot. ConcurrentBatchCount is normalised to -// >= 1 at state construction, so no zero-disable path is needed. +// in-flight batch ceiling and is free to spawn another batch. func (s *shardedWorkflowState) dispatchSlotAvailable() bool { return s.pendingDispatches < s.params.ConcurrentBatchCount } -// waitForDispatchSlot blocks the calling workflow coroutine until a -// dispatch slot frees up or lastErr trips. +// waitForDispatchSlot blocks until a dispatch slot frees up or lastErr trips. func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { _ = workflow.Await(ctx, func() bool { return s.lastErr != nil || s.pendingDispatches < s.params.ConcurrentBatchCount }) } -// resumeBatch is one packed dispatch plan: the BatchPayload that will -// become a batch's input, plus the matching per-shard no-progress -// durations. Built up front by packResumeBatchPlan so the dispatch -// loop can unpack any remainder back into ResumeShards if lastErr -// trips mid-dispatch. -type resumeBatch struct { - payload BatchPayload - noProgress map[int32]time.Duration -} - -// dispatchResumeBatches turns the prior cycle's drain payload into a -// fresh round of resume activities, packed across shards up to -// BatchSize per batch. Each shard appears at most once across the -// payload (shardInFlight enforces that only one batch holds a shard -// at a time, and a shard only lands in a drain return while its -// owning batch still has unverified execs on it), so per-shard -// contributions are taken whole and no MaxExecsPerShard cap applies — -// resume carries no inject load so the per-shard blast-radius the -// streaming packer guards against doesn't exist here. -// -// Plans every batch up front, then dispatches one at a time. If -// lastErr latches mid-dispatch, the remaining planned batches are -// unpacked back into s.params.ResumeShards so the recovery bundle -// (and the next CAN cycle) sees them — without the unpack step, a -// failing first resume batch would silently strand all subsequent -// resume entries. -func (s *shardedWorkflowState) dispatchResumeBatches(ctx workflow.Context) { - if len(s.params.ResumeShards) == 0 { - return - } - entries := make([]ResumeShard, 0, len(s.params.ResumeShards)) - for _, rs := range s.params.ResumeShards { - if runCount(rs.Execs) == 0 { - continue - } - entries = append(entries, rs) - } - // We've taken ownership of these entries — anything not - // dispatched gets restored below. - s.params.ResumeShards = nil - slices.SortFunc(entries, func(a, b ResumeShard) int { - return int(a.Shard - b.Shard) - }) - - batches := s.packResumeBatchPlan(entries) - for i, batch := range batches { - // Block until a dispatch slot is free so resume payloads - // can't overshoot ConcurrentBatchCount on cycles that - // carried many shards across CAN. The await also wakes - // immediately when lastErr is set. - s.waitForDispatchSlot(ctx) - if s.lastErr != nil { - s.params.ResumeShards = unpackResumeBatches(batches[i:]) - return - } - s.spawnBatch(ctx, batch.payload, true, batch.noProgress) - } -} - -// packResumeBatchPlan groups ResumeShard entries into batches sized -// at or below BatchSize. Entries are taken whole — shardInFlight only -// admits one batch per shard at a time, so a single shard's payload -// can't be split. -func (s *shardedWorkflowState) packResumeBatchPlan(entries []ResumeShard) []resumeBatch { - var batches []resumeBatch - current := resumeBatch{payload: BatchPayload{}, noProgress: map[int32]time.Duration{}} - packed := 0 - for _, rs := range entries { - rsCount := runCount(rs.Execs) - if packed+rsCount > s.params.BatchSize && packed > 0 { - batches = append(batches, current) - current = resumeBatch{payload: BatchPayload{}, noProgress: map[int32]time.Duration{}} - packed = 0 - } - current.payload[rs.Shard] = rs.Execs - current.noProgress[rs.Shard] = rs.NoProgressDuration - packed += rsCount - } - if packed > 0 { - batches = append(batches, current) - } - return batches -} - -// unpackResumeBatches reverses packResumeBatchPlan, turning planned -// batches back into a flat ResumeShard slice. Used when the dispatch -// loop aborts on lastErr so the undispatched remainder can be carried -// into the recovery bundle / next CAN cycle. The output is sorted by -// shard ID so the slice flowing into the CAN args is deterministic -// across replays. -func unpackResumeBatches(batches []resumeBatch) []ResumeShard { - if len(batches) == 0 { - return nil - } - var out []ResumeShard - for _, b := range batches { - out = append(out, resumeShardsFromPayload(b.payload, b.noProgress)...) - } - slices.SortFunc(out, func(a, b ResumeShard) int { - return int(a.Shard - b.Shard) - }) - return out -} - -// runCount sums runs across BIDs in a single shard's payload entry. -func runCount(byBID map[string][]RunEntry) int { - n := 0 - //workflowcheck:ignore (summation is order-independent) - for _, runs := range byBID { - n += len(runs) - } - return n -} - -// drainBuckets blocks until buckets are empty (success), lastErr -// trips (failure), or the cycle drain timer expires (CAN with -// leftover buckets). Each pass packs everything currently -// dispatchable, then awaits any change in pendingDispatches + -// shardInFlight so the next pass can attempt shards just freed by -// signal-release. -func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { - for { - if s.lastErr != nil || s.cycleDrainTimedOut { - return - } - for s.tryPackStreaming(ctx, true) { //nolint:revive - } - if s.bucketsEmpty() || s.lastErr != nil || s.cycleDrainTimedOut { - return - } - currentPending := s.pendingDispatches - if currentPending == 0 { - s.failDrainBucketsStuck() - return - } - _ = workflow.Await(ctx, s.drainBucketsAwaitPredicate(currentPending)) - } -} - -// failDrainBucketsStuck sets lastErr when buckets are non-empty but no -// batches are in flight — the shard-claim bookkeeping is corrupted, and -// returning silently would proceed to CAN with execs that were never -// dispatched (silent data loss). Failing forces lastErr to propagate -// through run(). -func (s *shardedWorkflowState) failDrainBucketsStuck() { - remaining := 0 - //workflowcheck:ignore (summation is order-independent) - for _, n := range s.bucketCounts { - remaining += n - } - s.setLastErr(temporal.NewNonRetryableApplicationError( - fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), - "DrainBucketsStuck", nil)) -} - -// drainBucketsAwaitPredicate returns true when the drainBuckets loop -// should wake up: lastErr tripped, cycle drain timer fired, a -// dispatch slot just freed, or a new free shard is ready to pack. A -// "free shard" wake-up only counts when there's also a dispatch slot -// to use it, otherwise the outer loop would busy-spin on -// tryPackStreaming returning false against the in-flight cap. -func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { - return func() bool { - if s.lastErr != nil || s.cycleDrainTimedOut { - return true - } - if s.pendingDispatches < currentPending { - return true - } - if !s.dispatchSlotAvailable() { - return false - } - //workflowcheck:ignore (existence check; iteration order does not affect result) - for sh, n := range s.bucketCounts { - if n > 0 && !s.shardInFlight[sh] { - return true - } - } - return false - } -} - -// drainForCAN cancels every in-flight batch and waits for them to -// return AND for the ReleaseShards signal channel to be drained. -// Called on lastErr (terminates with error; drain payload exposed via -// the status query) and on cycle drain timeout (CANs with drained -// state as carry-over). Activities honour cancellation by entering -// drain mode and returning a result whose InFlight carries their -// still-unverified execs; spawnBatch appends those entries to -// s.drainPayload. The signal channel drain is so a final ReleaseShards -// fired by an activity just before it returns doesn't get stranded -// mid-flight, which would leave shardInFlight set for shards the -// activity already considers complete. -// -// Channel.Len() is safe here because handleReleaseSignals has no -// yield points between Receive and the next blocking Receive, so -// Len() == 0 observed across an Await re-evaluation means every -// delivered signal has been processed. -// -// No explicit time bound: a well-behaved activity returns within -// req.DrainGrace (15s default) plus a small idle-cost slack; a -// misbehaved one is bounded by the activity's HeartbeatTimeout (1m). -// In practice this Await unblocks well under a minute. -func (s *shardedWorkflowState) drainForCAN(ctx workflow.Context) { - if s.pendingDispatches == 0 { +// recordVerified accumulates one batch's verified-exec delta into the +// child's running count and emits the per-batch counter metric. +// No-op when verified == 0 so batches with only inject (DisableVerification) +// don't add zeros to the count. +func (s *shardedWorkflowState) recordVerified(verified int64) { + if verified <= 0 { return } - s.cancelActivities() - releaseCh := workflow.GetSignalChannel(ctx, releaseShardsSignalName) - // Wait unconditionally for pendingDispatches to drain — lastErr may - // already be set on entry, but drainPayload, batchExecs, and status - // recovery fields only finalise once every in-flight goroutine has - // returned. - _ = workflow.Await(ctx, func() bool { - return s.pendingDispatches == 0 && releaseCh.Len() == 0 - }) + s.verifiedCount += verified + s.metricsHandler.Counter(metrics.ReplicatedWorkflowCount.Name()).Inc(verified) } // spawnBatch dispatches one batch on a new workflow.Go coroutine. -// Callers must have already marked every shard appearing as a -// top-level key in payload as shardInFlight (the "claim") so the -// packer can see them as busy while picking subsequent batches. The -// activity is run on s.activityCtx so drainForCAN can cancel every -// in-flight batch with a single call. -func (s *shardedWorkflowState) spawnBatch( - ctx workflow.Context, - payload BatchPayload, - resume bool, - noProgressByShard map[int32]time.Duration, -) { +// Callers must have already verified that every shard in payload is free +// (not in shardInFlight) — spawnBatch marks them in-flight here. +func (s *shardedWorkflowState) spawnBatch(ctx workflow.Context, payload BatchPayload) { if payload.totalRuns() == 0 { return } @@ -1033,12 +425,9 @@ func (s *shardedWorkflowState) spawnBatch( NamespaceID: s.namespaceID, Executions: payload, TargetClusterName: s.params.TargetClusterName, - Resume: resume, DisableVerification: s.params.DisableVerification, - NoProgressByShard: noProgressByShard, PerBatchGenerateRPS: s.params.PerBatchGenerateRPS, ShardNoProgress: s.params.ShardNoProgress, - DrainGrace: s.params.DrainGrace, IdleShardCost: s.params.IdleShardCost, } @@ -1049,68 +438,39 @@ func (s *shardedWorkflowState) spawnBatch( held[sh] = true } s.heldByBatch[batchID] = held - s.batchExecs[batchID] = payload s.pendingDispatches++ workflow.Go(ctx, func(coroCtx workflow.Context) { defer func() { s.pendingDispatches-- - // Clear any shards we still hold — signal-released - // shards have already been cleared from shardInFlight - // by handleReleaseSignals and may by now belong to a - // subsequent batch's claim. + // Clear any shards we still hold. Signal-released shards have + // already been cleared from shardInFlight by handleReleaseSignals + // and may by now belong to a subsequent batch's claim. //workflowcheck:ignore (deletes are commutative; order-independent) for sh := range s.heldByBatch[batchID] { delete(s.shardInFlight, sh) } delete(s.heldByBatch, batchID) }() - result, err := s.replicateBatch(coroCtx, s.activityCtx, req) + result, err := s.replicateBatch(coroCtx, req) if err != nil { - if temporal.IsCanceledError(err) { - // Cancel-before-start: the activity body never ran, - // so no result is available. Leaving batchExecs[batchID] - // intact lets collectRecoveredBucketsForCarryover - // re-bucket the execs as fresh inject+verify work next - // cycle (only meaningful on the CAN path; on lastErr - // they surface via the status query only). - return - } - // Activity errored after partial verify — the SDK discards - // the result on failure, but wrapBatchVerifyError on the - // activity side carries the partial doneCount through as - // ApplicationError details. Fold it into the running count - // so ReplicatedWorkflowCount reflects work actually done. - s.recordVerified(coroCtx, extractVerifiedCountFromError(err)) s.setLastErr(err) return } - // Activity body ran and returned cleanly — either a - // clean completion (empty InFlight) or a drained cancel - // (InFlight carries the still-unverified execs; reached - // only on lastErr or cycle drain timeout). CompletedShards - // is informational; the defer above clears heldByBatch + - // shardInFlight either way. - if len(result.InFlight) > 0 { - s.drainPayload = append(s.drainPayload, result.InFlight...) - } - s.recordVerified(coroCtx, result.VerifiedCount) - delete(s.batchExecs, batchID) + s.recordVerified(result.VerifiedCount) }) } -// tryPackStreaming attempts to pack and dispatch one batch from -// s.buckets. Returns true if a batch was dispatched. -// -// No per-shard or total-bucket threshold: as soon as any free shard -// has any execs and a dispatch slot is open, a batch fires. Safety -// is enforced outside the packer — MaxExecsPerShard caps a single -// shard's contribution to a batch (so one hot shard can't dominate), -// ConcurrentBatchCount caps in-flight batches, and PerBatchGenerateRPS -// caps the per-batch source RPS. Within those bounds the packer's -// sole job is to make progress every chance it gets. +// tryPackStreaming attempts to pack and dispatch one batch from s.buckets. +// Returns true if a batch was dispatched. Uses effectiveMaxExecsPerShard +// as the per-shard cap so the half-rate constraint propagates into every +// batch packed while the child is throttled or cut. // -// See shardIDsByPackPriority for the relax-mode ordering rationale. +// No per-shard or total-bucket threshold: as soon as any free shard has +// any execs and a dispatch slot is open, a batch fires. Safety is enforced +// outside the packer — effectiveMaxExecsPerShard caps a single shard's +// contribution, ConcurrentBatchCount caps in-flight batches, and +// PerBatchGenerateRPS caps the per-batch source RPS. func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool) bool { if s.lastErr != nil || s.params.BatchSize <= 0 || s.params.MaxExecsPerShard <= 0 { return false @@ -1123,6 +483,7 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool return false } + effectiveCap := s.effectiveMaxExecsPerShard() payload := BatchPayload{} packed := 0 for _, sh := range shardIDs { @@ -1130,7 +491,7 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool if room <= 0 { break } - take := min(s.params.MaxExecsPerShard, s.bucketCounts[sh], room) + take := min(effectiveCap, s.bucketCounts[sh], room) if take == 0 { continue } @@ -1140,12 +501,12 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool if packed == 0 { return false } - s.spawnBatch(ctx, payload, false, nil) + s.spawnBatch(ctx, payload) return true } -// bucketsEmpty reports whether every shard's bucket is empty. Reads -// from the sidecar count map so it's O(#shards), not O(#runs). +// bucketsEmpty reports whether every shard's bucket is empty. Reads from +// the sidecar count map so it's O(#shards), not O(#runs). func (s *shardedWorkflowState) bucketsEmpty() bool { //workflowcheck:ignore (existence check; iteration order does not affect result) for _, n := range s.bucketCounts { @@ -1156,23 +517,22 @@ func (s *shardedWorkflowState) bucketsEmpty() bool { return true } -// shardIDsByPackPriority returns free, non-empty shard IDs in the -// order the packer should consider them. Deterministic across -// replays: ordering is derived from workflow state (bucketCounts) -// with shard ID as a stable tiebreaker. +// shardIDsByPackPriority returns free, non-empty shard IDs in the order +// the packer should consider them. Deterministic across replays: ordering +// is derived from workflow state (bucketCounts) with shard ID as a stable +// tiebreaker. // -// relax=false (streaming): fullest first. Packer naturally produces -// large, predictable batches when work is plentiful and small ones -// when it isn't — either way it ships rather than waiting. +// relax=false (streaming): fullest first. Packer naturally produces large, +// predictable batches when work is plentiful and small ones when it isn't +// — either way it ships rather than waiting. // -// relax=true (drain): hot shards (count > MaxExecsPerShard) first, -// fullest within hot. These need >1 round trip to drain, so total -// drain wall-clock is bounded by the heaviest shard; starting their -// pipelines first is the dominant lever. After all hot shards are -// claimed, remaining batch capacity fills from smallest cold buckets -// (ascending count) so light shards clear out quickly — nothing is -// arriving in drain, so waiting for cold buckets to grow is wasted -// wall-clock. +// relax=true (drain): hot shards (count > effectiveMaxExecsPerShard) first, +// fullest within hot. These need >1 round trip to drain, so total drain +// wall-clock is bounded by the heaviest shard; starting their pipelines +// first is the dominant lever. After all hot shards are claimed, remaining +// batch capacity fills from smallest cold buckets (ascending count) so +// light shards clear out quickly — nothing is arriving in drain, so +// waiting for cold buckets to grow is wasted wall-clock. func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { out := make([]int32, 0, len(s.bucketCounts)) //workflowcheck:ignore (output is sorted below before any observable use) @@ -1183,7 +543,7 @@ func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { out = append(out, sh) } if relax { - maxPerShard := s.params.MaxExecsPerShard + maxPerShard := s.effectiveMaxExecsPerShard() slices.SortFunc(out, func(a, b int32) int { aHot, bHot := s.bucketCounts[a] > maxPerShard, s.bucketCounts[b] > maxPerShard switch { @@ -1212,3 +572,124 @@ func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { }) return out } + +// addToBucket appends one run to the (shard, BID) bucket and bumps +// the sidecar count. +func (s *shardedWorkflowState) addToBucket(shard int32, businessID string, run RunEntry) { + if s.buckets[shard] == nil { + s.buckets[shard] = map[string][]RunEntry{} + } + s.buckets[shard][businessID] = append(s.buckets[shard][businessID], run) + s.bucketCounts[shard]++ +} + +// takeFromBucket consumes up to n runs from the given shard and returns +// them grouped by BID. Walks BIDs in alphabetical order so the resulting +// payload is deterministic across replays; takes whole per-BID runs only +// as needed to reach n. Empties the shard from s.buckets / s.bucketCounts +// when nothing remains. +func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]RunEntry { + if n <= 0 { + return nil + } + byBID := s.buckets[shard] + if len(byBID) == 0 { + return nil + } + bids := make([]string, 0, len(byBID)) + //workflowcheck:ignore (bids is sorted before use) + for bid := range byBID { + bids = append(bids, bid) + } + slices.Sort(bids) + + out := map[string][]RunEntry{} + taken := 0 + for _, bid := range bids { + if taken >= n { + break + } + runs := byBID[bid] + take := min(len(runs), n-taken) + // append([]RunEntry(nil), ...) gives the output its own backing + // array — keeps the workflow's leftover slice (byBID[bid][take:]) + // and the activity's input independent in case either side appends. + out[bid] = append([]RunEntry(nil), runs[:take]...) + if take == len(runs) { + delete(byBID, bid) + } else { + byBID[bid] = runs[take:] + } + taken += take + } + s.bucketCounts[shard] -= taken + if s.bucketCounts[shard] <= 0 { + delete(s.bucketCounts, shard) + delete(s.buckets, shard) + } + return out +} + +// drainBuckets blocks until buckets are empty (success) or lastErr trips +// (failure). Each pass packs everything currently dispatchable, then +// awaits any change in pendingDispatches + shardInFlight so the next pass +// can attempt shards just freed by signal-release. +func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { + for { + if s.lastErr != nil { + return + } + for s.tryPackStreaming(ctx, true) { //nolint:revive + } + if s.bucketsEmpty() || s.lastErr != nil { + return + } + currentPending := s.pendingDispatches + if currentPending == 0 { + s.failDrainBucketsStuck() + return + } + _ = workflow.Await(ctx, s.drainBucketsAwaitPredicate(currentPending)) + } +} + +// failDrainBucketsStuck sets lastErr when buckets are non-empty but no +// batches are in flight — the shard-claim bookkeeping is corrupted, and +// returning silently would proceed with execs that were never dispatched +// (silent data loss). Failing forces the child to return the error. +func (s *shardedWorkflowState) failDrainBucketsStuck() { + remaining := 0 + //workflowcheck:ignore (summation is order-independent) + for _, n := range s.bucketCounts { + remaining += n + } + s.setLastErr(temporal.NewNonRetryableApplicationError( + fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), + "DrainBucketsStuck", nil)) +} + +// drainBucketsAwaitPredicate returns true when the drainBuckets loop +// should wake up: lastErr tripped, a dispatch slot just freed, or a new +// free shard is ready to pack. A "free shard" wake-up only counts when +// there's also a dispatch slot to use it, otherwise the outer loop would +// busy-spin on tryPackStreaming returning false against the in-flight cap. +func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { + return func() bool { + if s.lastErr != nil { + return true + } + if s.pendingDispatches < currentPending { + return true + } + if !s.dispatchSlotAvailable() { + return false + } + //workflowcheck:ignore (existence check; iteration order does not affect result) + for sh, n := range s.bucketCounts { + if n > 0 && !s.shardInFlight[sh] { + return true + } + } + return false + } +} diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 4ae20c9352a..cba3931294b 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -2,7 +2,6 @@ package migration import ( "context" - "errors" "fmt" "strconv" "sync" @@ -13,7 +12,6 @@ import ( "github.com/stretchr/testify/require" "go.temporal.io/api/workflowservice/v1" "go.temporal.io/sdk/activity" - "go.temporal.io/sdk/converter" "go.temporal.io/sdk/temporal" "go.temporal.io/sdk/testsuite" "go.temporal.io/sdk/workflow" @@ -110,7 +108,8 @@ func metadataResponseFor(shardCount int32) func(context.Context, MetadataRequest // registerShardedScaffolding registers the GetMetadata + CountWorkflow // stubs every sharded test needs before the page loop runs, plus the // task-queue-user-data child workflow + its activity so the parent's -// terminal Await on Done resolves. +// terminal Await on Done resolves. Also registers shardedForceReplicationWorker +// so parent tests can spawn it inline as a child. func registerShardedScaffolding(env *testsuite.TestWorkflowEnvironment, shardCount int32) { registerShardedScaffoldingWithSeed(env, shardCount, func(_ context.Context, _ TaskQueueUserDataReplicationParamsWithNamespace) error { return nil @@ -134,9 +133,47 @@ func registerShardedScaffoldingWithSeed( }, activity.RegisterOptions{Name: "CountWorkflow"}) env.RegisterWorkflowWithOptions(ForceTaskQueueUserDataReplicationWorkflow, workflow.RegisterOptions{Name: forceTaskQueueUserDataReplicationWorkflow}) env.RegisterActivityWithOptions(seed, activity.RegisterOptions{Name: "SeedReplicationQueueWithUserDataEntries"}) + // Register child worker so parent tests can spawn it inline. + env.RegisterWorkflow(shardedForceReplicationWorker) } -// ---- Tests ---- +// makeChildParams returns a shardedChildParams with sensible defaults for +// child-direct tests. Callers override fields as needed. +func makeChildParams(shardCount int32) shardedChildParams { + return shardedChildParams{ + Namespace: "test-ns", + NamespaceID: testNamespaceID, + TargetClusterName: "remote_cluster", + TargetShardCount: shardCount, + BatchSize: 100, + MaxExecsPerShard: 50, + ListWorkflowsPageSize: 1000, + ConcurrentBatchCount: 1, + PerBatchGenerateRPS: 30, + ShardNoProgress: time.Hour, + IdleShardCost: time.Hour, + } +} + +// hasAppErrType walks an error chain looking for a *temporal.ApplicationError +// with the given type. Use this when the error may be wrapped inside a +// child-workflow or fmt.Errorf chain. +func hasAppErrType(err error, wantType string) bool { + for err != nil { + if appErr, ok := err.(*temporal.ApplicationError); ok && appErr.Type() == wantType { + return true + } + type unwrapper interface{ Unwrap() error } + u, ok := err.(unwrapper) + if !ok { + break + } + err = u.Unwrap() + } + return false +} + +// ---- Parent workflow tests ---- // TestSharded_HappyPath_SingleCycle: a small workload exhausts in one // cycle, every batch returns clean completion, no CAN, no resume. @@ -180,139 +217,6 @@ func TestSharded_HappyPath_SingleCycle(t *testing.T) { require.Len(t, shardsSeen, 4, "every shard should be represented") } -// TestSharded_ResumeShards_Packed: a non-empty ResumeShards in params -// gets packed into multi-shard batches up to BatchSize. Asserts that -// no resume batch exceeds BatchSize and the per-shard contributions -// match the input. -func TestSharded_ResumeShards_Packed(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 8) - env.RegisterActivityWithOptions(pageThrough(nil, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) - - // 8 shards, 15 unverified execs each = 120 total. With BatchSize=100, - // we expect 2 batches: e.g., [0..5] (90 execs) + [5+, 6, 7] (30) or - // similar — depends on greedy packing. - resumeShards := make([]ResumeShard, 8) - for s := range 8 { - resumeShards[s] = ResumeShard{ - Shard: int32(s), - Execs: makeExecsForShard(int32(s), 15), - } - } - - var ( - mu sync.Mutex - batches [][]int32 - ) - env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { - require.True(t, req.Resume, "all resume-dispatched batches must have Resume=true") - require.LessOrEqual(t, req.Executions.totalRuns(), 100, "batch must not exceed BatchSize") - mu.Lock() - batches = append(batches, req.Executions.sortedShards()) - mu.Unlock() - return replicateBatchResult{}, nil - }, activity.RegisterOptions{Name: "ReplicateBatch"}) - - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - ResumeShards: resumeShards, - }) - - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) - - // Confirm every shard was covered exactly once across all batches. - covered := map[int32]int{} - for _, b := range batches { - for _, sh := range b { - covered[sh]++ - } - } - for s := range int32(8) { - require.Equal(t, 1, covered[s], "shard %d should be covered exactly once", s) - } - require.GreaterOrEqual(t, len(batches), 2, "should pack into at least 2 batches given 120 execs / BatchSize=100") -} - -// TestSharded_ReleaseShards_FreesShardForReuse: an activity sends a -// ReleaseShards signal mid-flight. The workflow must clear the shard -// from shardInFlight so the packer can dispatch a fresh batch -// targeting that shard *while the original activity is still running* -// — the slot in ConcurrentBatchCount is still claimed by batch 1, so -// batch 2 can only dispatch if signal-release worked. -// -// ConcurrentBatchCount=2 is explicit: defaultConcurrentBatchCount(2)=1 -// would gate batch 2 on the dispatch slot regardless of shard state, so -// the test couldn't distinguish "release-from-flight" from -// "batch 1 returned and freed the slot". -func TestSharded_ReleaseShards_FreesShardForReuse(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 2) - - // Two pages of execs on the same shards. The second batch can dispatch - // only when batch 1 releases its shards. - phase1 := makeExecs(2, 10) - phase2 := makeExecs(2, 10) - all := append(append([]*ExecutionInfo{}, phase1...), phase2...) - env.RegisterActivityWithOptions(pageThrough(all, 20), activity.RegisterOptions{Name: "ListWorkflows"}) - - var ( - mu sync.Mutex - batches []*shardedBatchReq - secondStarted = make(chan struct{}) - secondOnce sync.Once - releaseObserved atomic.Bool - ) - env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { - mu.Lock() - batches = append(batches, req) - idx := len(batches) - mu.Unlock() - - switch idx { - case 1: - // Signal-release, then block here until batch 2 actually - // dispatches. If signal-release wires through, the workflow - // dispatches batch 2 concurrently; if it doesn't, batch 2 - // can't run until this activity returns (the safety timeout - // below). - env.SignalWorkflow("ReleaseShards", releaseShardsPayload{ - BatchID: req.BatchID, - Shards: req.Executions.sortedShards(), - }) - select { - case <-secondStarted: - releaseObserved.Store(true) - case <-time.After(30 * time.Second): - // Safety release so the test fails the assertion rather - // than hanging indefinitely. Generous bound because CI - // can be slow; the happy path returns immediately. - } - case 2: - secondOnce.Do(func() { close(secondStarted) }) - default: - } - return replicateBatchResult{}, nil - }, activity.RegisterOptions{Name: "ReplicateBatch"}) - - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - ConcurrentBatchCount: 2, - }) - - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) - require.Len(t, batches, 2, "expected two batches across the two pages") - require.True(t, releaseObserved.Load(), - "signal-release should let batch 2 dispatch while batch 1 still holds its dispatch slot") -} - // TestSharded_ShardNoProgress_FailsWorkflow: activity returns a // non-retryable ShardNoProgress error → workflow fails with that // error, no CAN. @@ -337,215 +241,9 @@ func TestSharded_ShardNoProgress_FailsWorkflow(t *testing.T) { require.True(t, env.IsWorkflowCompleted()) err := env.GetWorkflowError() require.Error(t, err) - var appErr *temporal.ApplicationError - require.ErrorAs(t, err, &appErr) - require.Equal(t, "ShardNoProgress", appErr.Type()) -} - -// TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover: a -// non-empty InFlight in the activity's replicateBatchResult must -// populate drainPayload and end up as ResumeShards in the CAN -// carry-over. -// -// Tests the workflow plumbing only. In production, an activity -// returns InFlight after entering drain mode and grace-expiring; here -// we exercise the same code path by returning InFlight from a -// non-cancelled run, because the testsuite delivers cancellation as -// a CanceledError without preserving the activity's returned result. -// The dispatch coroutine's err == nil branch is what we're testing — -// it doesn't care whether the activity was cancelled or not, only -// whether the returned result has InFlight to fold into drainPayload. -func TestSharded_DrainResult_FromActivityResult_FeedsCANCarryover(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 10) - - // Page 1 returns a multi-shard population so the streaming packer - // has something to dispatch under the pinned BatchSize / - // MaxExecsPerShard the test sets below. The activity flips - // CAN-suggested from inside its body before returning, so by the - // time it has handed back its InFlight the workflow is committed - // to CAN — but without going through cancel, which the testsuite - // delivers as a CanceledError regardless of any result the - // activity returned. - pageExecs := makeExecs(10, 10) - // Drained exec mirrors a real input row so the simulated drain - // payload would be a valid response from a real activity. drainedBID - // is the first exec; drainedShard is the shard the workflow will - // hash it to. - drainedBID := pageExecs[0].BusinessID - const drainedRunID = "run-drained" - drainedShard := common.WorkflowIDToHistoryShard(testNamespaceID, drainedBID, 10) - env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { - return &listWorkflowsResponse{ - Executions: pageExecs, - NextPageToken: []byte("more"), - }, nil - }, activity.RegisterOptions{Name: "ListWorkflows"}) - - env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - // Hop to the main loop goroutine to flip CAN-suggested; the - // activity goroutine writing workflowInfo directly would race - // the workflow coroutine reading it at the top of its page loop. - env.RegisterDelayedCallback(func() { env.SetContinueAsNewSuggested(true) }, 0) - return replicateBatchResult{ - InFlight: []ResumeShard{{ - Shard: drainedShard, - Execs: map[string][]RunEntry{drainedBID: {{RunID: drainedRunID}}}, - NoProgressDuration: 42 * time.Second, - }}, - }, nil - }, activity.RegisterOptions{Name: "ReplicateBatch"}) - - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - BatchSize: 100, - MaxExecsPerShard: 10, - }) - - require.True(t, env.IsWorkflowCompleted()) - err := env.GetWorkflowError() - require.Error(t, err, "workflow should CAN, not return success") - - var canErr *workflow.ContinueAsNewError - require.ErrorAs(t, err, &canErr, "error should be ContinueAsNewError") - - var nextParams ShardedForceReplicationParams - require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) - require.NotEmpty(t, nextParams.ResumeShards, "InFlight from a returned activity should appear in resume payload") - require.Equal(t, drainedShard, nextParams.ResumeShards[0].Shard) - runs, ok := nextParams.ResumeShards[0].Execs[drainedBID] - require.True(t, ok, "drained business ID should appear in nested resume payload") - require.Equal(t, []RunEntry{{RunID: drainedRunID}}, runs) - require.Equal(t, 42*time.Second, nextParams.ResumeShards[0].NoProgressDuration) -} - -// TestSharded_ListingDoneAcrossCAN_SkipsPageLoop pins down that a -// post-listing CAN cycle does not re-enter ListWorkflows. The prior -// cycle finished pagination (NextPageToken empty) but still had carry- -// over (e.g. InFlight from a returned activity), so it CAN'd. Without -// the ContinuedAsNewCount > 0 guard in the page loop, the next cycle -// can't distinguish "haven't started listing" from "finished listing" -// and would re-list page 1, double-enqueueing every visible execution. -func TestSharded_ListingDoneAcrossCAN_SkipsPageLoop(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 4) - - var listCalls atomic.Int32 - env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { - listCalls.Add(1) - return &listWorkflowsResponse{}, nil - }, activity.RegisterOptions{Name: "ListWorkflows"}) - - // Resume-only batches; no fresh inject work. - env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { - require.True(t, req.Resume, "this cycle should only dispatch resume batches; no listing should occur") - return replicateBatchResult{ - CompletedShards: req.Executions.sortedShards(), - }, nil - }, activity.RegisterOptions{Name: "ReplicateBatch"}) - - // Mimic a CAN-cycle entry: prior cycle exhausted pagination - // (NextPageToken empty) and CAN'd because it had ResumeShards to - // carry over. ContinuedAsNewCount > 0 is the signal that empty - // NextPageToken means "done", not "haven't started". - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - ContinuedAsNewCount: 1, - ResumeShards: []ResumeShard{{ - Shard: 0, - Execs: map[string][]RunEntry{"wf-resume": {{RunID: "run-resume"}}}, - }}, - // On a real CAN cycle the prior cycle's child kicked off in - // cycle 0, so its signal arrives once; subsequent cycles see - // Done=true in their input params and skip both the kickoff - // and the terminal Await. The test mirrors that. - TaskQueueUserDataReplicationStatus: TaskQueueUserDataReplicationStatus{Done: true}, - }) - - require.True(t, env.IsWorkflowCompleted()) - require.NoError(t, env.GetWorkflowError()) - require.Zero(t, listCalls.Load(), "ListWorkflows must not be called on a CAN cycle that inherited an exhausted page token") -} - -// TestSharded_CancelBeforeStart_NoLostExecs pins down recovery when -// the cycle drain timer fires before any dispatched activity can run: -// the activities return CanceledError with no result, and the recovery -// path re-buckets the input execs into RecoveredBuckets so the next -// cycle dispatches them as fresh inject+verify batches. -func TestSharded_CancelBeforeStart_NoLostExecs(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 4) - - execs := makeExecs(4, 5) // 20 execs across 4 shards - pageServed := false - env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { - if pageServed { - return &listWorkflowsResponse{}, nil - } - pageServed = true - // Trigger CAN-suggested as soon as this page is served so the - // page loop bails and the cycle drain timer starts. Hop to the - // main loop goroutine — writing workflowInfo directly from the - // activity goroutine would race the workflow coroutine's read. - env.RegisterDelayedCallback(func() { env.SetContinueAsNewSuggested(true) }, 0) - return &listWorkflowsResponse{ - Executions: execs, - NextPageToken: []byte("more"), - }, nil - }, activity.RegisterOptions{Name: "ListWorkflows"}) - - // Activity that responds to ctx cancellation by returning a - // CanceledError with no result — simulating the - // "cancelled before any work done" path. - var ( - batchCount atomic.Int32 - cancelledIDs []int64 - muIDs sync.Mutex - ) - env.RegisterActivityWithOptions(func(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { - batchCount.Add(1) - <-ctx.Done() - muIDs.Lock() - cancelledIDs = append(cancelledIDs, req.BatchID) - muIDs.Unlock() - return replicateBatchResult{}, temporal.NewCanceledError() - }, activity.RegisterOptions{Name: "ReplicateBatch"}) - - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - // Short cycle drain timeout so the test hits the timer-fired - // cancel path quickly rather than waiting the production default. - CycleDrainTimeout: time.Millisecond, - }) - - require.True(t, env.IsWorkflowCompleted()) - err := env.GetWorkflowError() - require.Error(t, err) - var canErr *workflow.ContinueAsNewError - require.ErrorAs(t, err, &canErr, "expected CAN, got %v", err) - - var nextParams ShardedForceReplicationParams - require.NoError(t, converter.GetDefaultDataConverter().FromPayloads(canErr.Input, &nextParams)) - - // Count execs the next cycle would re-dispatch. The activity - // body never ran in this race, so execs land in RecoveredBuckets - // (fresh inject+verify) rather than ResumeShards (verify-only). - recovered := nextParams.RecoveredBuckets.totalRuns() - - t.Logf("dispatched %d batches, cancelled %d, recovered execs %d (expected %d)", - batchCount.Load(), len(cancelledIDs), recovered, len(execs)) - - require.Equal(t, len(execs), recovered, "every dispatched exec should land in RecoveredBuckets when its activity is cancelled before it can run") - require.Empty(t, nextParams.ResumeShards, "no ResumeShards — activity never ran, never injected, so no resume work") + // The error chain is: WorkflowExecutionError → childErr (wrapError) → ApplicationError{ShardNoProgress}. + // Walk the chain to find the first ApplicationError with type ShardNoProgress. + require.True(t, hasAppErrType(err, "ShardNoProgress"), "expected ShardNoProgress in error chain, got: %v", err) } // TestSharded_DisableVerification_NoVerifiedCount: with verification @@ -721,267 +419,188 @@ func TestSharded_TaskQueueReplicationFailure(t *testing.T) { require.Contains(t, status.TaskQueueUserDataReplicationStatus.FailureMessage, "namespace is required") } -// TestSharded_RecoveryBundle_OnBatchError: a batch returns a -// non-retryable error mid-cycle; the workflow latches lastErr, drains -// in-flight via cancellation, returns the error, and the status query -// reports a non-empty recovery bundle so an operator can start a -// fresh run with all unverified execs preserved. -func TestSharded_RecoveryBundle_OnBatchError(t *testing.T) { - suite := &testsuite.WorkflowTestSuite{} - env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 4) - - // One page of execs across 4 shards. - execs := makeExecs(4, 5) // 20 execs total - env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) +// ---- Child workflow tests ---- - // Every batch fails non-retryably so lastErr latches on the first - // return and subsequent in-flight batches are cancelled. - env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( - "batch failed", "BatchFailed", nil) - }, activity.RegisterOptions{Name: "ReplicateBatch"}) +// childWorkerID is the deterministic workflow ID given to the child worker +// spawned by childDirectRunner. Knowing this ID in advance lets the relay +// goroutine forward promotion signals to the child via SignalExternalWorkflow. +const childWorkerID = "test-sharded-child-worker" - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", +// childDirectRunner is a test-only wrapper workflow that spawns +// shardedForceReplicationWorker as a child via workflow.ExecuteChildWorkflow. +// +// Two problems it solves: +// +// 1. The production child workflow accesses +// workflow.GetInfo(ctx).ParentWorkflowExecution.ID; this panics when the +// workflow has no parent (nil pointer dereference). childDirectRunner +// provides a real parent so ParentWorkflowExecution is non-nil. +// +// 2. env.SignalWorkflow targets the top-level execution (childDirectRunner), +// not the child execution. childDirectRunner therefore relays the +// shardedResumeFullSignalName signal to the child via +// workflow.SignalExternalWorkflow so promotion signals from test +// callbacks reach the production code. +func childDirectRunner(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { + childCtx := workflow.WithChildOptions(ctx, workflow.ChildWorkflowOptions{ + WorkflowID: childWorkerID, }) - require.True(t, env.IsWorkflowCompleted()) - err := env.GetWorkflowError() - require.Error(t, err) - var appErr *temporal.ApplicationError - require.ErrorAs(t, err, &appErr) - require.Equal(t, "BatchFailed", appErr.Type()) - - envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) - require.NoError(t, qErr) - var status ForceReplicationStatus - require.NoError(t, envValue.Get(&status)) - - // Recovery bundle: the cancelled in-flight batches go into - // RecoveryBuckets (collectRecoveredBuckets on batchExecs). The - // failed batch's execs land there too — its activity attempt - // errored after running, so batchExecs[id] is still populated. - require.NotEmpty(t, status.RecoveryBuckets, - "failed run must expose its in-flight execs as RecoveryBuckets") - recovered := 0 - for _, byBID := range status.RecoveryBuckets { - for _, runs := range byBID { - recovered += len(runs) + // Relay promotion signal: when the test sends env.SignalWorkflow(shardedResumeFullSignalName, ...), + // it reaches this wrapper. Forward it to the child so the production + // handleReleaseSignals / resume goroutine inside shardedForceReplicationWorker receives it. + workflow.Go(ctx, func(gCtx workflow.Context) { + ch := workflow.GetSignalChannel(gCtx, shardedResumeFullSignalName) + var dummy struct{} + for ch.Receive(gCtx, &dummy) { + _ = workflow.SignalExternalWorkflow(gCtx, childWorkerID, "", shardedResumeFullSignalName, struct{}{}).Get(gCtx, nil) } - } - require.Equal(t, len(execs), recovered, - "every listed exec should be recoverable via the bundle") + }) + + var result shardedChildResult + err := workflow.ExecuteChildWorkflow(childCtx, shardedForceReplicationWorker, params).Get(ctx, &result) + return result, err } -// TestSharded_RecoveryBundle_PreservesUndispatchedResumeShards: when -// the first resume batch errors and latches lastErr, the remaining -// undispatched resume entries must be preserved in the recovery -// bundle rather than silently dropped. -func TestSharded_RecoveryBundle_PreservesUndispatchedResumeShards(t *testing.T) { +// TestChild_HappyPath_TerminalChild: a single page of execs fits in one +// batch; the child verifies them all and returns ReachedEnd=true. +func TestChild_HappyPath_TerminalChild(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 8) - env.RegisterActivityWithOptions(pageThrough(nil, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) - - // 8 shards × 60 execs each = 480 unverified. BatchSize=100 → - // dispatchResumeBatches plans ~5 batches; the first one fails, - // latches lastErr, and the remaining 4 must be preserved. - resumeShards := make([]ResumeShard, 8) - for s := range 8 { - resumeShards[s] = ResumeShard{ - Shard: int32(s), - Execs: makeExecsForShard(int32(s), 60), - } - } + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) - env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - return replicateBatchResult{}, temporal.NewNonRetryableApplicationError( - "resume batch failed", "ResumeFailed", nil) + execs := makeExecs(4, 5) // 20 execs across 4 shards, single page + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{ + VerifiedCount: 5, + CompletedShards: req.Executions.sortedShards(), + }, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - ConcurrentBatchCount: 1, // serialize so we can observe the early-bail behaviour - ResumeShards: resumeShards, - }) + env.ExecuteWorkflow(childDirectRunner, makeChildParams(4)) require.True(t, env.IsWorkflowCompleted()) - require.Error(t, env.GetWorkflowError()) - - envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) - require.NoError(t, qErr) - var status ForceReplicationStatus - require.NoError(t, envValue.Get(&status)) + require.NoError(t, env.GetWorkflowError()) - // Every shard must show up exactly once across RecoveryResumeShards - // (still-undispatched, batched-but-cancelled, or the failed batch - // itself — all paths fold back into the recovery bundle). - recoveredRuns := 0 - for _, rs := range status.RecoveryResumeShards { - for _, runs := range rs.Execs { - recoveredRuns += len(runs) - } - } - for _, byBID := range status.RecoveryBuckets { - for _, runs := range byBID { - recoveredRuns += len(runs) - } - } - require.Equal(t, 8*60, recoveredRuns, - "all 480 resume execs should be recoverable; got %d", recoveredRuns) + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.True(t, result.ReachedEnd, "single-page namespace should set ReachedEnd=true") + require.Equal(t, int64(5), result.VerifiedCount) } -// TestSharded_RecoveryBundle_TracksCurrentPageToken: a List error -// after the first page should preserve the page token of the next -// page to read, not the start-of-run token. -func TestSharded_RecoveryBundle_TracksCurrentPageToken(t *testing.T) { +// TestChild_VerifiedCountAccumulates: two batches (two pages of execs), +// each returning VerifiedCount=3; the child accumulates to 6. +func TestChild_VerifiedCountAccumulates(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 2) + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) - // First call succeeds; second call errors. Workflow processes page - // 1 successfully, then fails listing page 2 with NextPageToken - // pointing past page 1. - var listCalls atomic.Int32 - page1 := makeExecs(2, 3) - env.RegisterActivityWithOptions(func(_ context.Context, req *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { - n := listCalls.Add(1) - if n == 1 { - return &listWorkflowsResponse{ - Executions: page1, - NextPageToken: []byte("page-2"), - }, nil - } - return nil, temporal.NewNonRetryableApplicationError( - "list page 2 failed", "ListFailed", nil) - }, activity.RegisterOptions{Name: "ListWorkflows"}) - - // Batches succeed so page 1 doesn't pollute the recovery bundle — - // we want the page-token assertion isolated. + // Two pages of 10 execs each across 2 shards: makeExecs(2,5) produces + // 10 execs (5 per shard). With pageSize=5, the pager yields two fetches. + all := makeExecs(2, 5) // 10 execs; pageSize=5 → 2 pages of 5 + env.RegisterActivityWithOptions(pageThrough(all, 5), activity.RegisterOptions{Name: "ListWorkflows"}) env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - return replicateBatchResult{}, nil + return replicateBatchResult{VerifiedCount: 3}, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - }) + params := makeChildParams(2) + params.ConcurrentBatchCount = 1 + env.ExecuteWorkflow(childDirectRunner, params) require.True(t, env.IsWorkflowCompleted()) - require.Error(t, env.GetWorkflowError()) - - envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) - require.NoError(t, qErr) - var status ForceReplicationStatus - require.NoError(t, envValue.Get(&status)) + require.NoError(t, env.GetWorkflowError()) - require.Equal(t, []byte("page-2"), status.RecoveryNextPageToken, - "RecoveryNextPageToken should reflect where listing was about to resume, not start-of-run") + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.Equal(t, int64(6), result.VerifiedCount) + require.True(t, result.ReachedEnd) } -// TestWrapBatchVerifyError_RoundTrip: the workflow must be able to -// recover the partial VerifiedCount that the activity encoded on a -// failed batch return. Covers retryability preservation, the -// no-progress short-circuit, and inner-error reachability via -// Unwrap (so consumers that care about the underlying Type can walk -// past the wrapper). -func TestWrapBatchVerifyError_RoundTrip(t *testing.T) { - t.Run("non-zero count survives wrap; inner reachable via Unwrap", func(t *testing.T) { - cause := temporal.NewNonRetryableApplicationError("stuck", "ShardNoProgress", nil) - wrapped := wrapBatchVerifyError(cause, 42) - - // Outer wrapper carries the partial-verify tag and the count. - var appErr *temporal.ApplicationError - require.ErrorAs(t, wrapped, &appErr) - require.Equal(t, batchVerifyPartialErrorType, appErr.Type()) - require.True(t, appErr.NonRetryable()) - require.Equal(t, int64(42), extractVerifiedCountFromError(wrapped)) - - // Inner identity is preserved via Cause / Unwrap — consumers - // that key off the underlying type still get there. - var inner *temporal.ApplicationError - require.ErrorAs(t, appErr.Unwrap(), &inner) - require.Equal(t, "ShardNoProgress", inner.Type()) - }) +// TestChild_DisableVerification_ReachesEnd: with DisableVerification=true +// the child completes with VerifiedCount=0 and ReachedEnd=true. +func TestChild_DisableVerification_ReachesEnd(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) - t.Run("zero count returns cause unchanged", func(t *testing.T) { - cause := temporal.NewApplicationError("trivial", "X") - require.Same(t, cause, wrapBatchVerifyError(cause, 0)) - require.Equal(t, int64(0), extractVerifiedCountFromError(cause)) - }) + execs := makeExecs(2, 5) + env.RegisterActivityWithOptions(pageThrough(execs, 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: 0, + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) - t.Run("non-ApplicationError cause is wrapped and remains reachable", func(t *testing.T) { - cause := errors.New("plain failure") - wrapped := wrapBatchVerifyError(cause, 7) - var appErr *temporal.ApplicationError - require.ErrorAs(t, wrapped, &appErr) - require.Equal(t, batchVerifyPartialErrorType, appErr.Type()) - require.Equal(t, int64(7), extractVerifiedCountFromError(wrapped)) - require.ErrorIs(t, wrapped, cause) - }) + params := makeChildParams(2) + params.DisableVerification = true + env.ExecuteWorkflow(childDirectRunner, params) - t.Run("unrelated ApplicationError returns 0", func(t *testing.T) { - // An ApplicationError that wasn't produced by wrapBatchVerifyError - // — extractVerifiedCountFromError must not pull garbage out of it. - other := temporal.NewApplicationError("plain", "Other") - require.Equal(t, int64(0), extractVerifiedCountFromError(other)) - }) + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.True(t, result.ReachedEnd) + require.Equal(t, int64(0), result.VerifiedCount) } -// TestSharded_PartialVerifiedCount_RecordedOnError: when a batch -// errors out after partial verify, the wrapped error's count must be -// folded into ReplicatedWorkflowCount so the status query reflects -// work actually done rather than zeroing out a partially-successful -// batch. -func TestSharded_PartialVerifiedCount_RecordedOnError(t *testing.T) { +// TestChild_ThrottledPromotion_ReceivesSignal: a child started with +// StartThrottled=true runs at half rate but can receive a promotion signal +// via the childDirectRunner relay and complete a terminal page. This verifies +// that (a) childDirectRunner correctly relays the shardedResumeFullSignalName +// signal to the production child, and (b) a throttled child completes +// normally once promoted — i.e., that the promotion signal does not break +// execution when no CAN hint is active. +// +// Note: the "pause at CAN hint and wait for promotion" path is not testable +// via the child-workflow approach because env.SetContinueAsNewSuggested only +// affects the root (childDirectRunner) env, not the grandchild +// (shardedForceReplicationWorker) env. The CAN state is per-env and the +// SDK provides no API to set it on an inner child env from outside. +func TestChild_ThrottledPromotion_ReceivesSignal(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() - env.RegisterWorkflow(ShardedForceReplicationWorkflow) - registerShardedScaffolding(env, 2) - env.RegisterActivityWithOptions(pageThrough(makeExecs(2, 5), 1000), activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) - const partialCount int64 = 6 - env.RegisterActivityWithOptions(func(_ context.Context, _ *shardedBatchReq) (replicateBatchResult, error) { - return replicateBatchResult{}, wrapBatchVerifyError( - temporal.NewNonRetryableApplicationError("simulated mid-verify failure", "Simulated", nil), - partialCount, - ) + execs := makeExecs(2, 5) // 10 execs, single terminal page + env.RegisterActivityWithOptions(func(_ context.Context, _ *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + // From inside the activity body, schedule a delayed callback (runs on + // the workflow scheduler goroutine) to send the promotion signal. + // childDirectRunner's relay coroutine will forward it to the production + // child. Since the page is terminal, the child will proceed to + // completion rather than waiting for CAN-hint-triggered promotion. + env.RegisterDelayedCallback(func() { + env.SignalWorkflow(shardedResumeFullSignalName, struct{}{}) + }, 0) + return &listWorkflowsResponse{ + Executions: execs, + NextPageToken: nil, // terminal — child reaches end naturally + }, nil + }, activity.RegisterOptions{Name: "ListWorkflows"}) + + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + return replicateBatchResult{ + CompletedShards: req.Executions.sortedShards(), + VerifiedCount: 5, + }, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) - env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ - Namespace: "test-ns", - TargetClusterName: "remote_cluster", - }) + params := makeChildParams(2) + params.StartThrottled = true // child begins at half rate + env.ExecuteWorkflow(childDirectRunner, params) require.True(t, env.IsWorkflowCompleted()) - require.Error(t, env.GetWorkflowError()) - - envValue, qErr := env.QueryWorkflow(forceReplicationStatusQueryType) - require.NoError(t, qErr) - var status ForceReplicationStatus - require.NoError(t, envValue.Get(&status)) - - require.GreaterOrEqual(t, status.ReplicatedWorkflowCount, partialCount, - "failed batch's partial count should still be reflected in ReplicatedWorkflowCount") -} + require.NoError(t, env.GetWorkflowError()) -// ---- internal helpers ---- - -// makeExecsForShard produces `count` runs for the named shard's -// ResumeShard.Execs payload. Each run gets a distinct businessID so -// the resulting map has `count` entries with one run each — the -// simplest shape for tests that don't care about BID-reuse. -func makeExecsForShard(shard int32, count int) map[string][]RunEntry { - out := map[string][]RunEntry{} - for i := range count { - bid := "wf-" + strconv.Itoa(int(shard)) + "-" + strconv.Itoa(i) - out[bid] = []RunEntry{{RunID: "run-" + strconv.Itoa(int(shard)*1000+i)}} - } - return out + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + // Terminal page → child reaches end, even when starting throttled. + require.True(t, result.ReachedEnd, "terminal-page child should reach end regardless of throttle state") + require.Equal(t, int64(5), result.VerifiedCount) } From d4863655d234022333857f2192ceffc4ec571e2a Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Thu, 18 Jun 2026 15:18:46 +0100 Subject: [PATCH 32/35] Various fixes. --- .../migration/force_replication_workflow.go | 16 +- .../worker/migration/sharded_handover_test.go | 49 +- .../migration/sharded_parent_workflow.go | 63 ++- service/worker/migration/sharded_types.go | 10 +- .../worker/migration/sharded_types_test.go | 319 ++++++++++- service/worker/migration/sharded_workflow.go | 532 +++++++++++------- .../worker/migration/sharded_workflow_test.go | 62 +- 7 files changed, 786 insertions(+), 265 deletions(-) diff --git a/service/worker/migration/force_replication_workflow.go b/service/worker/migration/force_replication_workflow.go index 9ed400398bd..6f0023b33bf 100644 --- a/service/worker/migration/force_replication_workflow.go +++ b/service/worker/migration/force_replication_workflow.go @@ -90,13 +90,17 @@ type ( TotalWorkflowCount int64 ReplicatedWorkflowCount int64 ReplicatedWorkflowCountPerSecond float64 - PageTokenForRestart []byte - // RecoveryNextPageToken is the page token the sharded workflow - // was processing when it last continued-as-new or was interrupted. - // Feed this back into ShardedForceReplicationWorkflow.NextPageToken - // to resume from that position. Left zero by the legacy variants. - RecoveryNextPageToken []byte + // PageTokenForRestart is the ListWorkflows page token to restart + // from if the operation is interrupted; feed it back into the + // workflow's NextPageToken to resume near current progress. The + // legacy variants return the current continue-as-new run's start + // token, which advances as the parent frequently CANs. The sharded + // parent rarely CANs, so it instead returns the latest child + // checkpoint token (falling back to the run's start token before + // any child has checkpointed) so the value tracks progress the + // same way. + PageTokenForRestart []byte } ) diff --git a/service/worker/migration/sharded_handover_test.go b/service/worker/migration/sharded_handover_test.go index e7a591014da..8e4a289106e 100644 --- a/service/worker/migration/sharded_handover_test.go +++ b/service/worker/migration/sharded_handover_test.go @@ -53,7 +53,7 @@ func parentTestEnv(t *testing.T, shardCount int32) *testsuite.TestWorkflowEnviro // TestParent_OneFullHandover verifies the core handover sequence: // 1. Child 0 starts (not throttled) and sends a checkpoint to the parent after 1 minute. -// 2. Parent starts child 1 with StartThrottled=true. +// 2. Parent starts child 1 with Handover=true. // 3. Child 0 completes after 5 minutes. // 4. Parent sends resumeFullRate to child 1. // 5. Child 1 receives the signal and completes, parent returns nil. @@ -138,12 +138,12 @@ func TestParent_OneFullHandover(t *testing.T) { // Two children should have been started. require.Len(t, childParamsSeen, 2, "parent should start exactly two children") - // Child 0 is not throttled (first child, no predecessor). - require.False(t, childParamsSeen[0].StartThrottled, "child 0 must not be throttled") + // Child 0 does not start in handover (first child, no predecessor). + require.False(t, childParamsSeen[0].Handover, "child 0 must not start in handover") - // Child 1 is throttled (started alongside still-running child 0). - require.True(t, childParamsSeen[1].StartThrottled, - "child 1 must start throttled (predecessor still running)") + // Child 1 starts in handover (started alongside still-running child 0). + require.True(t, childParamsSeen[1].Handover, + "child 1 must start in handover (predecessor still running)") // Promotion (resumeFullRate to the successor) and live-count rollups are not // asserted here: the SDK test env assigns a mocked child a GetInfo run ID that @@ -243,6 +243,37 @@ func TestParent_ChildFailure(t *testing.T) { "parent error message should contain 'child worker' prefix") } +// TestParent_NoReachedEndNoCheckpoint_Fails verifies the parent fails loudly +// rather than returning nil when its last child completes without reaching the +// namespace end and without having checkpointed. That state leaves work past +// the last checkpoint with nothing scheduled to process it; returning nil +// would silently drop those executions, so the parent must surface a +// ShardedParentStuck error instead. +func TestParent_NoReachedEndNoCheckpoint_Fails(t *testing.T) { + env := parentTestEnv(t, 2) + + // Child completes ReachedEnd=false and sends no checkpoint, so no successor + // is started and the namespace is never marked exhausted. The CAN hint is + // left unset (parentTestEnv default), so the parent is not draining either. + env.OnWorkflow(shardedForceReplicationWorker, mock.Anything, mock.Anything).Return( + shardedChildResult{VerifiedCount: 3, ReachedEnd: false}, nil, + ).Once() + + env.ExecuteWorkflow(ShardedForceReplicationWorkflow, ShardedForceReplicationParams{ + Namespace: "test-ns", + TargetClusterName: "remote_cluster", + // Done=true skips the task-queue-user-data child so the parent reaches + // the terminal decision directly. + TaskQueueUserDataReplicationStatus: TaskQueueUserDataReplicationStatus{Done: true}, + }) + + require.True(t, env.IsWorkflowCompleted(), "parent should complete") + err := env.GetWorkflowError() + require.Error(t, err, "parent must fail when work remains but nothing is scheduled") + require.True(t, hasAppErrType(err, "ShardedParentStuck"), + "expected ShardedParentStuck in error chain, got: %v", err) +} + // handoverTestChildErr is the synthetic child error for TestParent_ChildFailure. var handoverTestChildErr = &handoverSyntheticError{msg: "synthetic child failure"} @@ -320,6 +351,8 @@ func TestParent_QueryAggregation(t *testing.T) { require.NoError(t, val.Get(&status)) require.Equal(t, int64(25), status.ReplicatedWorkflowCount, "final query must reflect retiredTotal of both children (20+5)") + require.Equal(t, []byte("next-page"), status.PageTokenForRestart, + "restart token must track the latest child checkpoint, not the run's start token") } // ---- Child tests — shardedForceReplicationWorker as root ---- @@ -390,7 +423,7 @@ func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { t.Run("blocks without promotion", func(t *testing.T) { env := newEnv() params := makeChildParams(2) - params.StartThrottled = true // not promoted: must await resumeFullRate before cutting + params.Handover = true // not promoted: must await resumeFullRate before cutting env.ExecuteWorkflow(shardedForceReplicationWorker, params) // A correctly-pausing child never completes on its own — it stays blocked // awaiting promotion, so the env surfaces a timeout rather than a result. A @@ -407,7 +440,7 @@ func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { env.SignalWorkflow(shardedResumeFullSignalName, struct{}{}) }, time.Minute) params := makeChildParams(2) - params.StartThrottled = true + params.Handover = true env.ExecuteWorkflow(shardedForceReplicationWorker, params) require.True(t, env.IsWorkflowCompleted(), "child should complete once promoted") require.NoError(t, env.GetWorkflowError()) diff --git a/service/worker/migration/sharded_parent_workflow.go b/service/worker/migration/sharded_parent_workflow.go index f912ca08684..94c9f5f813c 100644 --- a/service/worker/migration/sharded_parent_workflow.go +++ b/service/worker/migration/sharded_parent_workflow.go @@ -29,8 +29,9 @@ import ( // The parent keeps the registered name "force-replication-sharded" so // tooling that already targets that name continues to work. func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceReplicationParams) error { - // startPageToken is the token at workflow entry — returned as - // PageTokenForRestart in the status query for tooling compatibility. + // startPageToken is the token at workflow entry — the PageTokenForRestart + // fallback used until the first child checkpoint advances it (see the + // query handler below). startPageToken := params.NextPageToken // ps is set after setup; the query handler closes over the pointer so @@ -47,11 +48,17 @@ func ShardedForceReplicationWorkflow(ctx workflow.Context, params ShardedForceRe ReplicatedWorkflowCountPerSecond: params.ReplicatedWorkflowCountPerSecond, PageTokenForRestart: startPageToken, TaskQueueUserDataReplicationStatus: params.TaskQueueUserDataReplicationStatus, - RecoveryNextPageToken: params.NextPageToken, } if ps != nil { status.ReplicatedWorkflowCount = ps.retiredTotal + ps.liveTotalCount() status.ReplicatedWorkflowCountPerSecond = params.ReplicatedWorkflowCountPerSecond + // The parent rarely CANs, so startPageToken would stay pinned at + // the run's initial token. Once children begin checkpointing, + // advance the restart token to the latest checkpoint so a restart + // resumes near current progress instead of replaying the whole run. + if ps.lastCheckpointToken != nil { + status.PageTokenForRestart = ps.lastCheckpointToken + } } return status, nil }); err != nil { @@ -257,6 +264,12 @@ type shardedParentState struct { // child completes the parent CANs with the last checkpoint token. drainingForCAN bool + // reachedEnd is set when a child returns ReachedEnd=true — it exhausted + // the namespace (ListWorkflows returned an empty next-page token). This + // is the authoritative terminal signal: once set the parent finishes + // rather than continuing-as-new, because no work remains for a successor. + reachedEnd bool + // lastCheckpointToken is the NextPageToken from the most recently // received checkpoint signal. Carried into the CAN params so the // new parent run resumes children from the right position. @@ -321,15 +334,17 @@ func (ps *shardedParentState) liveTotalCount() int64 { // // Handover orchestration: // - On checkpoint: start a successor (at half rate) unless drainingForCAN. -// - On child completion: retire its verifiedCount, send resumeFullRate -// to its successor (if one was started), check for terminal exit. +// - On child completion: retire its verifiedCount, record ReachedEnd, and +// send resumeFullRate to its successor (if one was started). // - On GetContinueAsNewSuggested: set drainingForCAN. When the last live // child completes, CAN with lastCheckpointToken as NextPageToken. +// - On exit (no live children): finish if a child reached the namespace +// end; else CAN if draining; else fail (handover bookkeeping bug). // // At most two children are live at once (the cutting child plus its // successor) and at most one handover is in flight at any time. func (ps *shardedParentState) run(ctx workflow.Context) error { - // Start the first child. Not throttled — no predecessor exists. + // Start the first child. Not in handover — no predecessor exists. if err := ps.startChild(ctx, ps.params.NextPageToken, false); err != nil { return err } @@ -388,6 +403,15 @@ func (ps *shardedParentState) run(ctx workflow.Context) error { sel.Select(ctx) } + if ps.reachedEnd { + // A child exhausted the namespace — terminal. Record the final count + // for the terminal Await in ShardedForceReplicationWorkflow. Takes + // precedence over drainingForCAN: once the end is reached there is + // nothing left for a CAN'd successor to process. + ps.params.ReplicatedWorkflowCount = ps.retiredTotal + return nil + } + if ps.drainingForCAN { // Parent CAN: carry the last checkpoint token forward so the new // parent run resumes children from the right position. @@ -398,11 +422,13 @@ func (ps *shardedParentState) run(ctx workflow.Context) error { return workflow.NewContinueAsNewError(ctx, ShardedForceReplicationWorkflow, next) } - // Terminal: the last child set ReachedEnd=true, meaning the namespace - // is exhausted. Record final count in params for the terminal Await - // in ShardedForceReplicationWorkflow. - ps.params.ReplicatedWorkflowCount = ps.retiredTotal - return nil + // No live children, yet the namespace was never exhausted and we are not + // draining for CAN — work remains past the last checkpoint with nothing + // scheduled to process it (a handover-bookkeeping bug). Fail loudly rather + // than returning nil, which would silently drop those executions. + return temporal.NewNonRetryableApplicationError( + "parent has no live children but the namespace is not exhausted and is not draining for CAN (handover bookkeeping bug)", + "ShardedParentStuck", nil) } // startChild starts a new shardedForceReplicationWorker child and records @@ -410,9 +436,9 @@ func (ps *shardedParentState) run(ctx workflow.Context) error { // execution starts (GetChildWorkflowExecution().Get) so its run ID is // available for future signal addressing. // -// throttled=true means the child starts at half rate (StartThrottled=true), +// handover=true means the child starts at half rate (Handover=true), // awaiting a resumeFullRate signal from the parent before going full. -func (ps *shardedParentState) startChild(ctx workflow.Context, pageToken []byte, throttled bool) error { +func (ps *shardedParentState) startChild(ctx workflow.Context, pageToken []byte, handover bool) error { childParams := shardedChildParams{ Namespace: ps.params.Namespace, Query: ps.params.Query, @@ -428,7 +454,7 @@ func (ps *shardedParentState) startChild(ctx workflow.Context, pageToken []byte, IdleShardCost: ps.params.IdleShardCost, PerBatchGenerateRPS: ps.params.PerBatchGenerateRPS, StartPageToken: pageToken, - StartThrottled: throttled, + Handover: handover, } childOpts := workflow.ChildWorkflowOptions{ @@ -486,7 +512,7 @@ func (ps *shardedParentState) handleCheckpoint(ctx workflow.Context, sel workflo } ps.successorStarted[payload.ChildRunID] = true - // Start the successor at half rate (StartThrottled=true). The + // Start the successor at half rate (Handover=true). The // parent will promote it to full rate via resumeFullRate when the // current child (payload.ChildRunID) completes. if err := ps.startChild(ctx, payload.NextPageToken, true); err != nil { @@ -533,6 +559,13 @@ func (ps *shardedParentState) onChildCompleted(ctx workflow.Context, runID strin ps.retiredTotal += result.VerifiedCount ps.updateQPS(ctx) + // A child that exhausted the namespace is the chain's terminal child: + // it started no checkpoint and needs no successor. Record it so run() + // finishes instead of continuing-as-new. + if result.ReachedEnd { + ps.reachedEnd = true + } + // Promote the successor to full rate if one was started. The // successor's run ID is the last entry in liveRunIDs that is still // live (i.e., the one started from this child's checkpoint). diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index deb16852cb9..980b4a8d990 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -284,11 +284,11 @@ type shardedChildParams struct { // this child begins listing. Nil for the very first child. StartPageToken []byte - // StartThrottled, when true, means the child begins at half - // MaxExecsPerShard until the parent signals resumeFullRate. Used - // during handover overlap so a new child and its not-yet-drained - // predecessor together stay ≤ MaxExecsPerShard per shard. - StartThrottled bool + // Handover, when true, means the child begins in the startup handover + // phase: running at half MaxExecsPerShard until the parent signals + // resumeFullRate. Used during handover overlap so a new child and its + // not-yet-drained predecessor together stay ≤ MaxExecsPerShard per shard. + Handover bool } // shardedChildResult is the return value of each child worker workflow. diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 91f1ea4e69f..3ee84060e3c 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -152,30 +152,327 @@ func TestNewShardVerifyTracker_SeedsAllShards(t *testing.T) { require.False(t, sv1.lastProgress.IsZero(), "lastProgress must be seeded") } -// TestEffectiveMaxExecsPerShard: effectiveMaxExecsPerShard returns the -// full MaxExecsPerShard when promoted and not cut, and max(cap/2, 1) -// in all other cases. +// TestEffectiveMaxExecsPerShard: effectiveMaxExecsPerShard returns the full +// MaxExecsPerShard outside a handover phase, and max(cap/2, 1) while in handover. func TestEffectiveMaxExecsPerShard(t *testing.T) { cases := []struct { name string - promoted, cut bool + handover bool maxExecsPerShard int want int }{ - {"promoted and not cut returns full", true, false, 10, 10}, - {"not promoted returns half", false, false, 10, 5}, - {"promoted and cut returns half", true, true, 10, 5}, - {"minimum of 1 enforced", false, false, 1, 1}, - {"odd cap rounds down but floors at 1", false, false, 3, 1}, + {"not in handover returns full", false, 10, 10}, + {"in handover returns half", true, 10, 5}, + {"minimum of 1 enforced", true, 1, 1}, + {"odd cap rounds down but floors at 1", true, 3, 1}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { s := &shardedWorkflowState{ params: &shardedChildParams{MaxExecsPerShard: tc.maxExecsPerShard}, - promoted: tc.promoted, - cut: tc.cut, + handover: tc.handover, } require.Equal(t, tc.want, s.effectiveMaxExecsPerShard()) }) } } + +// ---- shardBuckets ---- + +// requireBucketInvariant asserts shardBuckets' core invariant: counts[sh] +// equals the total runs under byShard[sh], and a shard key is present in both +// maps or in neither (no empty buckets, no orphan counts). +func requireBucketInvariant(t *testing.T, b *shardBuckets) { + t.Helper() + for sh, byBID := range b.byShard { + n := 0 + for _, runs := range byBID { + n += len(runs) + } + require.Greaterf(t, n, 0, "byShard[%d] present but empty", sh) + require.Equalf(t, n, b.counts[sh], "counts[%d] out of sync with byShard[%d]", sh, sh) + } + for sh, c := range b.counts { + require.Greaterf(t, c, 0, "counts[%d] present but non-positive", sh) + _, ok := b.byShard[sh] + require.Truef(t, ok, "counts[%d] present with no byShard entry", sh) + } +} + +// TestShardBuckets_ZeroValueUsable: the zero value answers reads without +// panicking, and add lazily initialises the maps. +func TestShardBuckets_ZeroValueUsable(t *testing.T) { + var b shardBuckets + require.True(t, b.empty()) + require.Equal(t, 0, b.count(3)) + require.Equal(t, 0, b.totalRemaining()) + require.Nil(t, b.take(3, 5)) + require.Empty(t, b.shards()) + + b.add(3, "bid", RunEntry{RunID: "r1"}) + require.False(t, b.empty()) + require.Equal(t, 1, b.count(3)) + requireBucketInvariant(t, &b) +} + +// TestShardBuckets_AddAndCount: add bumps the per-shard count across BIDs and +// shards; totalRemaining sums them and shards lists non-empty shards sorted. +func TestShardBuckets_AddAndCount(t *testing.T) { + var b shardBuckets + b.add(1, "a", RunEntry{RunID: "r1"}) + b.add(1, "a", RunEntry{RunID: "r2"}) + b.add(1, "b", RunEntry{RunID: "r3"}) + b.add(2, "c", RunEntry{RunID: "r4"}) + + require.Equal(t, 3, b.count(1)) + require.Equal(t, 1, b.count(2)) + require.Equal(t, 0, b.count(99)) + require.Equal(t, 4, b.totalRemaining()) + require.False(t, b.empty()) + require.Equal(t, []int32{1, 2}, b.shards()) + requireBucketInvariant(t, &b) +} + +// TestShardBuckets_TakePartial: take walks BIDs alphabetically and takes whole +// per-BID groups only as far as needed to reach n, leaving the remainder intact. +func TestShardBuckets_TakePartial(t *testing.T) { + var b shardBuckets + b.add(1, "a", RunEntry{RunID: "r1"}) + b.add(1, "a", RunEntry{RunID: "r2"}) + b.add(1, "a", RunEntry{RunID: "r3"}) + b.add(1, "b", RunEntry{RunID: "r4"}) + b.add(1, "b", RunEntry{RunID: "r5"}) + b.add(1, "c", RunEntry{RunID: "r6"}) + + // n=4: take all of "a" (3), then 1 of "b"; stop before "c". + got := b.take(1, 4) + require.Equal(t, map[string][]RunEntry{ + "a": {{RunID: "r1"}, {RunID: "r2"}, {RunID: "r3"}}, + "b": {{RunID: "r4"}}, + }, got) + require.Equal(t, 2, b.count(1)) // r5, r6 remain + requireBucketInvariant(t, &b) + + // Next take drains the rest, still in BID order. + got = b.take(1, 100) + require.Equal(t, map[string][]RunEntry{ + "b": {{RunID: "r5"}}, + "c": {{RunID: "r6"}}, + }, got) + require.True(t, b.empty()) +} + +// TestShardBuckets_TakeAllDropsShard: taking everything (or more) removes the +// shard from both maps so empty/shards/count agree. +func TestShardBuckets_TakeAllDropsShard(t *testing.T) { + var b shardBuckets + b.add(5, "a", RunEntry{RunID: "r1"}) + b.add(5, "b", RunEntry{RunID: "r2"}) + + got := b.take(5, 10) // more than present + require.Len(t, got, 2) + require.Equal(t, 0, b.count(5)) + require.True(t, b.empty()) + require.Empty(t, b.shards()) + + _, okPayload := b.byShard[5] + require.False(t, okPayload, "byShard must drop emptied shard") + _, okCount := b.counts[5] + require.False(t, okCount, "counts must drop emptied shard") +} + +// TestShardBuckets_TakeEdgeCases: non-positive n and empty/missing shards +// return nil without mutating state. +func TestShardBuckets_TakeEdgeCases(t *testing.T) { + var b shardBuckets + b.add(1, "a", RunEntry{RunID: "r1"}) + + require.Nil(t, b.take(1, 0)) + require.Nil(t, b.take(1, -3)) + require.Nil(t, b.take(99, 5)) // no such shard + require.Equal(t, 1, b.count(1), "edge-case takes must not mutate") + requireBucketInvariant(t, &b) +} + +// TestShardBuckets_TakeOutputIndependentBacking: the returned slice has its +// own backing array, so appending to it cannot corrupt the bucket's leftover +// runs. Guards the append([]RunEntry(nil), ...) copy in take. +func TestShardBuckets_TakeOutputIndependentBacking(t *testing.T) { + var b shardBuckets + b.add(1, "a", RunEntry{RunID: "r1"}) + b.add(1, "a", RunEntry{RunID: "r2"}) + b.add(1, "a", RunEntry{RunID: "r3"}) + + got := b.take(1, 1) // takes r1, leaves r2,r3 + require.Equal(t, []RunEntry{{RunID: "r1"}}, got["a"]) + + // Appending to the returned slice must not overwrite the leftover r2/r3. + got["a"] = append(got["a"], RunEntry{RunID: "x"}) + + rest := b.take(1, 100) + require.Equal(t, []RunEntry{{RunID: "r2"}, {RunID: "r3"}}, rest["a"]) +} + +// ---- inFlightBatches ---- + +// requireBatchInvariant asserts inFlightBatches' core invariant: inFlight is +// exactly the union of every batch's held set, and each shard is held by at +// most one batch. +func requireBatchInvariant(t *testing.T, b *inFlightBatches) { + t.Helper() + holders := map[int32]int{} + for _, held := range b.held { + for sh := range held { + holders[sh]++ + } + } + for sh, n := range holders { + require.Equalf(t, 1, n, "shard %d held by %d batches (must be exactly 1)", sh, n) + require.Truef(t, b.inFlight[sh], "shard %d held but missing from inFlight", sh) + } + require.Equalf(t, len(holders), len(b.inFlight), + "inFlight (%d entries) must match the union of held sets (%d shards)", len(b.inFlight), len(holders)) +} + +// payloadForShards builds a minimal BatchPayload claiming the given shards. +// claim only inspects the top-level shard keys, so the inner runs are filler. +func payloadForShards(shards ...int32) BatchPayload { + p := BatchPayload{} + for _, sh := range shards { + p[sh] = map[string][]RunEntry{"b": {{RunID: "r"}}} + } + return p +} + +// TestInFlightBatches_ZeroValueUsable: reads on the zero value are safe, and +// releasing unknown batches is a no-op; claim lazily initialises the maps. +func TestInFlightBatches_ZeroValueUsable(t *testing.T) { + var b inFlightBatches + require.False(t, b.isInFlight(5)) + b.releaseShards(123, []int32{5}) // unknown batch: no-op, no panic + b.releaseAll(123) // unknown batch: no-op, no panic + + id := b.claim(payloadForShards(1, 2)) + require.Equal(t, int64(1), id) + require.True(t, b.isInFlight(1)) + require.True(t, b.isInFlight(2)) + requireBatchInvariant(t, &b) +} + +// TestInFlightBatches_ClaimAssignsMonotonicIDs: each claim gets the next ID +// and marks exactly its shards in-flight. +func TestInFlightBatches_ClaimAssignsMonotonicIDs(t *testing.T) { + var b inFlightBatches + require.Equal(t, int64(1), b.claim(payloadForShards(1))) + require.Equal(t, int64(2), b.claim(payloadForShards(2))) + require.Equal(t, int64(3), b.claim(payloadForShards(3))) + require.True(t, b.isInFlight(1)) + require.True(t, b.isInFlight(2)) + require.True(t, b.isInFlight(3)) + require.False(t, b.isInFlight(4)) + requireBatchInvariant(t, &b) +} + +// TestInFlightBatches_Count: count tracks claimed-but-not-released batches — +// the number of live dispatch coroutines. A batch that has signal-released all +// its shards still counts until releaseAll, because its activity is still +// running. +func TestInFlightBatches_Count(t *testing.T) { + var b inFlightBatches + require.Equal(t, 0, b.count()) + + b1 := b.claim(payloadForShards(1, 2)) + require.Equal(t, 1, b.count()) + b2 := b.claim(payloadForShards(3)) + require.Equal(t, 2, b.count()) + + // Signal-releasing every shard of b1 does NOT drop the batch: the + // activity is still running, so it stays counted. + b.releaseShards(b1, []int32{1, 2}) + require.Equal(t, 2, b.count(), "fully signal-released batch still counts until releaseAll") + + b.releaseAll(b1) + require.Equal(t, 1, b.count()) + b.releaseAll(b2) + require.Equal(t, 0, b.count()) +} + +// TestInFlightBatches_ReleaseAll: releaseAll frees the batch's shards and +// drops the batch. +func TestInFlightBatches_ReleaseAll(t *testing.T) { + var b inFlightBatches + id := b.claim(payloadForShards(1, 2)) + b.releaseAll(id) + + require.False(t, b.isInFlight(1)) + require.False(t, b.isInFlight(2)) + _, ok := b.held[id] + require.False(t, ok, "released batch must be dropped from held") + requireBatchInvariant(t, &b) +} + +// TestInFlightBatches_ReleaseShardsPartial: releaseShards frees only the named +// shards; the rest stay claimed until releaseAll. +func TestInFlightBatches_ReleaseShardsPartial(t *testing.T) { + var b inFlightBatches + id := b.claim(payloadForShards(1, 2, 3)) + + b.releaseShards(id, []int32{2}) + require.True(t, b.isInFlight(1)) + require.False(t, b.isInFlight(2)) + require.True(t, b.isInFlight(3)) + requireBatchInvariant(t, &b) + + b.releaseAll(id) + require.False(t, b.isInFlight(1)) + require.False(t, b.isInFlight(3)) + requireBatchInvariant(t, &b) +} + +// TestInFlightBatches_ReleaseShardsSkipsUnheld: shards not held by the batch +// are ignored (never claimed, or already released — releasing twice is safe). +func TestInFlightBatches_ReleaseShardsSkipsUnheld(t *testing.T) { + var b inFlightBatches + id := b.claim(payloadForShards(1)) + + // 2 and 3 were never claimed by this batch — releasing them is a no-op + // and must not disturb shard 1. + b.releaseShards(id, []int32{2, 3}) + require.True(t, b.isInFlight(1)) + require.False(t, b.isInFlight(2)) + requireBatchInvariant(t, &b) + + // Releasing the same shard twice is harmless. + b.releaseShards(id, []int32{1}) + b.releaseShards(id, []int32{1}) + require.False(t, b.isInFlight(1)) + requireBatchInvariant(t, &b) +} + +// TestInFlightBatches_ReleaseAllDoesNotStompReclaimedShard pins the +// signal-vs-defer invariant: after a batch signal-releases a shard that a +// later batch then re-claims, the first batch's releaseAll must clear only its +// own remaining shards, never the new claimant's. +func TestInFlightBatches_ReleaseAllDoesNotStompReclaimedShard(t *testing.T) { + var b inFlightBatches + + b1 := b.claim(payloadForShards(1, 2)) // b1 holds {1,2} + b.releaseShards(b1, []int32{1}) // b1 signal-releases shard 1 + require.False(t, b.isInFlight(1)) + + b2 := b.claim(payloadForShards(1)) // b2 re-claims shard 1 + require.True(t, b.isInFlight(1)) + requireBatchInvariant(t, &b) + + // b1's activity returns and clears its remaining claim (shard 2 only). + b.releaseAll(b1) + require.True(t, b.isInFlight(1), "shard 1 must stay claimed by b2") + require.False(t, b.isInFlight(2), "b1's remaining shard 2 must be freed") + requireBatchInvariant(t, &b) + + // b2 finishes; everything drains. + b.releaseAll(b2) + require.False(t, b.isInFlight(1)) + require.Empty(t, b.inFlight) + require.Empty(t, b.held) +} diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 95cb5777b80..b5cb55743b5 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -23,30 +23,30 @@ import ( // 1. List pages at effectiveMaxExecsPerShard rate until either the // namespace is exhausted (ReachedEnd=true, no checkpoint) or // GetContinueAsNewSuggested fires at a page boundary (send checkpoint -// signal to parent, set cut flag, stop listing). +// signal to parent, enter the shutdown handover, stop listing). // 2. Drain remaining bucket execs via drainBuckets. -// 3. Await pendingDispatches == 0. +// 3. Await all in-flight batches complete (batches.count() == 0). // 4. Return shardedChildResult{VerifiedCount, ReachedEnd}. // // A 60 s timer coroutine sends cumulative verifiedCount rollup signals to // the parent so the force-replication-status query stays up to date. // -// Throttled-hits-hint: a child started with StartThrottled=true (not yet -// promoted) pauses at a CAN hint and awaits the parent's resumeFullRate -// signal before cutting. This ensures at most one handover is in flight at -// any time and at most two children run concurrently. +// Handover-hits-hint: a child started with Handover=true (not yet promoted) +// pauses at a CAN hint and awaits the parent's resumeFullRate signal before +// cutting. This ensures at most one handover is in flight at any time and at +// most two children run concurrently. func shardedForceReplicationWorker(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { s := &shardedWorkflowState{ params: ¶ms, namespaceID: params.NamespaceID, targetShardCount: params.TargetShardCount, - buckets: BatchPayload{}, - bucketCounts: map[int32]int{}, - shardInFlight: map[int32]bool{}, - heldByBatch: map[int64]map[int32]bool{}, - // StartThrottled=true → wait for parent promotion before going full. - // StartThrottled=false → first child, already at full rate. - promoted: !params.StartThrottled, + // buckets and batches default to usable zero values (their maps are + // lazily initialised on first mutation). + // + // Handover=true → started mid-handover (predecessor still running); + // runs at half rate until promoted. Handover=false → first child, no + // predecessor, full rate from the start. + handover: params.Handover, metricsHandler: workflow.GetMetricsHandler(ctx).WithTags(map[string]string{ metrics.OperationTagName: metrics.MigrationWorkflowScope, NamespaceTagName: params.Namespace, @@ -69,79 +69,58 @@ type shardedWorkflowState struct { // Drives the per-exec shard hash and ConcurrentBatchCount derivation. targetShardCount int32 - // buckets accumulate execs that have been listed but not yet - // dispatched. Nested by destination shard then businessID. - buckets BatchPayload - - // bucketCounts mirrors len of all runs across BIDs for each shard. - // Kept as a sidecar so the packer's per-shard ordering decisions are - // O(1); consulted many times per cycle. - bucketCounts map[int32]int - - // shardInFlight is the per-shard exclusivity set: a shard's entry is - // set when it's part of any in-flight batch and cleared when that - // batch returns (either fully or via mid-flight signal-release). The - // packer uses this to ensure each shard is in at most one in-flight - // batch at a time. - shardInFlight map[int32]bool - - // heldByBatch tracks per-batch shard ownership. spawnBatch populates - // it with the batch's claimed shards; the signal handler removes - // entries as shards are released mid-flight; the dispatch coroutine's - // defer clears whatever remains after the activity returns. Required - // because a signal-released shard may have been re-claimed by a - // subsequent batch — the returning original batch must only clear its - // own remaining claims, not stomp on the new claimant. - heldByBatch map[int64]map[int32]bool - - // pendingDispatches counts spawned dispatch coroutines that have not - // yet returned. The main coroutine awaits this dropping to zero before - // returning the child result. - pendingDispatches int + // buckets buffer execs that have been listed but not yet dispatched, + // grouped by destination shard then businessID, with a sidecar + // per-shard count for O(1) packer ordering. See shardBuckets. + buckets shardBuckets + + // batches tracks the shards claimed by in-flight ReplicateBatch + // activities and enforces per-shard exclusivity (each shard is in at + // most one in-flight batch at a time). See inFlightBatches. + batches inFlightBatches // lastErr latches the first activity error. Once set, further dispatch // stops (tryPackStreaming and drainBuckets bail out) and the child // returns the error. lastErr error - // promoted is set when the parent sends the resumeFullRate signal: - // the predecessor child has completed and this child may run at the - // full MaxExecsPerShard rate. A child started with StartThrottled=false - // is promoted from the start (no predecessor). - promoted bool - - // cut is set when this child has sent its checkpoint signal to the - // parent and stopped listing new pages. After cut, effectiveMaxExecsPerShard - // drops to half so the successor can run at half alongside us without - // the combined per-shard in-flight count exceeding MaxExecsPerShard. - cut bool + // handover is set while this child is in a handover phase and must run + // at half rate so the other side of the handover can run alongside it + // without the combined per-shard in-flight count exceeding + // MaxExecsPerShard. Two distinct phases set it: + // - Startup: a child started with Handover=true begins in handover + // (predecessor still running) and clears it when the parent sends the + // resumeFullRate signal (predecessor done). The first child + // (Handover=false) has no predecessor and starts clear. + // - Shutdown: a child re-enters handover when it cuts — sends its + // checkpoint signal and stops listing — so the successor the parent + // starts from that checkpoint can run at half alongside it. + handover bool // verifiedCount accumulates verified executions for the final child // result and the 60 s progress rollup signal to the parent. verifiedCount int64 - nextBatchID int64 - // metricsHandler is tagged with the workflow's fixed scope + namespace // once at construction; recordVerified reuses it on every batch return. metricsHandler sdkclient.MetricsHandler } -// effectiveMaxExecsPerShard returns MaxExecsPerShard when this child is -// promoted (predecessor done) and not yet cut (checkpoint not yet sent): -// i.e., the normal full-rate steady state. In all other cases it returns -// max(MaxExecsPerShard/2, 1). +// effectiveMaxExecsPerShard returns the full MaxExecsPerShard in the steady +// state (not in a handover phase) and max(MaxExecsPerShard/2, 1) while in +// handover. // -// Half-rate applies in two distinct situations: -// - Before promotion: this child was started alongside a still-running +// Half-rate applies throughout a handover phase, which spans two situations +// (see the handover field): +// - Startup handover: this child was started alongside a still-running // predecessor; together they must stay ≤ MaxExecsPerShard per shard. -// - After cut: a successor has been started at half rate alongside us; -// same combined-load constraint. +// - Shutdown handover (after cut): a successor has been started at half +// rate alongside us; same combined-load constraint. // // Minimum of 1 ensures at least one exec can always be packed regardless // of BatchSize or MaxExecsPerShard settings. func (s *shardedWorkflowState) effectiveMaxExecsPerShard() int { - if s.promoted && !s.cut { + if !s.handover { return s.params.MaxExecsPerShard } return max(s.params.MaxExecsPerShard/2, 1) @@ -162,7 +141,7 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) (shardedChildResult, er } // Drain remaining buckets (at effectiveMaxExecsPerShard, which is - // half after cut so the successor can run alongside). + // half during the shutdown handover so the successor can run alongside). s.drainBuckets(ctx) if s.lastErr != nil { return shardedChildResult{}, s.lastErr @@ -170,7 +149,7 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) (shardedChildResult, er // Await all in-flight batches. No cancellation — children always run // to natural completion so their verified counts are exact. - _ = workflow.Await(ctx, func() bool { return s.pendingDispatches == 0 }) + _ = workflow.Await(ctx, func() bool { return s.batches.count() == 0 }) if s.lastErr != nil { return shardedChildResult{}, s.lastErr } @@ -185,14 +164,16 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) (shardedChildResult, er // 60 s progress rollup. The rollup runs only when the worker has a parent to // signal (skipped when the worker runs standalone, e.g. in unit tests). func (s *shardedWorkflowState) startBackgroundCoroutines(ctx workflow.Context, parentExec *workflow.Execution) { - // Parent promotion handler: a one-shot signal that sets promoted=true, - // allowing the child to advance to full rate. Runs until the signal - // arrives (or ctx is cancelled at workflow end). + // Parent promotion handler: a one-shot signal that ends the startup + // handover (clears handover), advancing the child to full rate once the + // predecessor has completed. Runs until the signal arrives (or ctx is + // cancelled at workflow end). Only throttled children ever receive it, and + // always before they cut, so it never races the shutdown handover. workflow.Go(ctx, func(gCtx workflow.Context) { ch := workflow.GetSignalChannel(gCtx, shardedResumeFullSignalName) var dummy struct{} if ch.Receive(gCtx, &dummy) { - s.promoted = true + s.handover = false } }) @@ -234,7 +215,7 @@ func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExe for _, ex := range executions { sh := common.WorkflowIDToHistoryShard(s.namespaceID, ex.BusinessID, s.targetShardCount) - s.addToBucket(sh, ex.BusinessID, RunEntry{ + s.buckets.add(sh, ex.BusinessID, RunEntry{ RunID: ex.RunID, ArchetypeID: ex.ArchetypeID, }) @@ -255,26 +236,34 @@ func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExe s.cutAtBoundary(ctx, parentExec, currentPageToken) return false } + + // Backpressure: don't fetch the next page until a dispatch slot + // is free. With every slot full, tryPackStreaming can't dispatch, + // so listing further would only inflate the buckets (and replay + // history) with execs we can't act on. Wakes on a freed slot or + // lastErr; the loop's lastErr guard then re-checks. + s.waitForDispatchSlot(ctx) } return false } -// cutAtBoundary performs the handover cut. If still throttled it first awaits -// promotion (so at most one handover is ever in flight), then sets cut (dropping -// to half rate) and signals the parent the next page token; the parent starts a -// successor from there and promotes it once this child completes. +// cutAtBoundary performs the handover cut. If still in the startup handover it +// first awaits promotion (so at most one handover is ever in flight), then +// re-enters handover for the shutdown phase (dropping to half rate) and signals +// the parent the next page token; the parent starts a successor from there and +// promotes it once this child completes. func (s *shardedWorkflowState) cutAtBoundary(ctx workflow.Context, parentExec *workflow.Execution, pageToken []byte) { - if !s.promoted { - // Throttled-hits-hint: the predecessor is still running. Pausing here - // bounds our history at ~hint size and guarantees at most one handover - // is in flight at a time. Await promotion (predecessor completion), then - // cut immediately. - _ = workflow.Await(ctx, func() bool { return s.promoted || s.lastErr != nil }) + if s.handover { + // Handover-hits-hint: still in the startup handover, predecessor + // running. Pausing here bounds our history at ~hint size and guarantees + // at most one handover is in flight at a time. Await promotion + // (predecessor completion, which clears handover), then cut immediately. + _ = workflow.Await(ctx, func() bool { return !s.handover || s.lastErr != nil }) if s.lastErr != nil { return } } - s.cut = true + s.handover = true if parentExec == nil { return } @@ -347,8 +336,9 @@ func (s *shardedWorkflowState) replicateBatch(ctx workflow.Context, req *sharded // handleReleaseSignals runs as a long-lived workflow coroutine, // consuming ReleaseShards signals from in-flight activities. Each // signal lists shards the activity considers complete; the handler -// clears them from heldByBatch[BatchID] (so the dispatch coroutine's -// defer won't double-release) and shardInFlight (so the packer can +// hands them to batches.releaseShards, which clears them from the +// batch's held set (so the dispatch coroutine's defer won't +// double-release) and from the in-flight set (so the packer can // dispatch new work against them while the activity stays running on // its still-pending shards). // @@ -363,16 +353,7 @@ func (s *shardedWorkflowState) handleReleaseSignals(ctx workflow.Context) { if !ch.Receive(ctx, &payload) { return } - held, ok := s.heldByBatch[payload.BatchID] - if !ok { - continue - } - for _, sh := range payload.Shards { - if held[sh] { - delete(held, sh) - delete(s.shardInFlight, sh) - } - } + s.batches.releaseShards(payload.BatchID, payload.Shards) } } @@ -387,13 +368,13 @@ func (s *shardedWorkflowState) setLastErr(err error) { // dispatchSlotAvailable returns true when the workflow is below the // in-flight batch ceiling and is free to spawn another batch. func (s *shardedWorkflowState) dispatchSlotAvailable() bool { - return s.pendingDispatches < s.params.ConcurrentBatchCount + return s.batches.count() < s.params.ConcurrentBatchCount } // waitForDispatchSlot blocks until a dispatch slot frees up or lastErr trips. func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { _ = workflow.Await(ctx, func() bool { - return s.lastErr != nil || s.pendingDispatches < s.params.ConcurrentBatchCount + return s.lastErr != nil || s.batches.count() < s.params.ConcurrentBatchCount }) } @@ -411,13 +392,12 @@ func (s *shardedWorkflowState) recordVerified(verified int64) { // spawnBatch dispatches one batch on a new workflow.Go coroutine. // Callers must have already verified that every shard in payload is free -// (not in shardInFlight) — spawnBatch marks them in-flight here. +// (batches.isInFlight false); claim marks them in-flight here. func (s *shardedWorkflowState) spawnBatch(ctx workflow.Context, payload BatchPayload) { if payload.totalRuns() == 0 { return } - s.nextBatchID++ - batchID := s.nextBatchID + batchID := s.batches.claim(payload) req := &shardedBatchReq{ BatchID: batchID, @@ -431,27 +411,12 @@ func (s *shardedWorkflowState) spawnBatch(ctx workflow.Context, payload BatchPay IdleShardCost: s.params.IdleShardCost, } - held := make(map[int32]bool, len(payload)) - //workflowcheck:ignore (building a set of flags is order-independent) - for sh := range payload { - s.shardInFlight[sh] = true - held[sh] = true - } - s.heldByBatch[batchID] = held - - s.pendingDispatches++ workflow.Go(ctx, func(coroCtx workflow.Context) { - defer func() { - s.pendingDispatches-- - // Clear any shards we still hold. Signal-released shards have - // already been cleared from shardInFlight by handleReleaseSignals - // and may by now belong to a subsequent batch's claim. - //workflowcheck:ignore (deletes are commutative; order-independent) - for sh := range s.heldByBatch[batchID] { - delete(s.shardInFlight, sh) - } - delete(s.heldByBatch, batchID) - }() + // releaseAll drops the batch (so batches.count falls) and clears + // whatever shards it still holds. Shards signal-released mid-flight + // are already gone (and may now belong to a later batch), so it only + // touches our remaining claims. + defer s.batches.releaseAll(batchID) result, err := s.replicateBatch(coroCtx, req) if err != nil { s.setLastErr(err) @@ -464,7 +429,7 @@ func (s *shardedWorkflowState) spawnBatch(ctx workflow.Context, payload BatchPay // tryPackStreaming attempts to pack and dispatch one batch from s.buckets. // Returns true if a batch was dispatched. Uses effectiveMaxExecsPerShard // as the per-shard cap so the half-rate constraint propagates into every -// batch packed while the child is throttled or cut. +// batch packed while the child is in a handover phase. // // No per-shard or total-bucket threshold: as soon as any free shard has // any execs and a dispatch slot is open, a batch fires. Safety is enforced @@ -472,10 +437,7 @@ func (s *shardedWorkflowState) spawnBatch(ctx workflow.Context, payload BatchPay // contribution, ConcurrentBatchCount caps in-flight batches, and // PerBatchGenerateRPS caps the per-batch source RPS. func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool) bool { - if s.lastErr != nil || s.params.BatchSize <= 0 || s.params.MaxExecsPerShard <= 0 { - return false - } - if !s.dispatchSlotAvailable() { + if s.lastErr != nil || !s.dispatchSlotAvailable() { return false } shardIDs := s.shardIDsByPackPriority(relax) @@ -491,11 +453,11 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool if room <= 0 { break } - take := min(effectiveCap, s.bucketCounts[sh], room) + take := min(effectiveCap, s.buckets.count(sh), room) if take == 0 { continue } - payload[sh] = s.takeFromBucket(sh, take) + payload[sh] = s.buckets.take(sh, take) packed += take } if packed == 0 { @@ -505,22 +467,10 @@ func (s *shardedWorkflowState) tryPackStreaming(ctx workflow.Context, relax bool return true } -// bucketsEmpty reports whether every shard's bucket is empty. Reads from -// the sidecar count map so it's O(#shards), not O(#runs). -func (s *shardedWorkflowState) bucketsEmpty() bool { - //workflowcheck:ignore (existence check; iteration order does not affect result) - for _, n := range s.bucketCounts { - if n > 0 { - return false - } - } - return true -} - // shardIDsByPackPriority returns free, non-empty shard IDs in the order -// the packer should consider them. Deterministic across replays: ordering -// is derived from workflow state (bucketCounts) with shard ID as a stable -// tiebreaker. +// the packer should consider them. Deterministic across replays: the +// candidate set comes from buckets.shards() (sorted) and ordering is +// derived from per-shard counts with shard ID as a stable tiebreaker. // // relax=false (streaming): fullest first. Packer naturally produces large, // predictable batches when work is plentiful and small ones when it isn't @@ -534,29 +484,29 @@ func (s *shardedWorkflowState) bucketsEmpty() bool { // light shards clear out quickly — nothing is arriving in drain, so // waiting for cold buckets to grow is wasted wall-clock. func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { - out := make([]int32, 0, len(s.bucketCounts)) - //workflowcheck:ignore (output is sorted below before any observable use) - for sh, n := range s.bucketCounts { - if n == 0 || s.shardInFlight[sh] { - continue + candidates := s.buckets.shards() + out := make([]int32, 0, len(candidates)) + for _, sh := range candidates { + if !s.batches.isInFlight(sh) { + out = append(out, sh) } - out = append(out, sh) } if relax { maxPerShard := s.effectiveMaxExecsPerShard() slices.SortFunc(out, func(a, b int32) int { - aHot, bHot := s.bucketCounts[a] > maxPerShard, s.bucketCounts[b] > maxPerShard + ca, cb := s.buckets.count(a), s.buckets.count(b) + aHot, bHot := ca > maxPerShard, cb > maxPerShard switch { case aHot && !bHot: return -1 case !aHot && bHot: return 1 case aHot && bHot: - if d := s.bucketCounts[b] - s.bucketCounts[a]; d != 0 { + if d := cb - ca; d != 0 { return d } default: - if d := s.bucketCounts[a] - s.bucketCounts[b]; d != 0 { + if d := ca - cb; d != 0 { return d } } @@ -565,7 +515,7 @@ func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { return out } slices.SortFunc(out, func(a, b int32) int { - if d := s.bucketCounts[b] - s.bucketCounts[a]; d != 0 { + if d := s.buckets.count(b) - s.buckets.count(a); d != 0 { return d } return int(a - b) @@ -573,26 +523,107 @@ func (s *shardedWorkflowState) shardIDsByPackPriority(relax bool) []int32 { return out } -// addToBucket appends one run to the (shard, BID) bucket and bumps -// the sidecar count. -func (s *shardedWorkflowState) addToBucket(shard int32, businessID string, run RunEntry) { - if s.buckets[shard] == nil { - s.buckets[shard] = map[string][]RunEntry{} +// drainBuckets blocks until buckets are empty (success) or lastErr trips +// (failure). Each pass packs everything currently dispatchable, then +// awaits any change in the in-flight batch count or shard claims so the next +// pass can attempt shards just freed by signal-release. +func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { + for { + if s.lastErr != nil { + return + } + for s.tryPackStreaming(ctx, true) { //nolint:revive + } + if s.buckets.empty() || s.lastErr != nil { + return + } + currentPending := s.batches.count() + if currentPending == 0 { + s.failDrainBucketsStuck() + return + } + _ = workflow.Await(ctx, s.drainBucketsAwaitPredicate(currentPending)) + } +} + +// failDrainBucketsStuck sets lastErr when buckets are non-empty but no +// batches are in flight — the shard-claim bookkeeping is corrupted, and +// returning silently would proceed with execs that were never dispatched +// (silent data loss). Failing forces the child to return the error. +func (s *shardedWorkflowState) failDrainBucketsStuck() { + remaining := s.buckets.totalRemaining() + s.setLastErr(temporal.NewNonRetryableApplicationError( + fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), + "DrainBucketsStuck", nil)) +} + +// drainBucketsAwaitPredicate returns true when the drainBuckets loop +// should wake up: lastErr tripped, a dispatch slot just freed, or a new +// free shard is ready to pack. A "free shard" wake-up only counts when +// there's also a dispatch slot to use it, otherwise the outer loop would +// busy-spin on tryPackStreaming returning false against the in-flight cap. +func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { + return func() bool { + if s.lastErr != nil { + return true + } + if s.batches.count() < currentPending { + return true + } + if !s.dispatchSlotAvailable() { + return false + } + for _, sh := range s.buckets.shards() { + if !s.batches.isInFlight(sh) { + return true + } + } + return false } - s.buckets[shard][businessID] = append(s.buckets[shard][businessID], run) - s.bucketCounts[shard]++ } -// takeFromBucket consumes up to n runs from the given shard and returns -// them grouped by BID. Walks BIDs in alphabetical order so the resulting -// payload is deterministic across replays; takes whole per-BID runs only -// as needed to reach n. Empties the shard from s.buckets / s.bucketCounts -// when nothing remains. -func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]RunEntry { +// shardBuckets accumulates listed-but-not-yet-dispatched execs, grouped by +// destination shard then businessID, with a sidecar per-shard count so the +// packer's ordering decisions stay O(1). +// +// Invariant: counts[sh] always equals the total runs under byShard[sh], and a +// shard key is present in both maps or in neither (a shard is dropped from +// both the moment its last run is taken). All mutation goes through add/take, +// so the invariant can't be violated from outside the type. The zero value is +// usable — add lazily initialises the maps on first use. +type shardBuckets struct { + // byShard holds the execs, nested shard → businessID → runs. This is the + // BatchPayload wire shape, reused so take can hand a shard's per-BID map + // straight into a batch payload. + byShard BatchPayload + // counts mirrors len of all runs across BIDs for each shard. Consulted + // many times per pack cycle, so it's kept as an O(1) sidecar rather than + // recomputed by walking byShard. + counts map[int32]int +} + +// add appends one run to the (shard, BID) bucket and bumps the sidecar count. +func (b *shardBuckets) add(shard int32, businessID string, run RunEntry) { + if b.byShard == nil { + b.byShard = BatchPayload{} + b.counts = map[int32]int{} + } + if b.byShard[shard] == nil { + b.byShard[shard] = map[string][]RunEntry{} + } + b.byShard[shard][businessID] = append(b.byShard[shard][businessID], run) + b.counts[shard]++ +} + +// take consumes up to n runs from the given shard and returns them grouped by +// BID. Walks BIDs in alphabetical order so the result is deterministic across +// replays; takes whole per-BID runs only as needed to reach n. Drops the shard +// from both maps when nothing remains. +func (b *shardBuckets) take(shard int32, n int) map[string][]RunEntry { if n <= 0 { return nil } - byBID := s.buckets[shard] + byBID := b.byShard[shard] if len(byBID) == 0 { return nil } @@ -622,74 +653,139 @@ func (s *shardedWorkflowState) takeFromBucket(shard int32, n int) map[string][]R } taken += take } - s.bucketCounts[shard] -= taken - if s.bucketCounts[shard] <= 0 { - delete(s.bucketCounts, shard) - delete(s.buckets, shard) + b.counts[shard] -= taken + if b.counts[shard] <= 0 { + delete(b.counts, shard) + delete(b.byShard, shard) } return out } -// drainBuckets blocks until buckets are empty (success) or lastErr trips -// (failure). Each pass packs everything currently dispatchable, then -// awaits any change in pendingDispatches + shardInFlight so the next pass -// can attempt shards just freed by signal-release. -func (s *shardedWorkflowState) drainBuckets(ctx workflow.Context) { - for { - if s.lastErr != nil { - return - } - for s.tryPackStreaming(ctx, true) { //nolint:revive - } - if s.bucketsEmpty() || s.lastErr != nil { - return - } - currentPending := s.pendingDispatches - if currentPending == 0 { - s.failDrainBucketsStuck() - return +// count returns the number of runs buffered for a shard (0 if none). +func (b *shardBuckets) count(shard int32) int { + return b.counts[shard] +} + +// empty reports whether every shard's bucket is empty. Reads the sidecar +// count map so it's O(#shards), not O(#runs). +func (b *shardBuckets) empty() bool { + //workflowcheck:ignore (existence check; iteration order does not affect result) + for _, n := range b.counts { + if n > 0 { + return false } - _ = workflow.Await(ctx, s.drainBucketsAwaitPredicate(currentPending)) } + return true } -// failDrainBucketsStuck sets lastErr when buckets are non-empty but no -// batches are in flight — the shard-claim bookkeeping is corrupted, and -// returning silently would proceed with execs that were never dispatched -// (silent data loss). Failing forces the child to return the error. -func (s *shardedWorkflowState) failDrainBucketsStuck() { +// totalRemaining sums runs across all shards. +func (b *shardBuckets) totalRemaining() int { remaining := 0 //workflowcheck:ignore (summation is order-independent) - for _, n := range s.bucketCounts { + for _, n := range b.counts { remaining += n } - s.setLastErr(temporal.NewNonRetryableApplicationError( - fmt.Sprintf("drainBuckets: %d execs in buckets but no batches in flight (shard-claim bookkeeping corrupted)", remaining), - "DrainBucketsStuck", nil)) + return remaining } -// drainBucketsAwaitPredicate returns true when the drainBuckets loop -// should wake up: lastErr tripped, a dispatch slot just freed, or a new -// free shard is ready to pack. A "free shard" wake-up only counts when -// there's also a dispatch slot to use it, otherwise the outer loop would -// busy-spin on tryPackStreaming returning false against the in-flight cap. -func (s *shardedWorkflowState) drainBucketsAwaitPredicate(currentPending int) func() bool { - return func() bool { - if s.lastErr != nil { - return true - } - if s.pendingDispatches < currentPending { - return true - } - if !s.dispatchSlotAvailable() { - return false +// shards returns every shard ID with a non-empty bucket, in ascending order. +// Sorted so callers can iterate deterministically across replays. +func (b *shardBuckets) shards() []int32 { + out := make([]int32, 0, len(b.counts)) + //workflowcheck:ignore (output is sorted below before any observable use) + for sh, n := range b.counts { + if n > 0 { + out = append(out, sh) } - //workflowcheck:ignore (existence check; iteration order does not affect result) - for sh, n := range s.bucketCounts { - if n > 0 && !s.shardInFlight[sh] { - return true - } + } + slices.Sort(out) + return out +} + +// inFlightBatches tracks which shards are claimed by in-flight ReplicateBatch +// activities, enforcing per-shard exclusivity: a shard is claimed by at most +// one batch at a time. +// +// Invariant: inFlight is the union of every batch's held set, so a shard is in +// inFlight iff exactly one batch holds it. All mutation goes through +// claim/releaseShards/releaseAll. The zero value is usable — claim lazily +// initialises the maps on first use. +// +// The per-batch held sets (rather than a single shared in-flight set) exist +// for the signal-vs-defer interplay: an activity may signal-release some of +// its shards mid-flight (releaseShards), after which a later batch can +// re-claim them; when the original activity finally returns, releaseAll must +// clear only the shards it still holds, not stomp the new claimant. Because +// releaseShards has already pruned held[batchID], releaseAll touching just +// held[batchID] makes that automatic. +type inFlightBatches struct { + // inFlight is the per-shard exclusivity set: the union of all held sets. + inFlight map[int32]bool + // held maps batchID → the shards that batch still owns. + held map[int64]map[int32]bool + // nextID hands out monotonic batch IDs. + nextID int64 +} + +// claim marks every shard in payload as in-flight under a fresh batch ID and +// returns that ID. Callers must have already confirmed each shard is free +// (isInFlight false) via the packer. +func (b *inFlightBatches) claim(payload BatchPayload) int64 { + if b.inFlight == nil { + b.inFlight = map[int32]bool{} + b.held = map[int64]map[int32]bool{} + } + b.nextID++ + batchID := b.nextID + held := make(map[int32]bool, len(payload)) + //workflowcheck:ignore (building a set of flags is order-independent) + for sh := range payload { + b.inFlight[sh] = true + held[sh] = true + } + b.held[batchID] = held + return batchID +} + +// isInFlight reports whether a shard is currently claimed by any batch. +func (b *inFlightBatches) isInFlight(shard int32) bool { + return b.inFlight[shard] +} + +// count returns the number of batches currently in flight: claimed but not +// yet releaseAll'd. This equals the number of live dispatch coroutines, +// because claim/releaseAll bracket each coroutine's lifetime. A batch that +// has signal-released all its shards still counts until its coroutine returns +// — releaseShards empties the batch's held set but leaves the batch entry, so +// the still-running activity remains counted. +func (b *inFlightBatches) count() int { + return len(b.held) +} + +// releaseShards clears the named shards from the batch's held set and from the +// in-flight set, so the packer can dispatch new work against them while the +// activity keeps running on its still-pending shards. Shards not held by this +// batch (already released, or re-claimed by another) are skipped. +func (b *inFlightBatches) releaseShards(batchID int64, shards []int32) { + held, ok := b.held[batchID] + if !ok { + return + } + for _, sh := range shards { + if held[sh] { + delete(held, sh) + delete(b.inFlight, sh) } - return false } } + +// releaseAll clears whatever shards the batch still holds and drops the batch. +// Signal-released shards are already gone from held (and may now belong to a +// later batch), so this only touches shards still owned here. +func (b *inFlightBatches) releaseAll(batchID int64) { + //workflowcheck:ignore (deletes are commutative; order-independent) + for sh := range b.held[batchID] { + delete(b.inFlight, sh) + } + delete(b.held, batchID) +} diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index cba3931294b..1cd4fb0a0d7 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -550,7 +550,7 @@ func TestChild_DisableVerification_ReachesEnd(t *testing.T) { } // TestChild_ThrottledPromotion_ReceivesSignal: a child started with -// StartThrottled=true runs at half rate but can receive a promotion signal +// Handover=true runs at half rate but can receive a promotion signal // via the childDirectRunner relay and complete a terminal page. This verifies // that (a) childDirectRunner correctly relays the shardedResumeFullSignalName // signal to the production child, and (b) a throttled child completes @@ -592,7 +592,7 @@ func TestChild_ThrottledPromotion_ReceivesSignal(t *testing.T) { }, activity.RegisterOptions{Name: "ReplicateBatch"}) params := makeChildParams(2) - params.StartThrottled = true // child begins at half rate + params.Handover = true // child begins at half rate env.ExecuteWorkflow(childDirectRunner, params) require.True(t, env.IsWorkflowCompleted()) @@ -604,3 +604,61 @@ func TestChild_ThrottledPromotion_ReceivesSignal(t *testing.T) { require.True(t, result.ReachedEnd, "terminal-page child should reach end regardless of throttle state") require.Equal(t, int64(5), result.VerifiedCount) } + +// TestChild_ListingBackpressure_BlocksUntilSlotFree: with ConcurrentBatchCount=1 +// the listing loop must not fetch the next page until the in-flight batch from +// the previous page has completed and freed the single dispatch slot. Without +// backpressure (waitForDispatchSlot), the loop would list every page up front, +// inflating the buckets with execs it can't dispatch. +// +// Four pages, each one exec on a distinct shard so per-shard exclusivity never +// binds and ConcurrentBatchCount is the only limiter. The events slice records +// each ListWorkflows ("L") and ReplicateBatch ("R") call; backpressure forces +// strict L,R,L,R,... alternation. waitForDispatchSlot Awaits the batch +// coroutine's releaseAll, which drops batches.count() (and runs only after +// ReplicateBatch returns) before the next list, so the recorded order is +// deterministic. +func TestChild_ListingBackpressure_BlocksUntilSlotFree(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(childDirectRunner) + env.RegisterWorkflow(shardedForceReplicationWorker) + + all := makeExecs(4, 1) // 4 execs across 4 distinct shards + pager := pageThrough(all, 1) + + var mu sync.Mutex + var events []string + env.RegisterActivityWithOptions(func(ctx context.Context, req *workflowservice.ListWorkflowExecutionsRequest) (*listWorkflowsResponse, error) { + mu.Lock() + events = append(events, "L") + mu.Unlock() + return pager(ctx, req) + }, activity.RegisterOptions{Name: "ListWorkflows"}) + env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { + mu.Lock() + events = append(events, "R") + mu.Unlock() + return replicateBatchResult{ + VerifiedCount: 1, + CompletedShards: req.Executions.sortedShards(), + }, nil + }, activity.RegisterOptions{Name: "ReplicateBatch"}) + + params := makeChildParams(4) + params.ConcurrentBatchCount = 1 + env.ExecuteWorkflow(childDirectRunner, params) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError()) + + var result shardedChildResult + require.NoError(t, env.GetWorkflowResult(&result)) + require.True(t, result.ReachedEnd) + require.Equal(t, int64(4), result.VerifiedCount) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, []string{"L", "R", "L", "R", "L", "R", "L", "R"}, events, + "listing must block on the single dispatch slot, alternating list and replicate") +} From c4b4c646272e0f3383d85a7cafbc198736821d91 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Thu, 18 Jun 2026 17:17:14 +0100 Subject: [PATCH 33/35] Better error handling. --- .../worker/migration/sharded_handover_test.go | 178 ++++++++++++++---- .../migration/sharded_parent_workflow.go | 49 +++-- service/worker/migration/sharded_types.go | 2 +- service/worker/migration/sharded_workflow.go | 66 ++++--- 4 files changed, 224 insertions(+), 71 deletions(-) diff --git a/service/worker/migration/sharded_handover_test.go b/service/worker/migration/sharded_handover_test.go index 8e4a289106e..b6992ba384b 100644 --- a/service/worker/migration/sharded_handover_test.go +++ b/service/worker/migration/sharded_handover_test.go @@ -1,7 +1,7 @@ package migration -// sharded_handover_test.go: focused unit tests for the handover/cut logic of the -// sharded force-replication child orchestration. +// sharded_handover_test.go: focused unit tests for the handover/checkpoint logic +// of the sharded force-replication child orchestration. // // Design rationale for the two test groups: // @@ -20,11 +20,12 @@ package migration // targets the production child's env directly so the CAN-hint code path executes. The // trade-off is that workflow.GetInfo(ctx).ParentWorkflowExecution is nil for root // workflows, which triggers a nil pointer dereference in the current production code -// (sharded_workflow.go). See TestChild_ThrottledHitsHint_SkipCase and the note in -// TestChild_CutAtHint_HalfRateAfterCut for details. +// (sharded_workflow.go). See TestChild_ThrottledHitsHint_PausesUntilPromoted and the +// note in TestChild_CheckpointAtHint_StopsListing for details. import ( "context" + "errors" "sync" "sync/atomic" "testing" @@ -33,6 +34,7 @@ import ( "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" "go.temporal.io/sdk/activity" + "go.temporal.io/sdk/temporal" "go.temporal.io/sdk/testsuite" "go.temporal.io/sdk/workflow" "go.temporal.io/server/common/payloads" @@ -85,7 +87,7 @@ func TestParent_OneFullHandover(t *testing.T) { if n == 1 { // Child 0: send a checkpoint to the parent after 1 minute (simulating - // "cut at page boundary"). Then complete after 5 minutes. + // "checkpoint at page boundary"). Then complete after 5 minutes. workflow.Go(ctx, func(gCtx workflow.Context) { _ = workflow.NewTimer(gCtx, 1*time.Minute).Get(gCtx, nil) if parentExec != nil { @@ -94,7 +96,7 @@ func TestParent_OneFullHandover(t *testing.T) { shardedCheckpointSignalName, shardedCheckpointPayload{ ChildRunID: myRunID, - NextPageToken: []byte("page-token-after-cut"), + NextPageToken: []byte("page-token-after-checkpoint"), }, ).Get(gCtx, nil) } @@ -290,7 +292,7 @@ func (e *handoverSyntheticError) Error() string { return e.msg } // ID, so handleProgress drops the rollup (the IDs match in real Temporal). End-to-end // rollup is covered by integration tests. // -// Sequence: child 0 cuts at 30s (→ child 1 starts) and completes at 1min with +// Sequence: child 0 checkpoints at 30s (→ child 1 starts) and completes at 1min with // VerifiedCount=20; child 1 completes with VerifiedCount=5. After both retire, the // final status query reports 25. func TestParent_QueryAggregation(t *testing.T) { @@ -306,7 +308,7 @@ func TestParent_QueryAggregation(t *testing.T) { myRunID := workflow.GetInfo(ctx).WorkflowExecution.RunID if n == 1 { - // Child 0: cut at 30s (parent starts the successor), complete at 1min. + // Child 0: checkpoint at 30s (parent starts the successor), complete at 1min. workflow.Go(ctx, func(gCtx workflow.Context) { _ = workflow.NewTimer(gCtx, 30*time.Second).Get(gCtx, nil) if parentExec != nil { @@ -357,21 +359,22 @@ func TestParent_QueryAggregation(t *testing.T) { // ---- Child tests — shardedForceReplicationWorker as root ---- -// TestChild_CutAtHint_StopsListing verifies that when GetContinueAsNewSuggested -// fires at a fully-consumed page boundary with pages remaining, a promoted child -// cuts: it stops listing, drains only the pages it already consumed, and returns -// ReachedEnd=false (the remaining pages belong to the successor the parent starts -// from the checkpoint token). The worker runs as the root workflow so -// env.SetContinueAsNewSuggested drives the production hint directly (the nil -// ParentWorkflowExecution is handled by the guard in run()). The per-shard -// half-rate value applied after the cut is unit-tested by TestEffectiveMaxExecsPerShard. -func TestChild_CutAtHint_StopsListing(t *testing.T) { +// TestChild_CheckpointAtHint_StopsListing verifies that when +// GetContinueAsNewSuggested fires at a fully-consumed page boundary with pages +// remaining, a promoted child checkpoints: it stops listing, drains only the +// pages it already consumed, and returns ReachedEnd=false (the remaining pages +// belong to the successor the parent starts from the checkpoint token). The +// worker runs as the root workflow so env.SetContinueAsNewSuggested drives the +// production hint directly (the nil ParentWorkflowExecution is handled by the +// guard in run()). The per-shard half-rate value applied after the checkpoint is +// unit-tested by TestEffectiveMaxExecsPerShard. +func TestChild_CheckpointAtHint_StopsListing(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() env.RegisterWorkflow(shardedForceReplicationWorker) // 20 execs paginated 5-at-a-time so page 1 returns a non-empty next-page token: - // the hint becomes a "cut" (work remains) rather than a terminal end. + // the hint becomes a "checkpoint" (work remains) rather than a terminal end. execs := makeExecs(2, 10) env.RegisterActivityWithOptions(pageThrough(execs, 5), activity.RegisterOptions{Name: "ListWorkflows"}) env.RegisterActivityWithOptions(func(_ context.Context, req *shardedBatchReq) (replicateBatchResult, error) { @@ -381,7 +384,7 @@ func TestChild_CutAtHint_StopsListing(t *testing.T) { }, nil }, activity.RegisterOptions{Name: "ReplicateBatch"}) - // Hint is true from the start: the child cuts after the first fully-consumed page. + // Hint is true from the start: the child checkpoints after the first fully-consumed page. env.SetContinueAsNewSuggested(true) params := makeChildParams(2) // not throttled → promoted; first child has no predecessor @@ -393,16 +396,16 @@ func TestChild_CutAtHint_StopsListing(t *testing.T) { var result shardedChildResult require.NoError(t, env.GetWorkflowResult(&result)) require.False(t, result.ReachedEnd, - "child must cut (ReachedEnd=false) when the hint fires with pages remaining") + "child must checkpoint (ReachedEnd=false) when the hint fires with pages remaining") require.Equal(t, int64(5), result.VerifiedCount, - "child verifies only the first fully-consumed page (5 execs) before cutting") + "child verifies only the first fully-consumed page (5 execs) before checkpointing") } // TestChild_ThrottledHitsHint_PausesUntilPromoted verifies the throttled-hits-hint // rule: a child started throttled (predecessor still running) that reaches the CAN -// hint must NOT cut until it is promoted — guaranteeing at most one handover in -// flight. The worker runs as root so env.SetContinueAsNewSuggested drives the hint -// and env.SignalWorkflow delivers resumeFullRate. +// hint must NOT checkpoint until it is promoted — guaranteeing at most one handover +// in flight. The worker runs as root so env.SetContinueAsNewSuggested drives the +// hint and env.SignalWorkflow delivers resumeFullRate. func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { newEnv := func() *testsuite.TestWorkflowEnvironment { suite := &testsuite.WorkflowTestSuite{} @@ -423,19 +426,19 @@ func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { t.Run("blocks without promotion", func(t *testing.T) { env := newEnv() params := makeChildParams(2) - params.Handover = true // not promoted: must await resumeFullRate before cutting + params.Handover = true // not promoted: must await resumeFullRate before checkpointing env.ExecuteWorkflow(shardedForceReplicationWorker, params) // A correctly-pausing child never completes on its own — it stays blocked // awaiting promotion, so the env surfaces a timeout rather than a result. A - // child that wrongly cut without promotion would instead complete cleanly - // with ReachedEnd=false, leaving GetWorkflowError nil. + // child that wrongly checkpointed without promotion would instead complete + // cleanly with ReachedEnd=false, leaving GetWorkflowError nil. require.Error(t, env.GetWorkflowError(), "throttled child must pause at the hint until promoted (never completing on its own)") }) t.Run("completes after promotion", func(t *testing.T) { env := newEnv() - // Deliver the promotion; the child leaves the pause, cuts, and completes. + // Deliver the promotion; the child leaves the pause, checkpoints, and completes. env.RegisterDelayedCallback(func() { env.SignalWorkflow(shardedResumeFullSignalName, struct{}{}) }, time.Minute) @@ -446,16 +449,16 @@ func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { require.NoError(t, env.GetWorkflowError()) var result shardedChildResult require.NoError(t, env.GetWorkflowResult(&result)) - require.False(t, result.ReachedEnd, "throttled child cuts after promotion (pages remain)") + require.False(t, result.ReachedEnd, "throttled child checkpoints after promotion (pages remain)") }) } -// TestChild_PromotedAndCuts_VerifiedCountAccumulated verifies the basic child -// lifecycle: a promoted child (not throttled) lists a namespace, dispatches +// TestChild_PromotedAndCheckpoints_VerifiedCountAccumulated verifies the basic +// child lifecycle: a promoted child (not throttled) lists a namespace, dispatches // batches, and returns ReachedEnd=true with the correct VerifiedCount. // This uses childDirectRunner (real child with parent) so ParentWorkflowExecution // is non-nil, sidestepping the production nil-guard bug. -func TestChild_PromotedAndCuts_VerifiedCountAccumulated(t *testing.T) { +func TestChild_PromotedAndCheckpoints_VerifiedCountAccumulated(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() env.RegisterWorkflow(childDirectRunner) @@ -505,3 +508,114 @@ func TestChild_PromotedAndCuts_VerifiedCountAccumulated(t *testing.T) { } mu.Unlock() } + +// checkpointProbe is a test-only workflow that drives checkpointAtBoundary +// directly with a non-nil parent execution. The production trigger for this path +// (a CAN hint at a page boundary) is not reproducible for a non-root child in +// the SDK test env — env.SetContinueAsNewSuggested only affects the root env, +// not an inner child (see the note on TestChild_ThrottledPromotion_ReceivesSignal) +// — so the probe exercises the checkpoint-signal handling without a full listing +// loop. It returns lastErr so the test can assert on it. +func checkpointProbe(ctx workflow.Context) error { + s := &shardedWorkflowState{} + s.checkpointAtBoundary(ctx, &workflow.Execution{ID: "parent-id", RunID: "parent-run"}, []byte("next-page")) + return s.lastErr +} + +// TestCheckpointAtBoundary_SignalFailureFailsChild verifies that checkpointAtBoundary +// latches lastErr when the checkpoint signal to the parent fails. A dropped +// checkpoint means the parent never starts a successor, so every execution past +// the page boundary would be silently lost; SignalExternalWorkflow has no +// transient failure modes, so the child must surface the error. The success +// case guards against the probe passing vacuously. +func TestCheckpointAtBoundary_SignalFailureFailsChild(t *testing.T) { + t.Run("signal fails → error latched", func(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(checkpointProbe) + + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedCheckpointSignalName, mock.Anything, + ).Return(errors.New("signal target gone")) + + env.ExecuteWorkflow(checkpointProbe) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err, "checkpointAtBoundary must latch lastErr when the signal fails") + require.Contains(t, err.Error(), "checkpoint signal to parent") + }) + + t.Run("signal succeeds → no error", func(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(checkpointProbe) + + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedCheckpointSignalName, mock.Anything, + ).Return(nil) + + env.ExecuteWorkflow(checkpointProbe) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError(), + "checkpointAtBoundary must not error when the signal succeeds") + }) +} + +// promoteSuccessorProbe is a test-only workflow that drives promoteSuccessor +// directly. The production promotion path is not reachable through the full +// parent flow in the SDK test env: the mocked child's GetInfo run ID differs +// from the run ID the parent tracks, so successorStarted never matches and the +// promotion block is skipped (see the note in TestParent_OneFullHandover). The +// probe returns childErr so the test can assert on it. +func promoteSuccessorProbe(ctx workflow.Context) error { + ps := &shardedParentState{} + ps.promoteSuccessor(ctx, workflow.Execution{ID: "successor-id", RunID: "successor-run"}) + return ps.childErr +} + +// TestPromoteSuccessor_SignalFailure verifies that an anomalous resumeFullRate +// failure fails the parent (a live, unpromoted successor would otherwise +// deadlock at its next checkpoint), while a "successor already completed" +// failure is tolerated (the successor finished without ever needing promotion). +func TestPromoteSuccessor_SignalFailure(t *testing.T) { + t.Run("anomalous failure → parent fails", func(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(promoteSuccessorProbe) + + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedResumeFullSignalName, mock.Anything, + ).Return(errors.New("unexpected signal failure")) + + env.ExecuteWorkflow(promoteSuccessorProbe) + + require.True(t, env.IsWorkflowCompleted()) + err := env.GetWorkflowError() + require.Error(t, err, "anomalous promotion-signal failure must fail the parent") + require.Contains(t, err.Error(), "resumeFullRate signal to successor") + }) + + t.Run("successor already completed → tolerated", func(t *testing.T) { + suite := &testsuite.WorkflowTestSuite{} + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(promoteSuccessorProbe) + + // UnknownExternalWorkflowExecutionError == target not found: the + // successor already completed, so it never needed promotion. + env.OnSignalExternalWorkflow( + mock.Anything, mock.Anything, mock.Anything, + shardedResumeFullSignalName, mock.Anything, + ).Return(&temporal.UnknownExternalWorkflowExecutionError{}) + + env.ExecuteWorkflow(promoteSuccessorProbe) + + require.True(t, env.IsWorkflowCompleted()) + require.NoError(t, env.GetWorkflowError(), + "a successor that already completed is benign — no promotion was needed") + }) +} diff --git a/service/worker/migration/sharded_parent_workflow.go b/service/worker/migration/sharded_parent_workflow.go index 94c9f5f813c..3c8f771e7af 100644 --- a/service/worker/migration/sharded_parent_workflow.go +++ b/service/worker/migration/sharded_parent_workflow.go @@ -1,6 +1,7 @@ package migration import ( + "errors" "fmt" "time" @@ -284,8 +285,9 @@ type shardedParentState struct { // (which cannot return an error) can surface it to the main loop. startErr error - // childErr captures the first child failure so the main loop can - // return it after the selector fires. + // childErr captures the first error surfaced from a child-completion + // callback — a child failure, or a failed promotion signal to a successor + // — so the main loop can return it after the selector fires. childErr error metricsHandler sdkclient.MetricsHandler @@ -341,7 +343,7 @@ func (ps *shardedParentState) liveTotalCount() int64 { // - On exit (no live children): finish if a child reached the namespace // end; else CAN if draining; else fail (handover bookkeeping bug). // -// At most two children are live at once (the cutting child plus its +// At most two children are live at once (the checkpointing child plus its // successor) and at most one handover is in flight at any time. func (ps *shardedParentState) run(ctx workflow.Context) error { // Start the first child. Not in handover — no predecessor exists. @@ -357,7 +359,7 @@ func (ps *shardedParentState) run(ctx workflow.Context) error { } ps.wiredCount = len(ps.liveRunIDs) - // Checkpoint signal: received when a child cuts at a page boundary. + // Checkpoint signal: received when a child checkpoints at a page boundary. // Re-added to the selector each iteration because AddReceive fires // once per message delivery; the channel itself persists. sel.AddReceive(ps.checkpointCh, func(c workflow.ReceiveChannel, _ bool) { @@ -585,22 +587,39 @@ func (ps *shardedParentState) onChildCompleted(ctx workflow.Context, runID strin if _, live := ps.liveChildren[candidateRunID]; !live { continue } - // Send resumeFullRate to the successor. Best-effort: - // if the signal fails (e.g., successor already completed), - // the child simply stays at half rate for its remaining work, - // which is safe. - successorExec := ps.liveExecs[candidateRunID] - _ = workflow.SignalExternalWorkflow(ctx, - successorExec.ID, - successorExec.RunID, - shardedResumeFullSignalName, - struct{}{}, - ).Get(ctx, nil) + ps.promoteSuccessor(ctx, ps.liveExecs[candidateRunID]) break } } } +// promoteSuccessor sends resumeFullRate to advance a throttled successor to +// full rate. This is load-bearing, not best-effort: a successor started +// throttled stays throttled until this signal arrives, and a throttled child +// blocks at its next checkpoint awaiting promotion (see checkpointAtBoundary), +// so a lost promotion to a still-running successor would deadlock the handover +// chain. +// +// The one benign failure is the successor having already completed +// (UnknownExternalWorkflowExecutionError, target not found): a throttled child +// can only complete by reaching the namespace end without hitting a checkpoint, +// so it never needed promotion. Any other failure means a live, unpromoted +// successor — and SignalExternalWorkflow has no transient failure modes — so it +// is latched into childErr to fail the parent rather than hang. +func (ps *shardedParentState) promoteSuccessor(ctx workflow.Context, successorExec workflow.Execution) { + if err := workflow.SignalExternalWorkflow(ctx, + successorExec.ID, + successorExec.RunID, + shardedResumeFullSignalName, + struct{}{}, + ).Get(ctx, nil); err != nil { + var gone *temporal.UnknownExternalWorkflowExecutionError + if !errors.As(err, &gone) && ps.childErr == nil { + ps.childErr = fmt.Errorf("resumeFullRate signal to successor %s: %w", successorExec.RunID, err) + } + } +} + // updateQPS feeds the current aggregate verified count into the QPSQueue // and refreshes the rate gauge. Called on each progress rollup and child // completion so the rate tracks actual throughput. diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 980b4a8d990..e1792280127 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -28,7 +28,7 @@ const ( releaseShardsSignalName = "ReleaseShards" // shardedCheckpointSignalName is sent by a child to the parent at - // its cut point — when GetContinueAsNewSuggested fires at a page + // its checkpoint — when GetContinueAsNewSuggested fires at a page // boundary. Payload: shardedCheckpointPayload. shardedCheckpointSignalName = "force-replication-sharded-checkpoint" diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index b5cb55743b5..ca60c897a77 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -33,8 +33,8 @@ import ( // // Handover-hits-hint: a child started with Handover=true (not yet promoted) // pauses at a CAN hint and awaits the parent's resumeFullRate signal before -// cutting. This ensures at most one handover is in flight at any time and at -// most two children run concurrently. +// checkpointing. This ensures at most one handover is in flight at any time and +// at most two children run concurrently. func shardedForceReplicationWorker(ctx workflow.Context, params shardedChildParams) (shardedChildResult, error) { s := &shardedWorkflowState{ params: ¶ms, @@ -92,7 +92,7 @@ type shardedWorkflowState struct { // (predecessor still running) and clears it when the parent sends the // resumeFullRate signal (predecessor done). The first child // (Handover=false) has no predecessor and starts clear. - // - Shutdown: a child re-enters handover when it cuts — sends its + // - Shutdown: a child re-enters handover when it checkpoints — sends its // checkpoint signal and stops listing — so the successor the parent // starts from that checkpoint can run at half alongside it. handover bool @@ -114,8 +114,8 @@ type shardedWorkflowState struct { // (see the handover field): // - Startup handover: this child was started alongside a still-running // predecessor; together they must stay ≤ MaxExecsPerShard per shard. -// - Shutdown handover (after cut): a successor has been started at half -// rate alongside us; same combined-load constraint. +// - Shutdown handover (after checkpointing): a successor has been started at +// half rate alongside us; same combined-load constraint. // // Minimum of 1 ensures at least one exec can always be packed regardless // of BatchSize or MaxExecsPerShard settings. @@ -135,7 +135,7 @@ func (s *shardedWorkflowState) run(ctx workflow.Context) (shardedChildResult, er parentExec := workflow.GetInfo(ctx).ParentWorkflowExecution s.startBackgroundCoroutines(ctx, parentExec) - reachedEnd := s.listUntilCutOrEnd(ctx, parentExec) + reachedEnd := s.listUntilCheckpointOrEnd(ctx, parentExec) if s.lastErr != nil { return shardedChildResult{}, s.lastErr } @@ -168,7 +168,7 @@ func (s *shardedWorkflowState) startBackgroundCoroutines(ctx workflow.Context, p // handover (clears handover), advancing the child to full rate once the // predecessor has completed. Runs until the signal arrives (or ctx is // cancelled at workflow end). Only throttled children ever receive it, and - // always before they cut, so it never races the shutdown handover. + // always before they checkpoint, so it never races the shutdown handover. workflow.Go(ctx, func(gCtx workflow.Context) { ch := workflow.GetSignalChannel(gCtx, shardedResumeFullSignalName) var dummy struct{} @@ -188,23 +188,33 @@ func (s *shardedWorkflowState) startBackgroundCoroutines(ctx workflow.Context, p if err := workflow.NewTimer(gCtx, 60*time.Second).Get(gCtx, nil); err != nil { return // ctx cancelled — workflow is completing } - // Best-effort: signal failure is not fatal for the child. - _ = workflow.SignalExternalWorkflow(gCtx, + // The parent is alive for the whole child lifetime — it awaits all + // children before completing or continuing-as-new, and a child is + // terminated with its parent (ParentClosePolicy TERMINATE). So a + // progress signal can only fail for a permanent, anomalous reason + // (SignalExternalWorkflow has no transient failure modes). Latch it + // and stop: continuing would mean operating against a parent we can + // no longer reach, with unpredictable results. + if err := workflow.SignalExternalWorkflow(gCtx, parentExec.ID, parentExec.RunID, shardedProgressSignalName, shardedProgressPayload{ ChildRunID: workflow.GetInfo(gCtx).WorkflowExecution.RunID, VerifiedCount: s.verifiedCount, - }).Get(gCtx, nil) + }).Get(gCtx, nil); err != nil { + s.setLastErr(fmt.Errorf("progress signal to parent: %w", err)) + return + } } }) } -// listUntilCutOrEnd drives the listing loop: page through ListWorkflows, +// listUntilCheckpointOrEnd drives the listing loop: page through ListWorkflows, // bucketing and opportunistically packing each page, until either the // namespace is exhausted (returns true) or a CAN hint at a fully-consumed page -// boundary triggers a cut (returns false). Latches s.lastErr on listing failure. -func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExec *workflow.Execution) bool { +// boundary triggers a checkpoint (returns false). Latches s.lastErr on listing +// failure. +func (s *shardedWorkflowState) listUntilCheckpointOrEnd(ctx workflow.Context, parentExec *workflow.Execution) bool { currentPageToken := s.params.StartPageToken for s.lastErr == nil { executions, nextPageToken, err := s.listWorkflowPageWithToken(ctx, currentPageToken) @@ -233,7 +243,7 @@ func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExe // Check the CAN hint at this fully-consumed page boundary, so // currentPageToken is exact and safe to hand to a successor. if workflow.GetInfo(ctx).GetContinueAsNewSuggested() { - s.cutAtBoundary(ctx, parentExec, currentPageToken) + s.checkpointAtBoundary(ctx, parentExec, currentPageToken) return false } @@ -247,17 +257,25 @@ func (s *shardedWorkflowState) listUntilCutOrEnd(ctx workflow.Context, parentExe return false } -// cutAtBoundary performs the handover cut. If still in the startup handover it -// first awaits promotion (so at most one handover is ever in flight), then -// re-enters handover for the shutdown phase (dropping to half rate) and signals -// the parent the next page token; the parent starts a successor from there and -// promotes it once this child completes. -func (s *shardedWorkflowState) cutAtBoundary(ctx workflow.Context, parentExec *workflow.Execution, pageToken []byte) { +// checkpointAtBoundary performs the handover checkpoint. If still in the startup +// handover it first awaits promotion (so at most one handover is ever in +// flight), then re-enters handover for the shutdown phase (dropping to half +// rate) and signals the parent the next page token; the parent starts a +// successor from there and promotes it once this child completes. +// +// The checkpoint signal is load-bearing: it is the only way the parent learns +// where to resume, so if it fails the parent never starts a successor and every +// execution past this page boundary is silently dropped. SignalExternalWorkflow +// has no transient failure modes (it fails only for permanent reasons, e.g. the +// parent no longer exists), so a failure is latched into lastErr to fail the +// child rather than continue. +func (s *shardedWorkflowState) checkpointAtBoundary(ctx workflow.Context, parentExec *workflow.Execution, pageToken []byte) { if s.handover { // Handover-hits-hint: still in the startup handover, predecessor // running. Pausing here bounds our history at ~hint size and guarantees // at most one handover is in flight at a time. Await promotion - // (predecessor completion, which clears handover), then cut immediately. + // (predecessor completion, which clears handover), then checkpoint + // immediately. _ = workflow.Await(ctx, func() bool { return !s.handover || s.lastErr != nil }) if s.lastErr != nil { return @@ -267,13 +285,15 @@ func (s *shardedWorkflowState) cutAtBoundary(ctx workflow.Context, parentExec *w if parentExec == nil { return } - _ = workflow.SignalExternalWorkflow(ctx, + if err := workflow.SignalExternalWorkflow(ctx, parentExec.ID, parentExec.RunID, shardedCheckpointSignalName, shardedCheckpointPayload{ ChildRunID: workflow.GetInfo(ctx).WorkflowExecution.RunID, NextPageToken: pageToken, - }).Get(ctx, nil) + }).Get(ctx, nil); err != nil { + s.setLastErr(fmt.Errorf("checkpoint signal to parent: %w", err)) + } } var ( From 2bc8ccd05022aac28b7aaa6423f3945390583299 Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Fri, 19 Jun 2026 11:53:51 +0100 Subject: [PATCH 34/35] Comment update. --- service/worker/migration/fx.go | 11 ++++------ .../worker/migration/sharded_activities.go | 8 +++---- .../worker/migration/sharded_handover_test.go | 8 +++---- .../migration/sharded_parent_workflow.go | 17 +++++--------- service/worker/migration/sharded_types.go | 22 +++++++++---------- service/worker/migration/sharded_workflow.go | 11 +++++----- .../worker/migration/sharded_workflow_test.go | 8 +++---- 7 files changed, 36 insertions(+), 49 deletions(-) diff --git a/service/worker/migration/fx.go b/service/worker/migration/fx.go index b4b12e38783..4baec187c8e 100644 --- a/service/worker/migration/fx.go +++ b/service/worker/migration/fx.go @@ -58,10 +58,9 @@ type ( } // shardedWorkerComponent registers the sharded force-replication - // workflow + ReplicateBatch activity on their dedicated TQ. Holds an - // *activities-sized clone so its activity registration is isolated - // from the default-TQ worker — sharded inject paths don't accidentally - // land on the legacy MigrationActivityTQ. + // workflows and the activities they call on a dedicated task queue. + // It holds its own *activities so its activity registration targets + // the sharded TQ rather than the legacy MigrationActivityTQ. shardedWorkerComponent struct { activities *activities } @@ -86,9 +85,7 @@ func NewResult(params initParams) (fxResult, error) { }, nil } -// NewShardedResult constructs the sharded WorkerComponent. The component -// owns its own *activities clone so registration against the sharded TQ -// doesn't bleed into the legacy worker. +// NewShardedResult constructs the sharded WorkerComponent. func NewShardedResult(params initParams) (fxResult, error) { a, err := newActivitiesFromParams(params, shardedForceReplicationWorkflowName) if err != nil { diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index 29d522d7270..a54476d1b71 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -32,9 +32,9 @@ type ( // via AdminService.AddOrUpdateRemoteCluster (which fetches HistoryShardCount // from the remote at registration time and stores it) — a prerequisite for // force replication, since the source generates replication tasks against -// it. ClusterMetadata refreshes the cache every minute, so the value can be -// at most that stale; shard count never changes for a live cluster so this -// is fine. +// it. ClusterMetadata refreshes the cache roughly every minute by default, +// so the value can be slightly stale; shard count never changes for a live +// cluster so this is fine. func (a *activities) DescribeTargetCluster(_ context.Context, req DescribeTargetClusterRequest) (*DescribeTargetClusterResponse, error) { info, ok := a.clusterMetadata.GetAllClusterInfo()[req.TargetClusterName] if !ok { @@ -170,7 +170,6 @@ func (a *activities) evaluateVerifyIteration( shards shardVerifyTracker, doneCount, execCount int, ) (bool, replicateBatchResult, error) { - // Clean completion — every exec verified. if doneCount >= execCount { return true, replicateBatchResult{ CompletedShards: shards.allCompleted(), @@ -178,7 +177,6 @@ func (a *activities) evaluateVerifyIteration( }, nil } - // Per-shard cumulative no-progress backstop. if sErr := a.checkStuckShard(req, shards, execs, verified, doneCount, execCount); sErr != nil { return false, replicateBatchResult{}, sErr } diff --git a/service/worker/migration/sharded_handover_test.go b/service/worker/migration/sharded_handover_test.go index b6992ba384b..9c575f228f8 100644 --- a/service/worker/migration/sharded_handover_test.go +++ b/service/worker/migration/sharded_handover_test.go @@ -19,9 +19,9 @@ package migration // CHILD tests run shardedForceReplicationWorker directly as ROOT. env.SetContinueAsNewSuggested(true) // targets the production child's env directly so the CAN-hint code path executes. The // trade-off is that workflow.GetInfo(ctx).ParentWorkflowExecution is nil for root -// workflows, which triggers a nil pointer dereference in the current production code -// (sharded_workflow.go). See TestChild_ThrottledHitsHint_PausesUntilPromoted and the -// note in TestChild_CheckpointAtHint_StopsListing for details. +// workflows (handled by the guard in run()), so the parent-directed checkpoint and +// progress signals are skipped. See TestChild_ThrottledHitsHint_PausesUntilPromoted and +// the note in TestChild_CheckpointAtHint_StopsListing for details. import ( "context" @@ -457,7 +457,7 @@ func TestChild_ThrottledHitsHint_PausesUntilPromoted(t *testing.T) { // child lifecycle: a promoted child (not throttled) lists a namespace, dispatches // batches, and returns ReachedEnd=true with the correct VerifiedCount. // This uses childDirectRunner (real child with parent) so ParentWorkflowExecution -// is non-nil, sidestepping the production nil-guard bug. +// is non-nil and the parent-directed checkpoint/progress signals fire. func TestChild_PromotedAndCheckpoints_VerifiedCountAccumulated(t *testing.T) { suite := &testsuite.WorkflowTestSuite{} env := suite.NewTestWorkflowEnvironment() diff --git a/service/worker/migration/sharded_parent_workflow.go b/service/worker/migration/sharded_parent_workflow.go index 3c8f771e7af..036a1eae9eb 100644 --- a/service/worker/migration/sharded_parent_workflow.go +++ b/service/worker/migration/sharded_parent_workflow.go @@ -250,7 +250,7 @@ type shardedParentState struct { liveExecs map[string]workflow.Execution // liveRunIDs is the insertion-ordered slice of live child run IDs. - // wireCount tracks how many of these have been wired into the + // wiredCount tracks how many of these have been wired into the // selector; after each sel.Select, newly appended children are wired. liveRunIDs []string wiredCount int @@ -568,18 +568,11 @@ func (ps *shardedParentState) onChildCompleted(ctx workflow.Context, runID strin ps.reachedEnd = true } - // Promote the successor to full rate if one was started. The - // successor's run ID is the last entry in liveRunIDs that is still - // live (i.e., the one started from this child's checkpoint). - // We use successorStarted to know whether a successor was launched - // for this particular child; if so, find it by scanning backwards - // through liveRunIDs for a still-live child that we haven't - // previously promoted. + // Promote the successor to full rate if one was started for this + // child. Handover is sequential (at most one successor in flight), so + // the successor is simply the remaining live child other than this + // one; successorStarted records whether one was launched. if ps.successorStarted[runID] { - // Find the successor: scan liveRunIDs for the first live child - // after this one's position. Since handover is sequential (at - // most one in flight), the successor is simply any remaining - // live child. for _, candidateRunID := range ps.liveRunIDs { if candidateRunID == runID { continue diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index e1792280127..9ae9fe14bb9 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -21,10 +21,9 @@ const ( // releaseShardsSignalName carries mid-flight ReleaseShards signals // from active replicate-batch activities back to their parent - // workflow (the child). Drain-mode shard completions are gone in - // the new design, so this signal only fires while the activity is - // still running normally and the child needs to free the shard early - // so a successor can pack against it. + // workflow (the child). It fires while the activity is still running + // so the child can free a completed shard early and let a successor + // pack against it. releaseShardsSignalName = "ReleaseShards" // shardedCheckpointSignalName is sent by a child to the parent at @@ -331,8 +330,8 @@ type shardedProgressPayload struct { } // shardedBatchReq is the per-batch activity input. Executions is the -// per-shard, per-BID nested payload — the workflow has marked every -// shard appearing as a top-level key in shardInFlight before dispatch, +// per-shard, per-BID nested payload — the workflow has claimed every +// shard appearing as a top-level key into inFlight before dispatch, // and the activity is responsible for either signal-releasing each shard // mid-flight or listing it in the return value's CompletedShards set. type shardedBatchReq struct { @@ -353,9 +352,10 @@ type shardedBatchReq struct { // replicateBatchResult is the activity's return payload. // -// CompletedShards is informational (the dispatch coroutine's defer clears -// heldByBatch + shardInFlight regardless), but keeping it in the result -// gives metrics a clean handle on "which shards this batch finished". +// CompletedShards is informational (the dispatch coroutine's defer, +// releaseAll, clears the batch's held set and its inFlight shards +// regardless), but keeping it in the result gives metrics a clean handle +// on "which shards this batch finished". type replicateBatchResult struct { CompletedShards []int32 @@ -376,8 +376,8 @@ type replicateBatchHeartbeat struct { // releaseShardsPayload is the body of the mid-flight ReleaseShards signal // an activity sends to its parent workflow when the cumulative idle cost // across its completed-but-not-yet-released shards crosses IdleShardCost. -// The workflow handler clears these shards from shardInFlight + -// heldByBatch[BatchID] so the packer can immediately dispatch new work +// The workflow handler (releaseShards) clears these shards from inFlight +// and held[BatchID] so the packer can immediately dispatch new work // against them while the activity stays running on its still-pending // shards. type releaseShardsPayload struct { diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index ca60c897a77..3fe6f5ee97a 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -40,9 +40,6 @@ func shardedForceReplicationWorker(ctx workflow.Context, params shardedChildPara params: ¶ms, namespaceID: params.NamespaceID, targetShardCount: params.TargetShardCount, - // buckets and batches default to usable zero values (their maps are - // lazily initialised on first mutation). - // // Handover=true → started mid-handover (predecessor still running); // runs at half rate until promoted. Handover=false → first child, no // predecessor, full rate from the start. @@ -186,7 +183,7 @@ func (s *shardedWorkflowState) startBackgroundCoroutines(ctx workflow.Context, p workflow.Go(ctx, func(gCtx workflow.Context) { for { if err := workflow.NewTimer(gCtx, 60*time.Second).Get(gCtx, nil); err != nil { - return // ctx cancelled — workflow is completing + return } // The parent is alive for the whole child lifetime — it awaits all // children before completing or continuing-as-new, and a child is @@ -231,7 +228,9 @@ func (s *shardedWorkflowState) listUntilCheckpointOrEnd(ctx workflow.Context, pa }) } - // Opportunistically pack and dispatch batches from the buckets. + // Ship whatever this page made packable in streaming priority + // (relax=false: fullest shards first) so dispatch slots stay busy + // while listing continues; the drain pass handles the remainder. for s.tryPackStreaming(ctx, false) { //nolint:revive // intentional empty body } @@ -400,7 +399,7 @@ func (s *shardedWorkflowState) waitForDispatchSlot(ctx workflow.Context) { // recordVerified accumulates one batch's verified-exec delta into the // child's running count and emits the per-batch counter metric. -// No-op when verified == 0 so batches with only inject (DisableVerification) +// No-op when verified <= 0 so batches with only inject (DisableVerification) // don't add zeros to the count. func (s *shardedWorkflowState) recordVerified(verified int64) { if verified <= 0 { diff --git a/service/worker/migration/sharded_workflow_test.go b/service/worker/migration/sharded_workflow_test.go index 1cd4fb0a0d7..28dc43340e7 100644 --- a/service/worker/migration/sharded_workflow_test.go +++ b/service/worker/migration/sharded_workflow_test.go @@ -431,10 +431,10 @@ const childWorkerID = "test-sharded-child-worker" // // Two problems it solves: // -// 1. The production child workflow accesses -// workflow.GetInfo(ctx).ParentWorkflowExecution.ID; this panics when the -// workflow has no parent (nil pointer dereference). childDirectRunner -// provides a real parent so ParentWorkflowExecution is non-nil. +// 1. The production child signals its parent execution (checkpoint and +// progress rollups), which the nil-parent guard in run() skips when the +// child runs as root. childDirectRunner provides a real parent so +// ParentWorkflowExecution is non-nil and those signal paths execute. // // 2. env.SignalWorkflow targets the top-level execution (childDirectRunner), // not the child execution. childDirectRunner therefore relays the From bb947bdac20e19ce2d28baefbd2852ed7d7dd2fd Mon Sep 17 00:00:00 2001 From: Rob Holland Date: Wed, 24 Jun 2026 12:33:21 +0100 Subject: [PATCH 35/35] Heartbeat includes verifies. --- .../worker/migration/sharded_activities.go | 110 ++++++++++++++---- service/worker/migration/sharded_types.go | 12 ++ .../worker/migration/sharded_types_test.go | 97 +++++++++++++++ service/worker/migration/sharded_workflow.go | 20 ++-- 4 files changed, 207 insertions(+), 32 deletions(-) diff --git a/service/worker/migration/sharded_activities.go b/service/worker/migration/sharded_activities.go index a54476d1b71..bf55b2adc9c 100644 --- a/service/worker/migration/sharded_activities.go +++ b/service/worker/migration/sharded_activities.go @@ -51,10 +51,13 @@ func (a *activities) DescribeTargetCluster(_ context.Context, req DescribeTarget } // ReplicateBatch is the per-batch activity body for the sharded force -// replication workflow. Runs inject (heartbeat-resumable via -// NextInjectIdx/InjectDone) then verify, signal-releasing completed +// replication workflow. Runs inject then verify, signal-releasing completed // shards mid-flight as their cumulative idle cost crosses IdleShardCost. -// Returns {CompletedShards, VerifiedCount} on success. +// Returns {CompletedShards, VerifiedCount} on success. Both phases are +// heartbeat-resumable — inject via NextInjectIdx/InjectDone, verify via the +// ReleasedShards/VerifiedExecs progress snapshot — so a retry (e.g. a worker +// lost to a deploy and detected via the heartbeat timeout) resumes in place +// rather than re-injecting or re-verifying completed work. func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) (replicateBatchResult, error) { // Flatten once so per-exec bookkeeping (verified[], attempts[], // nextRetryAt[]) can stay index-based. @@ -100,7 +103,7 @@ func (a *activities) ReplicateBatch(ctx context.Context, req *shardedBatchReq) ( return replicateBatchResult{}, fmt.Errorf("get remote admin client for %s: %w", req.TargetClusterName, err) } - return a.runVerifyPhase(ctx, req, execs, execCount, remoteAdminClient, ns) + return a.runVerifyPhase(ctx, req, execs, execCount, remoteAdminClient, ns, hb) } // runVerifyPhase is the verify-phase loop body of ReplicateBatch. It @@ -114,30 +117,17 @@ func (a *activities) runVerifyPhase( execCount int, remoteAdminClient adminservice.AdminServiceClient, ns *namespace.Namespace, + hb replicateBatchHeartbeat, ) (replicateBatchResult, error) { - verified := make([]bool, execCount) + // Seed from the resumed heartbeat so a retried attempt (e.g. a worker lost + // to a deploy, detected via the heartbeat timeout) picks up where the last + // left off rather than re-verifying every exec and re-releasing shards the + // workflow has already freed. A fresh attempt starts from a zero hb. + verified, doneCount, shards := seedVerifyState(execs, hb) attempts := make([]int, execCount) nextRetryAt := make([]time.Time, execCount) - doneCount := 0 - - shards := newShardVerifyTracker(execs) for { - // Worker shutdown short-circuits with a retryable error so the - // SDK reschedules on another worker. Returning here avoids the - // ~HeartbeatTimeout wait that silent worker death would incur - // before the server retries the attempt. Inject is already - // heartbeat-preserved (InjectDone), so retry skips it; verify - // re-runs from scratch but DMS reads are idempotent. - select { - case <-activity.GetWorkerStopChannel(ctx): - return replicateBatchResult{}, temporal.NewApplicationErrorWithOptions( - "worker shutdown", "WorkerShutdown", - temporal.ApplicationErrorOptions{NextRetryDelay: time.Second}, - ) - default: - } - passDelta, minNextRetry, vErr := a.runVerifyPass( ctx, remoteAdminClient, ns, req, execs, verified, attempts, nextRetryAt, shards) doneCount += passDelta @@ -145,7 +135,7 @@ func (a *activities) runVerifyPhase( return replicateBatchResult{}, vErr } - activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) + activity.RecordHeartbeat(ctx, verifyHeartbeat(execs, verified, shards)) if done, result, err := a.evaluateVerifyIteration( ctx, req, execs, verified, shards, doneCount, execCount); err != nil { @@ -158,6 +148,68 @@ func (a *activities) runVerifyPhase( } } +// seedVerifyState reconstructs verify-phase bookkeeping from a resumed +// heartbeat. Execs on already-released shards are marked verified (they were +// all applied before the shard was released) and the shards re-flagged +// released, so the resumed attempt neither re-verifies them nor re-signals a +// release the workflow already acted on; remaining verified execs are replayed +// from hb.VerifiedExecs. The tracker's lastProgress is seeded at now (via +// newShardVerifyTracker), so the no-progress backstop measures from this +// attempt's start — resumed work gets the benefit of the doubt. +func seedVerifyState(execs []*shardedExecutionInfo, hb replicateBatchHeartbeat) ([]bool, int, shardVerifyTracker) { + verified := make([]bool, len(execs)) + shards := newShardVerifyTracker(execs) + now := time.Now() + doneCount := 0 + + released := make(map[int32]bool, len(hb.ReleasedShards)) + for _, sh := range hb.ReleasedShards { + released[sh] = true + } + for i, ex := range execs { + if released[ex.Shard] { + verified[i] = true + doneCount++ + shards.recordVerified(ex.Shard, now) + } + } + for _, i := range hb.VerifiedExecs { + if i < 0 || i >= len(execs) || verified[i] { + continue + } + verified[i] = true + doneCount++ + shards.recordVerified(execs[i].Shard, now) + } + // markReleased after recording so the released shards land with + // pending==0, released==true, and doneAt cleared (no idle-cost accrual). + shards.markReleased(hb.ReleasedShards) + return verified, doneCount, shards +} + +// verifyHeartbeat snapshots verify progress for the activity heartbeat so a +// retry can resume via seedVerifyState. Released shards are recorded by ID; +// their execs are implicitly verified and omitted from VerifiedExecs to keep +// the payload compact. +func verifyHeartbeat(execs []*shardedExecutionInfo, verified []bool, shards shardVerifyTracker) replicateBatchHeartbeat { + released := shards.releasedShards() + releasedSet := make(map[int32]bool, len(released)) + for _, sh := range released { + releasedSet[sh] = true + } + var verifiedExecs []int + for i, ex := range execs { + if verified[i] && !releasedSet[ex.Shard] { + verifiedExecs = append(verifiedExecs, i) + } + } + return replicateBatchHeartbeat{ + InjectDone: true, + ReleasedShards: released, + VerifiedExecs: verifiedExecs, + } +} + // evaluateVerifyIteration runs the post-pass checks (clean completion, // stuck-shard backstop, mid-flight signal release) for a single // verify-loop iteration. Returns done=true with the result when the loop @@ -261,7 +313,7 @@ func (a *activities) runVerifyPass( minNextRetry = earliest(minNextRetry, nextRetryAt[i]) } - activity.RecordHeartbeat(ctx, replicateBatchHeartbeat{InjectDone: true}) + activity.RecordHeartbeat(ctx, verifyHeartbeat(execs, verified, shards)) } return verifiedDelta, minNextRetry, nil } @@ -502,6 +554,14 @@ func (t shardVerifyTracker) allCompleted() []int32 { }) } +// releasedShards returns the shards already signal-released this run, ascending +// so the heartbeat snapshot is deterministic across replays. +func (t shardVerifyTracker) releasedShards() []int32 { + return t.completedShards(func(sv shardVerify) bool { + return sv.released + }) +} + // pickStuck returns (shard, age, true) for the lowest-numbered shard // whose cumulative no-progress duration meets or exceeds its effective // threshold. A shard that hasn't yet produced any verified outcome gets diff --git a/service/worker/migration/sharded_types.go b/service/worker/migration/sharded_types.go index 9ae9fe14bb9..af5f62c631b 100644 --- a/service/worker/migration/sharded_types.go +++ b/service/worker/migration/sharded_types.go @@ -371,6 +371,18 @@ type replicateBatchHeartbeat struct { NextInjectIdx int // InjectDone marks the inject phase as complete; retries skip inject. InjectDone bool + + // ReleasedShards lists shards already signal-released to the workflow. + // Carried so a retried activity neither re-verifies their (already + // applied) execs nor re-sends a ReleaseShards signal the workflow has + // already acted on. Nil during inject. + ReleasedShards []int32 + // VerifiedExecs lists the flatten-order indices of execs verified on + // shards that are not yet released, letting a retry resume verify in + // place rather than re-checking every exec. Execs on ReleasedShards are + // implicitly verified and omitted here to keep the payload small. Nil + // during inject. + VerifiedExecs []int } // releaseShardsPayload is the body of the mid-flight ReleaseShards signal diff --git a/service/worker/migration/sharded_types_test.go b/service/worker/migration/sharded_types_test.go index 3ee84060e3c..480e4f27513 100644 --- a/service/worker/migration/sharded_types_test.go +++ b/service/worker/migration/sharded_types_test.go @@ -152,6 +152,103 @@ func TestNewShardVerifyTracker_SeedsAllShards(t *testing.T) { require.False(t, sv1.lastProgress.IsZero(), "lastProgress must be seeded") } +// seedTestExecs builds a fixed five-exec flattened slice spanning shards +// 0 (indices 0,1), 1 (indices 2,3), and 2 (index 4) — the layout shared by +// the seed/heartbeat round-trip tests below. +func seedTestExecs() []*shardedExecutionInfo { + return []*shardedExecutionInfo{ + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-a", RunID: "r0"}, Shard: 0}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-b", RunID: "r1"}, Shard: 0}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-c", RunID: "r2"}, Shard: 1}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-d", RunID: "r3"}, Shard: 1}, + {ExecutionInfo: &ExecutionInfo{BusinessID: "wf-e", RunID: "r4"}, Shard: 2}, + } +} + +// TestSeedVerifyState_Fresh: a zero heartbeat (first attempt) yields nothing +// verified and a fresh tracker with full pending counts. +func TestSeedVerifyState_Fresh(t *testing.T) { + execs := seedTestExecs() + verified, doneCount, shards := seedVerifyState(execs, replicateBatchHeartbeat{}) + + require.Equal(t, []bool{false, false, false, false, false}, verified) + require.Equal(t, 0, doneCount) + require.Empty(t, shards.releasedShards()) + require.Equal(t, 2, shards[0].pending) + require.Equal(t, 2, shards[1].pending) + require.Equal(t, 1, shards[2].pending) +} + +// TestSeedVerifyState_ResumesReleasedAndVerified: a resumed heartbeat marks +// every exec on a released shard verified (without listing them), replays the +// listed per-exec progress, and flags the released shard so it is neither +// re-verified nor re-released. +func TestSeedVerifyState_ResumesReleasedAndVerified(t *testing.T) { + execs := seedTestExecs() + hb := replicateBatchHeartbeat{ + InjectDone: true, + ReleasedShards: []int32{0}, + VerifiedExecs: []int{2}, + } + + verified, doneCount, shards := seedVerifyState(execs, hb) + + require.Equal(t, []bool{true, true, true, false, false}, verified) + require.Equal(t, 3, doneCount) + + // Shard 0 is released: pending drained, released flagged, doneAt cleared so + // it accrues no idle cost and isn't offered for release again. + require.Equal(t, 0, shards[0].pending) + require.True(t, shards[0].released) + require.True(t, shards[0].doneAt.IsZero()) + require.Equal(t, []int32{0}, shards.releasedShards()) + require.NotContains(t, shards.awaitingRelease(), int32(0), "released shard must not be re-offered") + + // Shard 1 has one of two execs verified; shard 2 is untouched. + require.Equal(t, 1, shards[1].pending) + require.True(t, shards[1].verifiedAny) + require.Equal(t, 1, shards[2].pending) + require.False(t, shards[2].verifiedAny) +} + +// TestSeedVerifyState_IgnoresOutOfRangeIndices: a VerifiedExecs index past the +// exec slice is skipped rather than panicking (defends against tracker/slice +// drift across a resume). +func TestSeedVerifyState_IgnoresOutOfRangeIndices(t *testing.T) { + execs := seedTestExecs() + hb := replicateBatchHeartbeat{InjectDone: true, VerifiedExecs: []int{2, 99, -1}} + + verified, doneCount, _ := seedVerifyState(execs, hb) + + require.Equal(t, []bool{false, false, true, false, false}, verified) + require.Equal(t, 1, doneCount) +} + +// TestVerifyHeartbeat_RoundTrip: a heartbeat snapshot omits released shards' +// execs from VerifiedExecs, and seeding from it reconstructs the same verified +// set and done count. +func TestVerifyHeartbeat_RoundTrip(t *testing.T) { + execs := seedTestExecs() + + // Verify state: shard 0 fully done + released, one exec on shard 1 done. + verified := []bool{true, true, true, false, false} + now := time.Now() + shards := newShardVerifyTracker(execs) + shards.recordVerified(0, now) + shards.recordVerified(0, now) + shards.recordVerified(1, now) + shards.markReleased([]int32{0}) + + hb := verifyHeartbeat(execs, verified, shards) + require.Equal(t, []int32{0}, hb.ReleasedShards) + require.Equal(t, []int{2}, hb.VerifiedExecs, "released shard execs must be omitted") + require.True(t, hb.InjectDone) + + gotVerified, gotDone, _ := seedVerifyState(execs, hb) + require.Equal(t, verified, gotVerified) + require.Equal(t, 3, gotDone) +} + // TestEffectiveMaxExecsPerShard: effectiveMaxExecsPerShard returns the full // MaxExecsPerShard outside a handover phase, and max(cap/2, 1) while in handover. func TestEffectiveMaxExecsPerShard(t *testing.T) { diff --git a/service/worker/migration/sharded_workflow.go b/service/worker/migration/sharded_workflow.go index 3fe6f5ee97a..3814a851038 100644 --- a/service/worker/migration/sharded_workflow.go +++ b/service/worker/migration/sharded_workflow.go @@ -307,20 +307,26 @@ var ( } // Per-exec backoff owns the per-exec retry; MaximumAttempts lets a - // transient activity failure recover via heartbeat-resume without - // losing inject progress. + // transient activity failure (or a worker lost to a deploy) recover via + // heartbeat-resume without losing inject or verify progress. // // 10 attempts is sized for fleet rollouts: a rolling deploy of the - // activity workers can burn several attempts per batch (each shutdown - // surfaces as a retryable WorkerShutdown error). 3 was tight enough - // that two unlucky deploys could exhaust the budget. + // activity workers can burn several attempts per batch (a lost worker is + // detected once its heartbeats lapse and the attempt is retried + // elsewhere). 3 was tight enough that two unlucky deploys could exhaust + // the budget. shardedReplicateBatchRetryPolicy = &temporal.RetryPolicy{ MaximumAttempts: 10, } shardedReplicateBatchActivityOptions = workflow.ActivityOptions{ StartToCloseTimeout: 24 * time.Hour, - HeartbeatTimeout: time.Minute, - RetryPolicy: shardedReplicateBatchRetryPolicy, + // 30s rather than the usual minute: a worker lost to a deploy or crash + // is only detected once its heartbeats lapse, and until then the shards + // this batch holds are stalled. The verify loop heartbeats per exec, so + // 30s sits comfortably above the real inter-heartbeat interval and won't + // false-positive on a slow pass. + HeartbeatTimeout: 30 * time.Second, + RetryPolicy: shardedReplicateBatchRetryPolicy, } )