diff --git a/common/dynamicconfig/constants.go b/common/dynamicconfig/constants.go index ea8a747f4a..67e9dfcaff 100644 --- a/common/dynamicconfig/constants.go +++ b/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, diff --git a/common/limiter/keyed_bytes_limiter.go b/common/limiter/keyed_bytes_limiter.go new file mode 100644 index 0000000000..302e453f35 --- /dev/null +++ b/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{} +} diff --git a/common/limiter/keyed_bytes_limiter_test.go b/common/limiter/keyed_bytes_limiter_test.go new file mode 100644 index 0000000000..1a3f5671e1 --- /dev/null +++ b/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()) +} diff --git a/common/metrics/metric_defs.go b/common/metrics/metric_defs.go index c16da974da..8a7e676d78 100644 --- a/common/metrics/metric_defs.go +++ b/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 diff --git a/service/history/api/create_workflow_util.go b/service/history/api/create_workflow_util.go index 15f0f55258..c62b3be4df 100644 --- a/service/history/api/create_workflow_util.go +++ b/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, diff --git a/service/history/api/respondworkflowtaskcompleted/api.go b/service/history/api/respondworkflowtaskcompleted/api.go index 11caeb673c..58eb2888f1 100644 --- a/service/history/api/respondworkflowtaskcompleted/api.go +++ b/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, ) diff --git a/service/history/chasm_engine.go b/service/history/chasm_engine.go index e5ff2d5318..e86cf2480e 100644 --- a/service/history/chasm_engine.go +++ b/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, diff --git a/service/history/configs/config.go b/service/history/configs/config.go index 3009f7c48d..e04dce32ac 100644 --- a/service/history/configs/config.go +++ b/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 diff --git a/service/history/ndc/history_replicator.go b/service/history/ndc/history_replicator.go index 3c451f640a..f84e09c308 100644 --- a/service/history/ndc/history_replicator.go +++ b/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( diff --git a/service/history/ndc/mutable_state_initializer.go b/service/history/ndc/mutable_state_initializer.go index 93c91e8ea9..56bbf66b7c 100644 --- a/service/history/ndc/mutable_state_initializer.go +++ b/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 { diff --git a/service/history/ndc/resetter_test.go b/service/history/ndc/resetter_test.go index 4dbce13eb7..d42d56fcae 100644 --- a/service/history/ndc/resetter_test.go +++ b/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() diff --git a/service/history/ndc/workflow_resetter.go b/service/history/ndc/workflow_resetter.go index c8f9287801..cba7b8cc8e 100644 --- a/service/history/ndc/workflow_resetter.go +++ b/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( diff --git a/service/history/ndc/workflow_state_replicator.go b/service/history/ndc/workflow_state_replicator.go index a659a51f43..9dfd4e7807 100644 --- a/service/history/ndc/workflow_state_replicator.go +++ b/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( diff --git a/service/history/shard/context_testutil.go b/service/history/shard/context_testutil.go index 4bdc9081e9..7ef7ecbefa 100644 --- a/service/history/shard/context_testutil.go +++ b/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 } diff --git a/service/history/statemachine_environment_test.go b/service/history/statemachine_environment_test.go index 8d35d23b86..0183508a67 100644 --- a/service/history/statemachine_environment_test.go +++ b/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 } diff --git a/service/history/timer_queue_active_task_executor.go b/service/history/timer_queue_active_task_executor.go index a293434b5e..7b17f0a70d 100644 --- a/service/history/timer_queue_active_task_executor.go +++ b/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, ) diff --git a/service/history/workflow/cache/cache.go b/service/history/workflow/cache/cache.go index 70352fde26..8688a766bb 100644 --- a/service/history/workflow/cache/cache.go +++ b/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 diff --git a/service/history/workflow/cache/cache_test.go b/service/history/workflow/cache/cache_test.go index dbd9937888..522f6428be 100644 --- a/service/history/workflow/cache/cache_test.go +++ b/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, diff --git a/service/history/workflow/context.go b/service/history/workflow/context.go index 412bbc9efe..fd355721c8 100644 --- a/service/history/workflow/context.go +++ b/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 diff --git a/service/history/workflow/context_test.go b/service/history/workflow/context_test.go index 46a1948d2d..e177f7fe4c 100644 --- a/service/history/workflow/context_test.go +++ b/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()) } diff --git a/tests/workflow_completion_pagination_test.go b/tests/workflow_completion_pagination_test.go index 8cab6d29d1..401dee0574 100644 --- a/tests/workflow_completion_pagination_test.go +++ b/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