schedulertest: add buffer-bound and action-budget invariants

Adds two invariants to CheckInvariants, with the driver instrumentation
they need:

- buffer-bound: total buffered starts must never exceed MaxBufferSize
  (the generator sizes each fill against this cap).
- action-budget: a single InvokerExecuteTask must not take more actions
  (start + terminate + cancel RPCs) than MaxActionsPerExecution. Guards
  the double-spend class where the budget resets across the
  terminate/cancel/start phases.

Driver now counts action RPCs via the mock clients, tracks the max taken
by any single side-effect task dispatch, and exposes BufferedStartsCount
/ MaxBufferSize / MaxActionsObserved / MaxActionsPerExecution on Snapshot.
Adds WithConfig so tests can set small Tweakables to exercise the caps.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: davidporter-id-au

chasm/lib/scheduler/schedulertest/driver.go

@@ -56,13 +56,27 @@ type Driver struct {
 	Frontend *workflowservicemock.MockWorkflowServiceClient
 	History  *historyservicemock.MockHistoryServiceClient
 
+	// Config is the scheduler config the library was wired with. Snapshot reads
+	// Tweakables (MaxBufferSize, MaxActionsPerExecution) from it so the
+	// buffer-bound and action-budget invariants can be evaluated.
+	Config *scheduler.Config
+
 	ref     chasm.ComponentRef
 	execKey chasm.ExecutionKey
 
 	// dispatchedSideEffects tracks physical side-effect tasks already executed,
 	// keyed by pointer identity, since MockNodeBackend accumulates tasks append-only.
 	dispatchedSideEffects map[*tasks.ChasmTask]struct{}
 
+	// actionRPCs points at the running count of action RPCs (StartWorkflowExecution
+	// + TerminateWorkflowExecution + RequestCancelWorkflowExecution) made through
+	// the mock clients. It is allocated before the clients so the stub closures can
+	// increment it. maxActionsObserved is the high-water mark of action RPCs made by
+	// any single side-effect task dispatch (one InvokerExecuteTask), which backs the
+	// action-budget invariant.
+	actionRPCs         *int
+	maxActionsObserved int
+
 	// AfterStep, if set, is called after every step's CloseTransaction with the
 	// driver, so invariant checks can observe each settled state.
 	AfterStep func(d *Driver)
@@ -72,6 +86,7 @@ type driverConfig struct {
 	schedule  *schedulepb.Schedule
 	startTime time.Time
 	afterStep func(d *Driver)
+	config    *scheduler.Config
 }
 
 // DriverOption configures a Driver at construction.
@@ -92,6 +107,13 @@ func WithAfterStep(fn func(d *Driver)) DriverOption {
 	return func(c *driverConfig) { c.afterStep = fn }
 }
 
+// WithConfig overrides the scheduler config (default: DefaultConfig). Use it to
+// set small Tweakables — e.g. MaxBufferSize or MaxActionsPerExecution — so the
+// buffer-bound and action-budget invariants are meaningfully exercised.
+func WithConfig(cfg *scheduler.Config) DriverOption {
+	return func(c *driverConfig) { c.config = cfg }
+}
+
 // NewDriver constructs a Driver and starts a scheduler execution via the
 // production CreateScheduler path. startTime defaults to a fixed instant for
 // determinism.
@@ -134,6 +156,7 @@ func newDriverScaffold(t *testing.T, opts ...DriverOption) (*Driver, *driverConf
 	cfg := &driverConfig{
 		schedule:  DefaultSchedule(),
 		startTime: time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
+		config:    DefaultConfig(),
 	}
 	for _, opt := range opts {
 		opt(cfg)
@@ -147,9 +170,29 @@ func newDriverScaffold(t *testing.T, opts ...DriverOption) (*Driver, *driverConf
 
 	frontend := workflowservicemock.NewMockWorkflowServiceClient(ctrl)
 	history := historyservicemock.NewMockHistoryServiceClient(ctrl)
+
+	// actionRPCs counts every action RPC (start/terminate/cancel) the invoker
+	// makes through the mock clients, so Step can measure how many actions a single
+	// ExecuteTask takes (the action-budget invariant).
+	actionRPCs := new(int)
+	countAction := func() { *actionRPCs++ }
+
 	// Default happy-path stubs; tests can layer more specific expectations.
 	frontend.EXPECT().StartWorkflowExecution(gomock.Any(), gomock.Any()).
-		Return(&workflowservice.StartWorkflowExecutionResponse{RunId: "run-id"}, nil).AnyTimes()
+		DoAndReturn(func(context.Context, *workflowservice.StartWorkflowExecutionRequest, ...grpc.CallOption) (*workflowservice.StartWorkflowExecutionResponse, error) {
+			countAction()
+			return &workflowservice.StartWorkflowExecutionResponse{RunId: "run-id"}, nil
+		}).AnyTimes()
+	history.EXPECT().TerminateWorkflowExecution(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(context.Context, *historyservice.TerminateWorkflowExecutionRequest, ...grpc.CallOption) (*historyservice.TerminateWorkflowExecutionResponse, error) {
+			countAction()
+			return &historyservice.TerminateWorkflowExecutionResponse{}, nil
+		}).AnyTimes()
+	history.EXPECT().RequestCancelWorkflowExecution(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(context.Context, *historyservice.RequestCancelWorkflowExecutionRequest, ...grpc.CallOption) (*historyservice.RequestCancelWorkflowExecutionResponse, error) {
+			countAction()
+			return &historyservice.RequestCancelWorkflowExecutionResponse{}, nil
+		}).AnyTimes()
 	// The callbacks task (used for migrated running workflows) describes each
 	// running workflow; report it completed so watchers resolve deterministically.
 	history.EXPECT().DescribeWorkflowExecution(gomock.Any(), gomock.Any()).
@@ -168,7 +211,7 @@ func newDriverScaffold(t *testing.T, opts ...DriverOption) (*Driver, *driverConf
 	if err := registry.Register(&chasm.CoreLibrary{}); err != nil {
 		t.Fatalf("failed to register core library: %v", err)
 	}
-	if err := registry.Register(NewTestLibrary(DefaultConfig(), logger, specProcessor, frontend, history)); err != nil {
+	if err := registry.Register(NewTestLibrary(cfg.config, logger, specProcessor, frontend, history)); err != nil {
 		t.Fatalf("failed to register scheduler library: %v", err)
 	}
 
@@ -196,7 +239,9 @@ func newDriverScaffold(t *testing.T, opts ...DriverOption) (*Driver, *driverConf
 		Logger:                logger,
 		Frontend:              frontend,
 		History:               history,
+		Config:                cfg.config,
 		dispatchedSideEffects: make(map[*tasks.ChasmTask]struct{}),
+		actionRPCs:            actionRPCs,
 		AfterStep:             cfg.afterStep,
 	}
 	return d, cfg
@@ -352,6 +397,7 @@ func (d *Driver) Step() bool {
 			continue
 		}
 		d.dispatchedSideEffects[ct] = struct{}{}
+		actionsBefore := *d.actionRPCs
 		execErr := node.ExecuteSideEffectTask(
 			engineCtx,
 			d.execKey,
@@ -368,6 +414,11 @@ func (d *Driver) Step() bool {
 			}
 			d.T.Fatalf("ExecuteSideEffectTask failed: %v", execErr)
 		}
+		// Track the most actions any single side-effect task (i.e. one
+		// InvokerExecuteTask) took, for the action-budget invariant.
+		if taken := *d.actionRPCs - actionsBefore; taken > d.maxActionsObserved {
+			d.maxActionsObserved = taken
+		}
 	}
 	if _, err := node.CloseTransaction(); err != nil {
 		d.T.Fatalf("CloseTransaction after side-effect tasks failed: %v", err)
@@ -489,14 +540,31 @@ type Snapshot struct {
 	// GeneratorLPT / InvokerLPT are the generator and invoker high-water marks.
 	GeneratorLPT time.Time
 	InvokerLPT   time.Time
+
+	// BufferedStartsCount is len(invoker.GetBufferedStarts()) — the total buffered
+	// starts (pending + retained). MaxBufferSize is the configured cap; the
+	// buffer-bound invariant requires BufferedStartsCount <= MaxBufferSize.
+	BufferedStartsCount int
+	MaxBufferSize       int
+
+	// MaxActionsObserved is the high-water mark of action RPCs (start + terminate +
+	// cancel) taken by any single InvokerExecuteTask so far. MaxActionsPerExecution
+	// is the configured per-execution cap; the action-budget invariant requires
+	// MaxActionsObserved <= MaxActionsPerExecution.
+	MaxActionsObserved     int
+	MaxActionsPerExecution int
 }
 
 // Snapshot reads current scheduler state into a Snapshot.
 func (d *Driver) Snapshot() Snapshot {
 	d.T.Helper()
+	tweakables := d.Config.Tweakables(Namespace)
 	snap := Snapshot{
-		Now:            d.Now(),
-		HasPendingTask: d.HasPendingTask(),
+		Now:                    d.Now(),
+		HasPendingTask:         d.HasPendingTask(),
+		MaxBufferSize:          tweakables.MaxBufferSize,
+		MaxActionsObserved:     d.maxActionsObserved,
+		MaxActionsPerExecution: tweakables.MaxActionsPerExecution,
 	}
 	d.ReadScheduler(func(s *scheduler.Scheduler, ctx chasm.Context) {
 		snap.Closed = s.Closed
@@ -513,8 +581,11 @@ func (d *Driver) Snapshot() Snapshot {
 			if g := s.Generator.Get(ctx); g != nil && g.LastProcessedTime != nil {
 				snap.GeneratorLPT = g.LastProcessedTime.AsTime()
 			}
-			if inv := s.Invoker.Get(ctx); inv != nil && inv.LastProcessedTime != nil {
-				snap.InvokerLPT = inv.LastProcessedTime.AsTime()
+			if inv := s.Invoker.Get(ctx); inv != nil {
+				snap.BufferedStartsCount = len(inv.GetBufferedStarts())
+				if inv.LastProcessedTime != nil {
+					snap.InvokerLPT = inv.LastProcessedTime.AsTime()
+				}
 			}
 		}
 	})

chasm/lib/scheduler/schedulertest/invariants.go

@@ -48,6 +48,27 @@ func CheckInvariants(prev *Snapshot, cur Snapshot) []Violation {
 			cur.Paused, cur.BackfillerCount, cur.IdleCloseTime.Format("2006-01-02T15:04:05"))
 	}
 
+	// (2a) Buffer bound: the total buffered starts must never exceed MaxBufferSize.
+	// The generator sizes each fill against this cap (remaining = MaxBufferSize -
+	// len(buffered)); exceeding it means the limit was not enforced and CHASM state
+	// can grow without bound (oversized state, persistence failures, OOM).
+	if cur.MaxBufferSize > 0 && cur.BufferedStartsCount > cur.MaxBufferSize {
+		add("buffer-bound",
+			"buffered starts = %d, exceeding MaxBufferSize=%d",
+			cur.BufferedStartsCount, cur.MaxBufferSize)
+	}
+
+	// (2b) Action budget: a single InvokerExecuteTask must not take more actions
+	// (start + terminate + cancel RPCs) than MaxActionsPerExecution. Exceeding it
+	// means the per-execution rate protection was bypassed — e.g. the action budget
+	// reset across the terminate/cancel/start phases, letting one execution take up
+	// to a multiple of the configured maximum.
+	if cur.MaxActionsPerExecution > 0 && cur.MaxActionsObserved > cur.MaxActionsPerExecution {
+		add("action-budget",
+			"a single ExecuteTask took %d actions, exceeding MaxActionsPerExecution=%d",
+			cur.MaxActionsObserved, cur.MaxActionsPerExecution)
+	}
+
 	if prev != nil {
 		// (3) Closed is terminal: a closed schedule must never reopen.
 		if prev.Closed && !cur.Closed {

chasm/lib/scheduler/schedulertest/invariants_test.go

@@ -113,3 +113,29 @@ func TestCheckInvariants_DetectsHwmRegression(t *testing.T) {
 	require.NotEmpty(t, violations)
 	require.Equal(t, "hwm-monotonic-generator", violations[0].Name)
 }
+
+// TestCheckInvariants_DetectsBufferOverrun confirms the buffer-bound invariant
+// fires when buffered starts exceed MaxBufferSize, and that being exactly at the
+// cap is allowed.
+func TestCheckInvariants_DetectsBufferOverrun(t *testing.T) {
+	over := schedulertest.Snapshot{HasPendingTask: true, MaxBufferSize: 1000, BufferedStartsCount: 1001}
+	violations := schedulertest.CheckInvariants(nil, over)
+	require.NotEmpty(t, violations, "buffer overrun must be flagged")
+	require.Equal(t, "buffer-bound", violations[0].Name)
+
+	atCap := schedulertest.Snapshot{HasPendingTask: true, MaxBufferSize: 1000, BufferedStartsCount: 1000}
+	require.Empty(t, schedulertest.CheckInvariants(nil, atCap), "exactly at the cap is allowed")
+}
+
+// TestCheckInvariants_DetectsBudgetOverrun confirms the action-budget invariant
+// fires when a single ExecuteTask takes more actions than MaxActionsPerExecution
+// (the double-spend bug class), and that being exactly at the cap is allowed.
+func TestCheckInvariants_DetectsBudgetOverrun(t *testing.T) {
+	over := schedulertest.Snapshot{HasPendingTask: true, MaxActionsPerExecution: 5, MaxActionsObserved: 6}
+	violations := schedulertest.CheckInvariants(nil, over)
+	require.NotEmpty(t, violations, "budget overrun must be flagged")
+	require.Equal(t, "action-budget", violations[0].Name)
+
+	atCap := schedulertest.Snapshot{HasPendingTask: true, MaxActionsPerExecution: 5, MaxActionsObserved: 5}
+	require.Empty(t, schedulertest.CheckInvariants(nil, atCap), "exactly at the cap is allowed")
+}