dynamic partitioning: use probability of AddTask to root partition to estimate overall add task rate

[carlydf]

What changed?

• use probability of AddTask to root partition to estimate overall add task rate, instead of assuming uniform distribution
• to enable this, change the AddTask loadbalancing to ensure that there is at least 1% chance of hitting the root, regardless of how much backlog the root has

Why?

A previous PR changed AddTask load-balancing from uniform random to backlog-aware

How did you test it?

• built
• run locally and tested manually
• covered by existing tests
• added new unit test(s)
• added new functional test(s)

Potential risks

Changes AddTask computation and AddTask load balancing, but improves the estimation of the former, and we are ok with the tradeoff for the latter

---

Note

Medium Risk
Changes write-partition routing (1% root floor) and the task-rate signal that drives partition scaling. Incorrect estimates could delay scale-up/down, but this is not auth or data-path security.

Overview
Replaces the uniform-write assumption used by dynamic partition scaling with an estimate derived from the probability that an AddTask lands on the root.

PickWritePartition now returns estimatedTasksAllPartitions as round(1 / rootProbability). Write routing still prefers partitions with more gap to the backlog cap, but forces at least a 1% chance of hitting the root so the scaler keeps a sample of overall write rate even when the root is full. That estimate is sent only on root writes via the etap-bin gRPC header; the root scaleManager uses it (falling back to write-partition count if missing) and batches by BatchSize * currentWritePartitions.

^{Reviewed by Cursor Bugbot for commit 4014f36. Bugbot is set up for automated code reviews on this repo. Configure here.}

[github-actions]

Claude finished @carlydf's task in 6m 52s —— View job

---

Reviewed 9 files, 3 findings.
· cdf/improve-add-tasks-estimate

[carlydf]

could delete this

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 9a4ee05. Configure here.}

[cursor]

Root check uses wrong partition

Low Severity

appendEstimatedTasksAllPartitions is gated on p.IsRoot(), but p is still the pre-balance root partition. pickClientForWrite only reassigns its local p when choosing a write partition, so the caller always sees the original root. The estimate header is therefore attached on every load-balanced write, including requests sent to non-root partitions, even though the value is only meaningful for root samples.

Additional Locations (2)

• client/matching/client.go#L214-L215
• client/matching/client.go#L325-L326

 

^{Reviewed by Cursor Bugbot for commit 9a4ee05. Configure here.}

[Lakshaymiddha]

Batch threshold is disconnected from the per-call estimate - scaler fires on every root hit when the 1% floor is active

When rootProbability < 0.01, estimatedTasksAllPartitions becomes 100 (from math.Round(1/0.01)). But the batch threshold in AddedTasks is still based on batchSize * currentWritePartitions:

temporal/service/matching/scale_manager.go

Lines 133 to 139 in 9a4ee05

	func (sm *scaleManager) AddedTasks(estimatedTasksAllPartitions int) {
	if sm == nil {
	return
	}
	// wait until the batch has accumulated ~batchSize tasks per partition
	if sm.batch.Add(int64(estimatedTasksAllPartitions)) < sm.batchSize*int64(sm.getLatestWritePartitions()) {

With batchSize=5 and currentWrite=4, the threshold is 5 * 4 = 20, but each root hit adds 100. The threshold is exceeded on the very first call - the batch never accumulates.

The old code kept the threshold proportional to the per-call value:

batchSize := int64(numTasks) * sm.batchSize
if sm.batch.Add(int64(numTasks)) < batchSize {

This always took batchSize calls to fire, regardless of the magnitude of numTasks. The new code breaks that proportionality by using partitionCount on the threshold side and 1/rootProbability on the addend side.

Downstream impact: In

temporal/service/matching/scale_manager.go

Lines 170 to 195 in 9a4ee05

	}
	}
	}
	// callScaler runs the scaler on the accumulated batch and persists the resulting state if it
	// changed.
	// Called from backgroundWork only.
	func (sm *scaleManager) callScaler() {
	// don't bother calling during cooldown period
	if !sm.nextDecision.IsZero() && sm.timeSource.Now().Before(sm.nextDecision) {
	return
	}
	settings := sm.settings()
	shadowMode := settings.ShadowModeLogInterval > 0
	// Entering shadow mode on top of a previously-applied managed target releases
	// control back to the dynamic-config baseline: zero the managed target once so
	// the write side follows dynamic config again (and tracks future config
	// changes), and cold-start the shadow simulation from the baseline. BacklogState
	// is preserved, so read partitions are not dropped here (reclaiming them is left
	// to the drain path). This is the only state shadow mode writes; it still never
	// applies the scaler's hypothetical decisions.
	if shadowMode && sm.scaleState.GetTarget() != 0 {
	sm.releaseManagedState(settings)
	}

, batch.Swap(0) drains whatever has accumulated. Since the batch fires after 1 hit
instead of 5, NumTasks per scaler call is ~100 instead of the intended ~500 (batchSize * estimate). The scaler sees a 5x weaker rate signal, which could cause under-scaling.

Suggested fix - restore the proportional relationship:

  func (sm *scaleManager) AddedTasks(estimatedTasksAllPartitions int) {                                                                                                                                                                       
      if sm == nil {
          return                                                                                                                                                                                                                              
      }           
      batchThreshold := int64(estimatedTasksAllPartitions) * sm.batchSize
      if sm.batch.Add(int64(estimatedTasksAllPartitions)) < batchThreshold {
          return
      }
      select {
      case sm.wakeup <- struct{}{}:                                                                                                                                                                                                           
      default:
      }                                                                                                                                                                                                                                       
  }               

This matches the old behavior: always takes batchSize samples to fire, regardless of the estimate magnitude. The floor case (estimate=100) would need 5 root hits to accumulate 500 before waking, just like the non-floor case (estimate=4) needs 5 hits to reach 20.
or
Option B: Use max of both to handle mixed estimates
Since different calls might have different estimates (100 from floor, 4 from normal routing), use the larger:

  batchThreshold := max(                                                                                                                                                                                                                      
      int64(estimatedTasksAllPartitions) * sm.batchSize,                                                                                                                                                                                      
      sm.batchSize * int64(sm.getLatestWritePartitions()),
  ) 

[Lakshaymiddha]

No metric when the root-probability floor activates
When rootProbability < 0.01, the load balancer overrides normal backlog-aware routing to force 1% of traffic to root - even though root is nearly full. It's the only branch in pickWritePartitionByGap that actively pushes traffic against the natural routing decision, and it does so silently.

If root backlog starts growing unexpectedly in production, an operator today would have to read the source code and manually calculate rootGap / totalGap across all partitions to figure out whether the floor is the cause. A simple counter (e.g. write_partition_root_floor_engaged) would make this immediately visible on a dashboard and let them separate "root is getting forced traffic for sampling" from "root is getting traffic because something else is wrong."