WCP 6/X: Process-wide limit for paged RespondWorkflowTaskCompleted requests (#11151)

## What changed?
Adds two limits on the memory held by in-flight buffer for paginated
RespondWorkflowTaskCompleted requests:

- Process-wide limit on the total bytes buffered across all workflows on
a history host. Dynamic config is
`WorkflowTaskCompletionBufferTotalSizeLimit`
- Per-namespace share of that process-wide limit (expressed as a ratio),
so a single namespace can't consume the entire host budget.
`WorkflowTaskCompletionBufferNamespaceRatio`

When buffering a page would push either limit over the top, the page is
rejected with the existing transient buffer-lost signal and the
in-progress buffer is dropped.

## Why?
Pagination lets one workflow task ship large volume of commands split
across several requests, which the server holds in memory until the
final page arrives. Without a ceiling, many concurrent large completions
or one heavy namespace could exhaust a history host's memory. The
process-wide limit bounds total exposure, and the per-namespace share
keeps one namespace from starving the others.

## How did you test it?
- [X] built
- [X] run locally and tested manually
- [X] covered by existing tests
- [X] added new unit test(s)
- [X] added new functional test(s)

## Potential risks
* Gated behind the pagination feature flag (off by default), so there's
no effect until pagination is enabled for a namespace.
* The limit is enforced per process; a namespace that stays over its
share will keep hitting buffer-lost and retrying until its in-flight
buffers drain. This is recoverable but could show up as retry churn if a
limit is set too low.

common/dynamicconfig/constants.go

@@ -1785,12 +1785,26 @@ See DynamicRateLimitingParams comments for more details.`,
 		`EnableWorkflowTaskCompletionPagination enables the pagination of RespondWorkflowTaskCompleted requests.
 		When false, paginated requests (the ones with intermediate_page set to true) are rejected.`,
 	)
+	WorkflowTaskCompletionBufferTotalSizeLimit = NewGlobalIntSetting(
+		"history.workflowTaskCompletionBufferTotalSizeLimit",
+		1024*1024*1024,
+		`WorkflowTaskCompletionBufferTotalSizeLimit is the process wide limit in bytes on the total
+		size of buffers allocated for in-flight pages of RespondWorkflowTaskCompleted requests. A page that would push
+		the total over this limit is rejected. Setting to 0 disables the limit.`,
+	)
 	WorkflowTaskCompletionBufferSizeLimit = NewNamespaceIntSetting(
 		"history.workflowTaskCompletionBufferSizeLimit",
 		40*1024*1024,
 		`WorkflowTaskCompletionBufferSizeLimit is the limit in bytes on the total
 		size of buffered pages in paginated RespondWorkflowTaskCompleted requests for a single workflow task.`,
 	)
+	WorkflowTaskCompletionBufferNamespaceRatio = NewNamespaceFloatSetting(
+		"history.workflowTaskCompletionBufferNamespaceRatio",
+		0.5,
+		`WorkflowTaskCompletionBufferNamespaceRatio is the fraction of the process-wide pagination buffer
+		limit (WorkflowTaskCompletionBufferTotalSizeLimit) that a single namespace may hold at once, so one
+		namespace cannot exhaust the whole process budget.`,
+	)
 	HistoryLongPollExpirationInterval = NewNamespaceDurationSetting(
 		"history.longPollExpirationInterval",
 		time.Second*20,

common/limiter/keyed_bytes_limiter.go

@@ -0,0 +1,60 @@
+package limiter
+
+import "sync"
+
+// KeyedBytesLimiter tracks in-flight bytes and enforces a global total plus a per-key
+// sub-limit, so one key cannot exhaust the whole total. The limits are supplied by the
+// caller on each reservation, so the limiter itself holds no configuration and only does
+// the accounting. It is safe for concurrent use.
+type KeyedBytesLimiter struct {
+	mu     sync.Mutex
+	used   int64
+	perKey map[string]int64
+}
+
+// TryReserve adds n bytes to the totals if both the global total stays within totalLimit
+// and the key's total within keyLimit. A non-positive limit disables that check. Returns
+// success and the resulting global total.
+func (l *KeyedBytesLimiter) TryReserve(key string, n, totalLimit, keyLimit int64) (bool, int64) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if totalLimit > 0 && l.used+n > totalLimit {
+		return false, l.used
+	}
+	if keyLimit > 0 && l.perKey[key]+n > keyLimit {
+		return false, l.used
+	}
+	l.used += n
+	if l.perKey == nil {
+		l.perKey = make(map[string]int64)
+	}
+	l.perKey[key] += n
+	return true, l.used
+}
+
+// Release returns n bytes for the given key and reports the resulting global total.
+func (l *KeyedBytesLimiter) Release(key string, n int64) int64 {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if n != 0 {
+		l.used -= n
+		if remaining := l.perKey[key] - n; remaining > 0 {
+			l.perKey[key] = remaining
+		} else {
+			delete(l.perKey, key)
+		}
+	}
+	return l.used
+}
+
+// Used reports the current global in-flight byte total.
+func (l *KeyedBytesLimiter) Used() int64 {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.used
+}
+
+// NewKeyedBytesLimiter returns an empty limiter.
+func NewKeyedBytesLimiter() *KeyedBytesLimiter {
+	return &KeyedBytesLimiter{}
+}

common/limiter/keyed_bytes_limiter_test.go

@@ -0,0 +1,81 @@
+package limiter
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestKeyedBytesLimiter_Reserve(t *testing.T) {
+	l := &KeyedBytesLimiter{}
+	const key = "k"
+
+	// Reservation within limit succeeds and tracks the running total.
+	ok, used := l.TryReserve(key, 100, 250, 0)
+	require.True(t, ok)
+	require.Equal(t, int64(100), used)
+
+	ok, used = l.TryReserve(key, 100, 250, 0)
+	require.True(t, ok)
+	require.Equal(t, int64(200), used)
+
+	// A reservation that would exceed the limit is rejected.
+	ok, used = l.TryReserve(key, 100, 250, 0)
+	require.False(t, ok)
+	require.Equal(t, int64(200), used)
+
+	// A non-positive limit disables the check.
+	ok, _ = l.TryReserve(key, 1<<40, 0, 0)
+	require.True(t, ok)
+
+	// Release brings the counter back down.
+	l.Release(key, 1<<40)
+	l.Release(key, 200)
+	require.Equal(t, int64(0), l.Used())
+	require.Equal(t, int64(0), l.perKey[key])
+}
+
+// TestKeyedBytesLimiter_KeyLimit verifies the per-key limit is enforced independently of
+// the global total and that one key's usage does not block another.
+func TestKeyedBytesLimiter_KeyLimit(t *testing.T) {
+	l := &KeyedBytesLimiter{}
+
+	// k1 can fill up to its key limit even though the total has room.
+	ok, _ := l.TryReserve("k1", 100, 1000, 150)
+	require.True(t, ok)
+	ok, _ = l.TryReserve("k1", 100, 1000, 150)
+	require.False(t, ok, "second reservation exceeds k1's 150-byte key limit")
+	require.Equal(t, int64(100), l.perKey["k1"])
+
+	// A different key has its own budget and is unaffected.
+	ok, _ = l.TryReserve("k2", 100, 1000, 150)
+	require.True(t, ok)
+	require.Equal(t, int64(200), l.Used(), "global total tracks both keys")
+
+	l.Release("k1", 100)
+	l.Release("k2", 100)
+	require.Equal(t, int64(0), l.Used())
+}
+
+// TestKeyedBytesLimiter_ConcurrentNetsToZero exercises concurrent reserve/release on a
+// shared limiter.
+func TestKeyedBytesLimiter_ConcurrentNetsToZero(t *testing.T) {
+	l := &KeyedBytesLimiter{}
+	const goroutines = 50
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	for range goroutines {
+		go func() {
+			defer wg.Done()
+			for range 100 {
+				if ok, _ := l.TryReserve("k", 1024, 1<<40, 0); ok {
+					l.Release("k", 1024)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+	require.Equal(t, int64(0), l.Used())
+}

common/metrics/metric_defs.go

@@ -1214,9 +1214,13 @@ var (
 		"workflow_task_completion_paginated_bytes",
 		WithDescription("Total wire size of successfully completed paginated RespondWorkflowTaskCompleted requests. count givens the total number of successful paginated requests."),
 	)
+	WorkflowTaskCompletionBufferInflightBytes = NewGaugeDef(
+		"workflow_task_completion_buffer_inflight_bytes",
+		WithDescription("Process-wide total bytes currently held across all in-flight workflow task completion pagination buffers."),
+	)
 	WorkflowTaskCompletionBufferLost = NewCounterDef(
 		"workflow_task_completion_buffer_lost",
-		WithDescription("Paginated workflow task completions aborted because the buffer was lost (evicted or a page was missing)."),
+		WithDescription("Paginated workflow task completions aborted because the buffer was lost (evicted, process limit exceeded, or a page was missing)."),
 	)
 
 	// Matching

service/history/api/create_workflow_util.go

@@ -147,6 +147,7 @@ func NewWorkflowLeaseAndContext(
 			shardCtx.GetLogger(),
 			shardCtx.GetThrottledLogger(),
 			shardCtx.GetMetricsHandler(),
+			nil, // no pagination buffer limiter as it is a transient context
 		),
 		wcache.NoopReleaseFn,
 		ms,

service/history/api/respondworkflowtaskcompleted/api.go

@@ -657,6 +657,7 @@ func (handler *WorkflowTaskCompletedHandler) Invoke(
 				handler.logger,
 				handler.shardContext.GetThrottledLogger(),
 				handler.shardContext.GetMetricsHandler(),
+				nil, // no pagination buffer limiter as it is a transient context
 			),
 			newMutableState,
 		)

service/history/chasm_engine.go

@@ -906,6 +906,7 @@ func (e *ChasmEngine) createNewExecutionWithUpdate(
 			shardContext.GetLogger(),
 			shardContext.GetThrottledLogger(),
 			shardContext.GetMetricsHandler(),
+			nil, // no pagination buffer limiter as it is a transient context
 		),
 		mutableState: mutableState,
 		snapshot:     snapshot,

service/history/configs/config.go

@@ -44,15 +44,17 @@ type Config struct {
 	VisibilityAllowList                     dynamicconfig.BoolPropertyFnWithNamespaceFilter
 	SuppressErrorSetSystemSearchAttribute   dynamicconfig.BoolPropertyFnWithNamespaceFilter
 
-	EmitShardLagLog                        dynamicconfig.BoolPropertyFn
-	EnableDataLossMetrics                  dynamicconfig.BoolPropertyFn
-	ThrottledLogRPS                        dynamicconfig.IntPropertyFn
-	EnableStickyQuery                      dynamicconfig.BoolPropertyFnWithNamespaceFilter
-	EnableWorkflowTaskCompletionPagination dynamicconfig.BoolPropertyFnWithNamespaceFilter
-	AlignMembershipChange                  dynamicconfig.DurationPropertyFn
-	WorkflowTaskCompletionBufferSizeLimit  dynamicconfig.IntPropertyFnWithNamespaceFilter
-	ShutdownDrainDuration                  dynamicconfig.DurationPropertyFn
-	StartupMembershipJoinDelay             dynamicconfig.DurationPropertyFn
+	EmitShardLagLog                            dynamicconfig.BoolPropertyFn
+	EnableDataLossMetrics                      dynamicconfig.BoolPropertyFn
+	ThrottledLogRPS                            dynamicconfig.IntPropertyFn
+	EnableStickyQuery                          dynamicconfig.BoolPropertyFnWithNamespaceFilter
+	EnableWorkflowTaskCompletionPagination     dynamicconfig.BoolPropertyFnWithNamespaceFilter
+	AlignMembershipChange                      dynamicconfig.DurationPropertyFn
+	WorkflowTaskCompletionBufferTotalSizeLimit dynamicconfig.IntPropertyFn
+	WorkflowTaskCompletionBufferSizeLimit      dynamicconfig.IntPropertyFnWithNamespaceFilter
+	WorkflowTaskCompletionBufferNamespaceRatio dynamicconfig.FloatPropertyFnWithNamespaceFilter
+	ShutdownDrainDuration                      dynamicconfig.DurationPropertyFn
+	StartupMembershipJoinDelay                 dynamicconfig.DurationPropertyFn
 
 	// Workflow reset related settings.
 	AllowResetWithPendingChildren dynamicconfig.BoolPropertyFnWithNamespaceFilter
@@ -832,8 +834,10 @@ func NewConfig(
 		RoutingInfoCacheMaxSize:              dynamicconfig.RoutingInfoCacheMaxSize.Get(dc),
 
 		// Workflow task completion pagination
-		EnableWorkflowTaskCompletionPagination: dynamicconfig.EnableWorkflowTaskCompletionPagination.Get(dc),
-		WorkflowTaskCompletionBufferSizeLimit:  dynamicconfig.WorkflowTaskCompletionBufferSizeLimit.Get(dc),
+		EnableWorkflowTaskCompletionPagination:     dynamicconfig.EnableWorkflowTaskCompletionPagination.Get(dc),
+		WorkflowTaskCompletionBufferTotalSizeLimit: dynamicconfig.WorkflowTaskCompletionBufferTotalSizeLimit.Get(dc),
+		WorkflowTaskCompletionBufferSizeLimit:      dynamicconfig.WorkflowTaskCompletionBufferSizeLimit.Get(dc),
+		WorkflowTaskCompletionBufferNamespaceRatio: dynamicconfig.WorkflowTaskCompletionBufferNamespaceRatio.Get(dc),
 	}
 
 	return cfg

service/history/ndc/history_replicator.go

@@ -599,6 +599,7 @@ func (r *HistoryReplicatorImpl) applyNonStartEventsToCurrentBranch(
 			r.logger,
 			r.shardContext.GetThrottledLogger(),
 			r.shardContext.GetMetricsHandler(),
+			nil, // no pagination buffer limiter as it is a transient context
 		)
 
 		newWorkflow = NewWorkflow(

service/history/ndc/mutable_state_initializer.go

@@ -161,6 +161,7 @@ func (r *MutableStateInitializerImpl) InitializeFromToken(
 		r.logger,
 		r.shardContext.GetThrottledLogger(),
 		r.shardContext.GetMetricsHandler(),
+		nil, // no pagination buffer limiter as it is a transient context
 	)
 	mutableStateRow, dbRecordVersion, dbHistorySize, existsInDB, err := r.deserializeBackfillToken(token)
 	if err != nil {

service/history/ndc/resetter_test.go

@@ -93,6 +93,7 @@ func (s *resetterSuite) SetupTest() {
 		s.logger,
 		s.mockShard.GetThrottledLogger(),
 		s.mockShard.GetMetricsHandler(),
+		nil,
 	)
 	s.newRunID = uuid.NewString()
 

service/history/ndc/workflow_resetter.go

@@ -533,6 +533,7 @@ func (r *workflowResetterImpl) replayResetWorkflow(
 		r.logger,
 		r.shardContext.GetLogger(),
 		r.shardContext.GetMetricsHandler(),
+		nil, // no pagination buffer limiter as it is a transient context
 	)
 
 	resetMutableState, resetStats, err := r.stateRebuilder.Rebuild(

service/history/ndc/workflow_state_replicator.go

@@ -1130,6 +1130,7 @@ func (r *WorkflowStateReplicatorImpl) getNewRunWorkflow(
 		r.logger,
 		r.shardContext.GetThrottledLogger(),
 		r.shardContext.GetMetricsHandler(),
+		nil, // no pagination buffer limiter as it is a transient context
 	)
 
 	return NewWorkflow(

service/history/shard/context_testutil.go

@@ -12,6 +12,7 @@ import (
 	"go.temporal.io/server/common/cache"
 	"go.temporal.io/server/common/clock"
 	"go.temporal.io/server/common/cluster"
+	"go.temporal.io/server/common/finalizer"
 	"go.temporal.io/server/common/future"
 	"go.temporal.io/server/common/locks"
 	"go.temporal.io/server/common/log"
@@ -238,6 +239,12 @@ func (s *ContextTest) StopForTest() {
 	s.FinishStop()
 }
 
+// SetFinalizerForTest overrides the shard's finalizer. Production shards always have one, so
+// tests that exercise cache paths gated on a finalizer being present must set it explicitly.
+func (s *ContextTest) SetFinalizerForTest(f *finalizer.Finalizer) {
+	s.finalizer = f
+}
+
 func (s *StubContext) GetEngine(_ context.Context) (historyi.Engine, error) {
 	return s.engine, nil
 }

service/history/statemachine_environment_test.go

@@ -275,7 +275,7 @@ func TestValidateStateMachineRef(t *testing.T) {
 			tc.mutateNode(node)
 			tc.mutateRef(&ref)
 
-			workflowContext := workflow.NewContext(s.mockShard.GetConfig(), mutableState.GetWorkflowKey(), chasm.WorkflowArchetypeID, log.NewTestLogger(), log.NewTestLogger(), metrics.NoopMetricsHandler)
+			workflowContext := workflow.NewContext(s.mockShard.GetConfig(), mutableState.GetWorkflowKey(), chasm.WorkflowArchetypeID, log.NewTestLogger(), log.NewTestLogger(), metrics.NoopMetricsHandler, nil)
 			if tc.clearTransitionHistory {
 				mutableState.GetExecutionInfo().TransitionHistory = nil
 			}

service/history/timer_queue_active_task_executor.go

@@ -788,6 +788,7 @@ func (t *timerQueueActiveTaskExecutor) executeWorkflowRunTimeoutTask(
 			t.logger,
 			t.shardContext.GetThrottledLogger(),
 			t.shardContext.GetMetricsHandler(),
+			nil, // no pagination buffer limiter as it is a transient context
 		),
 		newMutableState,
 	)

service/history/workflow/cache/cache.go

@@ -15,6 +15,7 @@ import (
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/finalizer"
 	"go.temporal.io/server/common/headers"
+	"go.temporal.io/server/common/limiter"
 	"go.temporal.io/server/common/locks"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
@@ -63,6 +64,8 @@ type (
 		onPut                     func(wfContext *historyi.WorkflowContext)
 		onEvict                   func(wfContext *historyi.WorkflowContext)
 		nonUserContextLockTimeout time.Duration
+		// paginationLimiter limits pagination buffer size across all workflow cache contexts
+		paginationLimiter *limiter.KeyedBytesLimiter
 	}
 	cacheItem struct {
 		shardId   int32
@@ -126,17 +129,21 @@ func NewHostLevelCache(
 		OnEvict: func(val any) {
 			//revive:disable-next-line:unchecked-type-assertion
 			item := val.(*cacheItem)
-			if item.finalizer == nil {
-				return // should only happen in unit tests
-			}
-			wfKey := item.wfContext.GetWorkflowKey()
-			err := item.finalizer.Deregister(wfKey.String())
-			if err != nil {
-				// debug level since this is very common: the cache item was registered with a finalizer
-				// that has been finalized since then and is therefore no longer accepting any calls
-				logger.Debug("cache failed to de-register callback in finalizer",
-					tag.Error(err), tag.ShardID(item.shardId))
+			if item.finalizer != nil {
+				wfKey := item.wfContext.GetWorkflowKey()
+				err := item.finalizer.Deregister(wfKey.String())
+				if err != nil {
+					// debug level since this is very common: the cache item was registered with a finalizer
+					// that has been finalized since then and is therefore no longer accepting any calls
+					logger.Debug("cache failed to de-register callback in finalizer",
+						tag.Error(err), tag.ShardID(item.shardId))
+					return
+				}
 			}
+			// We removed the finalizer callback before it ran, so this eviction now owns clearing
+			// the context. Without this, resources it holds, for example the bytes it reserved on the
+			// shared pagination buffer limiter would leak until process restart.
+			item.wfContext.Clear()
 		},
 	}
 
@@ -145,6 +152,7 @@ func NewHostLevelCache(
 	return &cacheImpl{
 		Cache:                     c,
 		nonUserContextLockTimeout: config.HistoryCacheNonUserContextLockTimeout(),
+		paginationLimiter:         limiter.NewKeyedBytesLimiter(),
 	}
 }
 
@@ -289,6 +297,7 @@ func (c *cacheImpl) getOrCreateWorkflowExecutionInternal(
 			shardContext.GetLogger(),
 			shardContext.GetThrottledLogger(),
 			shardContext.GetMetricsHandler(),
+			c.paginationLimiter,
 		)
 
 		var err error

service/history/workflow/cache/cache_test.go

@@ -11,15 +11,18 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	commandpb "go.temporal.io/api/command/v1"
 	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/api/workflowservice/v1"
 	enumsspb "go.temporal.io/server/api/enums/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/common/cache"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/dynamicconfig"
+	"go.temporal.io/server/common/finalizer"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/locks"
 	"go.temporal.io/server/common/metrics"
@@ -225,6 +228,78 @@ func (s *workflowCacheSuite) TestHistoryCachePinning() {
 	release(err4)
 }
 
+// TestHistoryCacheEvictionReleasesPaginationBuffer verifies that when a cached context
+// holding an in-flight pagination buffer is evicted, the bytes it
+// reserved on the shared pagination limiter are released. Without a Clear() on the eviction
+// path those bytes leak and eventually reject all paginated completions on the host.
+func (s *workflowCacheSuite) TestHistoryCacheEvictionReleasesPaginationBuffer() {
+	// Count-based cache with room for a single entry, so the next insert forces eviction.
+	s.mockShard.GetConfig().HistoryHostLevelCacheMaxSize = dynamicconfig.GetIntPropertyFn(1)
+	// Production shards always have a finalizer, so give the test shard one too
+	s.mockShard.SetFinalizerForTest(finalizer.New(s.mockShard.GetLogger(), metrics.NoopMetricsHandler))
+	namespaceID := namespace.ID("test_namespace_id")
+	s.cache = NewHostLevelCache(s.mockShard.GetConfig(), s.mockShard.GetLogger(), metrics.NoopMetricsHandler)
+	limiter := s.cache.(*cacheImpl).paginationLimiter
+
+	we := commonpb.WorkflowExecution{
+		WorkflowId: "wf-cache-test-eviction-buffer",
+		RunId:      uuid.NewString(),
+	}
+	ctx, release, err := s.cache.GetOrCreateWorkflowExecution(
+		context.Background(),
+		s.mockShard,
+		namespaceID,
+		&we,
+		locks.PriorityHigh,
+	)
+	s.NoError(err)
+
+	// Give the context a MutableState so it can buffer a page and later Clear().
+	mock := historyi.NewMockMutableState(s.controller)
+	mock.EXPECT().IsDirty().Return(false).AnyTimes()
+	mock.EXPECT().GetNamespaceEntry().Return(tests.LocalNamespaceEntry).AnyTimes()
+	mock.EXPECT().GetStartedWorkflowTask().Return(&historyi.WorkflowTaskInfo{}).AnyTimes()
+	mock.EXPECT().GetQueryRegistry().Return(workflow.NewQueryRegistry()).AnyTimes()
+	mock.EXPECT().RemoveSpeculativeWorkflowTaskTimeoutTask().AnyTimes()
+	wfCtx := ctx.(*workflow.ContextImpl)
+	wfCtx.MutableState = mock
+
+	// Buffer an intermediate page; this reserves bytes on the shared limiter.
+	page := &workflowservice.RespondWorkflowTaskCompletedRequest{
+		IntermediatePage: true,
+		PageNumber:       0,
+		Commands: []*commandpb.Command{{
+			CommandType: enumspb.COMMAND_TYPE_RECORD_MARKER,
+			Attributes: &commandpb.Command_RecordMarkerCommandAttributes{
+				RecordMarkerCommandAttributes: &commandpb.RecordMarkerCommandAttributes{MarkerName: "m"},
+			},
+		}},
+	}
+	s.NoError(wfCtx.AppendTaskCompletionPage(10, 1, page))
+	s.Positive(limiter.Used())
+
+	// Unpin the entry so it is eligible for eviction.
+	release(nil)
+
+	// Insert a second workflow, forcing eviction of the first from the cache.
+	we2 := commonpb.WorkflowExecution{
+		WorkflowId: "wf-cache-test-eviction-buffer-2",
+		RunId:      uuid.NewString(),
+	}
+	_, release2, err := s.cache.GetOrCreateWorkflowExecution(
+		context.Background(),
+		s.mockShard,
+		namespaceID,
+		&we2,
+		locks.PriorityHigh,
+	)
+	s.NoError(err)
+	release2(nil)
+
+	// The evicted context must have released its reserved pagination bytes.
+	s.Equal(int64(0), limiter.Used())
+}
+
 func (s *workflowCacheSuite) TestHistoryCacheClear() {
 	s.mockShard.GetConfig().HistoryHostLevelCacheMaxSize = dynamicconfig.GetIntPropertyFn(20)
 	namespaceID := namespace.ID("test_namespace_id")
@@ -559,6 +634,7 @@ func (s *workflowCacheSuite) TestCacheImpl_lockWorkflowExecution() {
 				s.mockShard.GetLogger(),
 				s.mockShard.GetThrottledLogger(),
 				s.mockShard.GetMetricsHandler(),
+				nil,
 			)
 			ctx := headers.SetCallerType(context.Background(), tt.callerType)
 			ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
@@ -663,6 +739,9 @@ func (s *workflowCacheSuite) TestCacheImpl_RejectsRequestWhenAtLimitMultiple() {
 	}
 	mockMS1 := historyi.NewMockMutableState(s.controller)
 	mockMS1.EXPECT().IsDirty().Return(false).AnyTimes()
+	// Eviction now clears the evicted context, which drains these on its MutableState.
+	mockMS1.EXPECT().GetQueryRegistry().Return(workflow.NewQueryRegistry()).AnyTimes()
+	mockMS1.EXPECT().RemoveSpeculativeWorkflowTaskTimeoutTask().AnyTimes()
 
 	ctx, release1, err := s.cache.GetOrCreateWorkflowExecution(
 		context.Background(),
@@ -694,6 +773,8 @@ func (s *workflowCacheSuite) TestCacheImpl_RejectsRequestWhenAtLimitMultiple() {
 	}
 	mockMS2 := historyi.NewMockMutableState(s.controller)
 	mockMS2.EXPECT().IsDirty().Return(false).AnyTimes()
+	mockMS2.EXPECT().GetQueryRegistry().Return(workflow.NewQueryRegistry()).AnyTimes()
+	mockMS2.EXPECT().RemoveSpeculativeWorkflowTaskTimeoutTask().AnyTimes()
 	ctx, release2, err := s.cache.GetOrCreateWorkflowExecution(
 		context.Background(),
 		mockShard,
@@ -722,6 +803,8 @@ func (s *workflowCacheSuite) TestCacheImpl_RejectsRequestWhenAtLimitMultiple() {
 	}
 	mockMS3 := historyi.NewMockMutableState(s.controller)
 	mockMS3.EXPECT().IsDirty().Return(false).AnyTimes()
+	mockMS3.EXPECT().GetQueryRegistry().Return(workflow.NewQueryRegistry()).AnyTimes()
+	mockMS3.EXPECT().RemoveSpeculativeWorkflowTaskTimeoutTask().AnyTimes()
 	_, _, err = s.cache.GetOrCreateWorkflowExecution(
 		context.Background(),
 		mockShard,

service/history/workflow/context.go

@@ -18,6 +18,7 @@ import (
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/cluster"
 	"go.temporal.io/server/common/definition"
+	"go.temporal.io/server/common/limiter"
 	"go.temporal.io/server/common/locks"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
@@ -51,6 +52,9 @@ type (
 		// pagination of RespondWorkflowTaskCompleted requests; nil when no
 		// pagination is in progress
 		taskCompletionBuffer *TaskCompletionBuffer
+		// paginationLimiter enforces the process-wide and per-namespace limits on the
+		// total size of all in-flight pagination buffers. nil is treated as "no limit".
+		paginationLimiter *limiter.KeyedBytesLimiter
 	}
 
 	// workflowTaskIdentity identifies a specific workflow task attempt
@@ -64,8 +68,9 @@ type (
 	// pagination of RespondWorkflowTaskCompleted requests.
 	TaskCompletionBuffer struct {
 		pages     map[int32][]*commandpb.Command // page_number (0-based) -> commands
-		totalSize int64                          // cumulative buffered bytes
+		totalSize int64                          // cumulative size of buffered commands
 		identity  workflowTaskIdentity           // the workflow task this buffer belongs to
+		namespace string                         // namespace name
 	}
 )
 
@@ -80,6 +85,10 @@ const maxWorkflowTaskCompletionPages int32 = 1024
 // pagination and the handler fails the workflow task
 var ErrTaskCompletionBufferSizeExceeded = errors.New("workflow task completion buffer size exceeds the per-workflow limit")
 
+// NewContext builds a workflow context. paginationLimiter enforces the process-wide and
+// per-namespace pagination buffer limits; it is only needed for cached contexts
+// that buffer paginated RespondWorkflowTaskCompleted requests, so transient contexts
+// (replication, reset, new-run creation) pass nil, which disables the limit.
 func NewContext(
 	config *configs.Config,
 	workflowKey definition.WorkflowKey,
@@ -87,6 +96,7 @@ func NewContext(
 	logger log.Logger,
 	throttledLogger log.ThrottledLogger,
 	metricsHandler metrics.Handler,
+	paginationLimiter *limiter.KeyedBytesLimiter,
 ) *ContextImpl {
 	tags := func() []tag.Tag {
 		return []tag.Tag{
@@ -96,13 +106,14 @@ func NewContext(
 		}
 	}
 	contextImpl := &ContextImpl{
-		workflowKey:     workflowKey,
-		archetypeID:     archetypeID,
-		logger:          log.NewLazyLogger(logger, tags),
-		throttledLogger: log.NewLazyLogger(throttledLogger, tags),
-		metricsHandler:  metricsHandler.WithTags(metrics.OperationTag(metrics.WorkflowContextScope)),
-		config:          config,
-		lock:            locks.NewPrioritySemaphore(1),
+		workflowKey:       workflowKey,
+		archetypeID:       archetypeID,
+		logger:            log.NewLazyLogger(logger, tags),
+		throttledLogger:   log.NewLazyLogger(throttledLogger, tags),
+		metricsHandler:    metricsHandler.WithTags(metrics.OperationTag(metrics.WorkflowContextScope)),
+		config:            config,
+		lock:              locks.NewPrioritySemaphore(1),
+		paginationLimiter: paginationLimiter,
 	}
 	softassert.That(
 		contextImpl.throttledLogger,
@@ -145,11 +156,16 @@ func (c *ContextImpl) Clear() {
 	c.clearTaskCompletionBuffer()
 }
 
-// clearTaskCompletionBuffer drops the in-progress buffer
+// clearTaskCompletionBuffer drops the in-progress buffer and returns its reserved
+// bytes to the limiter
 func (c *ContextImpl) clearTaskCompletionBuffer() {
 	if c.taskCompletionBuffer == nil {
 		return
 	}
+	if c.paginationLimiter != nil {
+		used := c.paginationLimiter.Release(c.taskCompletionBuffer.namespace, c.taskCompletionBuffer.totalSize)
+		metrics.WorkflowTaskCompletionBufferInflightBytes.With(c.metricsHandler).Record(float64(used))
+	}
 	c.taskCompletionBuffer = nil
 }
 
@@ -199,6 +215,7 @@ func (c *ContextImpl) AppendTaskCompletionPage(
 		c.clearTaskCompletionBuffer()
 		return err
 	}
+	nsName := c.MutableState.GetNamespaceEntry().Name().String()
 	// The request's token supplies schedID/attempt; the version comes from the started
 	// workflow task
 	identity := workflowTaskIdentity{schedID: schedID, attempt: attempt, version: c.startedWorkflowTaskIdentity().version}
@@ -208,8 +225,9 @@ func (c *ContextImpl) AppendTaskCompletionPage(
 	}
 	if c.taskCompletionBuffer == nil {
 		c.taskCompletionBuffer = &TaskCompletionBuffer{
-			pages:    make(map[int32][]*commandpb.Command),
-			identity: identity,
+			pages:     make(map[int32][]*commandpb.Command),
+			identity:  identity,
+			namespace: nsName,
 		}
 	}
 	// Keep existing page if it is already buffered
@@ -220,13 +238,31 @@ func (c *ContextImpl) AppendTaskCompletionPage(
 	pageBytes := taskCompletionPageBytes(request.Commands)
 
 	// Apply per-workflow task limit
-	nsName := c.MutableState.GetNamespaceEntry().Name().String()
 	perWorkflowLimitBytes := int64(c.config.WorkflowTaskCompletionBufferSizeLimit(nsName))
 	if perWorkflowLimitBytes > 0 && c.taskCompletionBuffer.totalSize+pageBytes > perWorkflowLimitBytes {
 		c.clearTaskCompletionBuffer()
 		return ErrTaskCompletionBufferSizeExceeded
 	}
 
+	// Apply the process-wide limit and its per-namespace share so one
+	// namespace cannot exhaust the whole process budget.
+	if c.paginationLimiter != nil {
+		processLimit := int64(c.config.WorkflowTaskCompletionBufferTotalSizeLimit())
+		nsRatio := c.config.WorkflowTaskCompletionBufferNamespaceRatio(nsName)
+		nsLimit := int64(nsRatio * float64(processLimit))
+		ok, used := c.paginationLimiter.TryReserve(nsName, pageBytes, processLimit, nsLimit)
+		if !ok {
+			// BufferLost makes the SDK resend from page 0, so retaining the
+			// partial buffer buys. Clear it to release the reserved bytes back
+			// to the budget.
+			c.clearTaskCompletionBuffer()
+			metrics.WorkflowTaskCompletionBufferLost.With(c.metricsHandler).Record(1)
+			return serviceerror.NewWorkflowTaskCompletionBufferLostf(
+				"workflow task completion buffer memory limit reached while buffering page %d", request.GetPageNumber())
+		}
+		metrics.WorkflowTaskCompletionBufferInflightBytes.With(c.metricsHandler).Record(float64(used))
+	}
+
 	c.taskCompletionBuffer.pages[request.GetPageNumber()] = request.Commands
 	c.taskCompletionBuffer.totalSize += pageBytes
 	return nil

service/history/workflow/context_test.go

@@ -20,6 +20,7 @@ import (
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/cluster"
 	"go.temporal.io/server/common/definition"
+	"go.temporal.io/server/common/limiter"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/persistence"
@@ -80,6 +81,7 @@ func (s *contextSuite) SetupTest() {
 		log.NewNoopLogger(),
 		log.NewNoopLogger(),
 		metrics.NoopMetricsHandler,
+		limiter.NewKeyedBytesLimiter(),
 	)
 }
 
@@ -679,6 +681,8 @@ func (s *contextSuite) setTaskCompletionBufferSizeLimit(limit int) {
 	mockMutableState := historyi.NewMockMutableState(gomock.NewController(s.T()))
 	mockMutableState.EXPECT().GetNamespaceEntry().Return(tests.LocalNamespaceEntry).AnyTimes()
 	mockMutableState.EXPECT().GetStartedWorkflowTask().Return(&historyi.WorkflowTaskInfo{}).AnyTimes()
+	mockMutableState.EXPECT().GetQueryRegistry().Return(NewQueryRegistry()).AnyTimes()
+	mockMutableState.EXPECT().RemoveSpeculativeWorkflowTaskTimeoutTask().AnyTimes()
 	s.workflowContext.MutableState = mockMutableState
 }
 
@@ -746,12 +750,100 @@ func (s *contextSuite) TestTaskCompletionBuffer_PageCountLimit() {
 
 // TestTaskCompletionBuffer_PerWorkflowCapTerminates verifies that a page pushing the
 // cumulative buffer past the per-workflow limit returns the terminate sentinel and
-// drops the buffer.
+// releases the process budget.
 func (s *contextSuite) TestTaskCompletionBuffer_PerWorkflowCapTerminates() {
 	s.setTaskCompletionBufferSizeLimit(1)
+	s.workflowContext.paginationLimiter = limiter.NewKeyedBytesLimiter()
 
 	// 1-byte per-workflow limit: any non-empty page exceeds it.
 	err := s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "big"))
 	s.ErrorIs(err, ErrTaskCompletionBufferSizeExceeded)
 	s.Nil(s.workflowContext.taskCompletionBuffer)
+	s.Equal(int64(0), s.workflowContext.paginationLimiter.Used())
+}
+
+// TestTaskCompletionBuffer_ProcessLimitRejects verifies that a page which would push the
+// process total over the limit is rejected with buffer-lost without storing anything.
+func (s *contextSuite) TestTaskCompletionBuffer_ProcessLimitRejects() {
+	budget := limiter.NewKeyedBytesLimiter()
+	s.workflowContext.paginationLimiter = budget
+
+	s.setTaskCompletionBufferSizeLimit(0)
+
+	processLimit := int64(s.workflowContext.config.WorkflowTaskCompletionBufferTotalSizeLimit())
+	ok, _ := budget.TryReserve("filler", processLimit, 0, 0)
+	s.True(ok)
+
+	err := s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "p0"))
+	var bufferLost *serviceerror.WorkflowTaskCompletionBufferLost
+	s.ErrorAs(err, &bufferLost)
+	s.Nil(s.workflowContext.taskCompletionBuffer)
+}
+
+// TestTaskCompletionBuffer_NamespaceRatioRejects verifies that a page which would
+// push a namespace over its share (ratio * process limit) is rejected with
+// buffer-lost even though the process itself has plenty of room.
+func (s *contextSuite) TestTaskCompletionBuffer_NamespaceRatioRejects() {
+	budget := limiter.NewKeyedBytesLimiter()
+	s.workflowContext.paginationLimiter = budget
+
+	s.setTaskCompletionBufferSizeLimit(0)
+	// Process budget is effectively unbounded; the namespace share is ~1 byte, so any
+	// real page trips the per-namespace cap and not the process cap.
+	s.workflowContext.config.WorkflowTaskCompletionBufferTotalSizeLimit = func() int { return 1 << 30 }
+	s.workflowContext.config.WorkflowTaskCompletionBufferNamespaceRatio = func(string) float64 { return 1e-9 }
+
+	err := s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "p0"))
+	var bufferLost *serviceerror.WorkflowTaskCompletionBufferLost
+	s.ErrorAs(err, &bufferLost)
+	s.Nil(s.workflowContext.taskCompletionBuffer)
+	s.Equal(int64(0), budget.Used())
+}
+
+// TestTaskCompletionBuffer_ProcessLimitDropsPartialBuffer verifies that a process limit
+// rejection of a later page drops the whole partial buffer and releases its
+// already-reserved bytes
+func (s *contextSuite) TestTaskCompletionBuffer_ProcessLimitDropsPartialBuffer() {
+	budget := limiter.NewKeyedBytesLimiter()
+	s.workflowContext.paginationLimiter = budget
+
+	s.setTaskCompletionBufferSizeLimit(0)
+	processLimit := int64(s.workflowContext.config.WorkflowTaskCompletionBufferTotalSizeLimit())
+
+	s.NoError(s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "p0")))
+	page0Size := s.workflowContext.taskCompletionBuffer.totalSize
+	s.Positive(page0Size)
+
+	ok, _ := budget.TryReserve("filler", processLimit-page0Size, 0, 0)
+	s.True(ok)
+
+	err := s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(1, "p1"))
+	var bufferLost *serviceerror.WorkflowTaskCompletionBufferLost
+	s.ErrorAs(err, &bufferLost)
+	s.Nil(s.workflowContext.taskCompletionBuffer)
+	s.Equal(processLimit-page0Size, budget.Used())
+}
+
+// TestTaskCompletionBuffer_BudgetReleasedOnMergeAndClear verifies the process counter
+// goes back to zero after a successful merge
+func (s *contextSuite) TestTaskCompletionBuffer_BudgetReleasedOnMergeAndClear() {
+	budget := limiter.NewKeyedBytesLimiter()
+	s.workflowContext.paginationLimiter = budget
+
+	// Disable the per-workflow cap so it never interferes with the process-counter assertions.
+	s.setTaskCompletionBufferSizeLimit(0)
+
+	s.NoError(s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "p0")))
+	s.NoError(s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(1, "p1")))
+	s.Positive(budget.Used())
+
+	_, err := s.workflowContext.GetMergedTaskCompletionPages(10, 1, finalPage(2, nil))
+	s.NoError(err)
+	s.Equal(int64(0), budget.Used())
+
+	// And after Clear() of a fresh buffer.
+	s.NoError(s.workflowContext.AppendTaskCompletionPage(10, 1, intermediatePage(0, "p0")))
+	s.Positive(budget.Used())
+	s.workflowContext.Clear()
+	s.Equal(int64(0), budget.Used())
 }

tests/workflow_completion_pagination_test.go

@@ -48,7 +48,8 @@ func (s *WorkflowCompletionPaginationTestSuite) newPaginationEnv() *testcore.Tes
 	return testcore.NewEnv(s.T(),
 		testcore.WithDynamicConfig(dynamicconfig.MaximumEventBatchSizeInBytes, 1024*1024),
 		testcore.WithDynamicConfig(dynamicconfig.EnableWorkflowTaskCompletionPagination, true),
-		testcore.WithDynamicConfig(dynamicconfig.WorkflowTaskCompletionBufferSizeLimit, 32*1024*1024),
+		testcore.WithDynamicConfig(dynamicconfig.WorkflowTaskCompletionBufferTotalSizeLimit, 1024*1024*1024),
+		testcore.WithDynamicConfig(dynamicconfig.WorkflowTaskCompletionBufferSizeLimit, 256*1024*1024),
 	)
 }
 
@@ -410,6 +411,90 @@ func (s *WorkflowCompletionPaginationTestSuite) TestResendAfterBufferLost() {
 	s.Equal(0, countEvents(history, enumspb.EVENT_TYPE_WORKFLOW_TASK_FAILED))
 }
 
+// TestBufferLimits drives two intermediate marker pages against different total /
+// per-namespace limit combinations
+func (s *WorkflowCompletionPaginationTestSuite) TestBufferLimits() {
+	// process total exhausted
+	s.Run("process total exhausted", func(s *WorkflowCompletionPaginationTestSuite) {
+		s.runBufferLimitCase(markerPayloadSize+markerPayloadSize/2, 1.0, true)
+	})
+	// namespace share exhausted
+	s.Run("namespace share exhausted", func(s *WorkflowCompletionPaginationTestSuite) {
+		s.runBufferLimitCase(markerPayloadSize*3, 0.5, true)
+	})
+	// both sufficiently large
+	s.Run("both limits sufficiently large", func(s *WorkflowCompletionPaginationTestSuite) {
+		s.runBufferLimitCase(markerPayloadSize*100, 1.0, false)
+	})
+}
+
+// runBufferLimitCase buffers page 0 (always fits), then buffers page 1. When
+// expectRejected, page 1 must fail with a transient buffer-lost error and write no
+// events; otherwise page 1 buffers and a final page completes the workflow.
+func (s *WorkflowCompletionPaginationTestSuite) runBufferLimitCase(
+	totalSizeLimit int,
+	namespaceRatio float64,
+	expectRejected bool,
+) {
+	env := testcore.NewEnv(s.T(),
+		testcore.WithDynamicConfig(dynamicconfig.EnableWorkflowTaskCompletionPagination, true),
+		testcore.WithDynamicConfig(dynamicconfig.WorkflowTaskCompletionBufferTotalSizeLimit, totalSizeLimit),
+		testcore.WithDynamicConfig(dynamicconfig.WorkflowTaskCompletionBufferNamespaceRatio, namespaceRatio),
+	)
+	we, _, taskToken := s.startWorkflowAndStartWFT(env, time.Minute)
+
+	bufferMarkerPage := func(page int32) error {
+		_, err := env.FrontendClient().RespondWorkflowTaskCompleted(s.Context(), &workflowservice.RespondWorkflowTaskCompletedRequest{
+			Namespace:        env.Namespace().String(),
+			TaskToken:        taskToken,
+			Commands:         makeMarkerCommands(1),
+			IntermediatePage: true,
+			PageNumber:       page,
+			Identity:         "worker1",
+		})
+		return err
+	}
+
+	// Page 0 always fits.
+	s.NoError(bufferMarkerPage(0))
+
+	err := bufferMarkerPage(1)
+	if expectRejected {
+		// buffer-lost is transient: no WorkflowTaskFailed event, unlike the
+		// per-workflow limit which fails the workflow task.
+		var bufferLost *serviceerror.WorkflowTaskCompletionBufferLost
+		s.ErrorAs(err, &bufferLost)
+
+		history := env.GetHistory(env.Namespace().String(), we)
+		s.Equal(0, countEvents(history, enumspb.EVENT_TYPE_WORKFLOW_TASK_FAILED))
+		s.Equal(0, countEvents(history, enumspb.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED))
+		return
+	}
+
+	// Both pages buffered; the final page completes the workflow.
+	s.NoError(err)
+	_, err = env.FrontendClient().RespondWorkflowTaskCompleted(s.Context(), &workflowservice.RespondWorkflowTaskCompletedRequest{
+		Namespace: env.Namespace().String(),
+		TaskToken: taskToken,
+		Commands: []*commandpb.Command{{
+			CommandType: enumspb.COMMAND_TYPE_COMPLETE_WORKFLOW_EXECUTION,
+			Attributes: &commandpb.Command_CompleteWorkflowExecutionCommandAttributes{
+				CompleteWorkflowExecutionCommandAttributes: &commandpb.CompleteWorkflowExecutionCommandAttributes{
+					Result: payloads.EncodeString("done"),
+				},
+			},
+		}},
+		PageNumber: 2,
+		Identity:   "worker1",
+	})
+	s.NoError(err)
+
+	history := env.GetHistory(env.Namespace().String(), we)
+	s.Equal(2, countEvents(history, enumspb.EVENT_TYPE_MARKER_RECORDED))
+	s.Equal(1, countEvents(history, enumspb.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED))
+	s.Equal(0, countEvents(history, enumspb.EVENT_TYPE_WORKFLOW_TASK_FAILED))
+}
+
 // countEvents returns the number of history events of the given type.
 func countEvents(history []*historypb.HistoryEvent, eventType enumspb.EventType) int {
 	count := 0