From c3e72a212e5f4dc43419bd2cc7930a994200e61f Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Wed, 3 Jun 2026 14:21:36 -0400 Subject: [PATCH 1/3] CHASM return same timestamp from context.Now() within same MutableCtx --- chasm/context.go | 26 ++++++++++++- chasm/tree_test.go | 93 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 2 deletions(-) diff --git a/chasm/context.go b/chasm/context.go index 3f829d231b4..bf1ca11c36e 100644 --- a/chasm/context.go +++ b/chasm/context.go @@ -98,6 +98,8 @@ type immutableCtx struct { type mutableCtx struct { *immutableCtx + + nowByComponent map[*Node]time.Time } // NewContext creates a new Context from an existing Context and root Node. @@ -209,8 +211,27 @@ func NewMutableContext( node *Node, ) MutableContext { return &mutableCtx{ - immutableCtx: newContext(ctx, node), + immutableCtx: newContext(ctx, node), + nowByComponent: make(map[*Node]time.Time), + } +} + +func (c *mutableCtx) Now(component Component) time.Time { + node, ok := c.root.valueToNode[component] + if !ok || !node.isComponent() { + return c.root.Now(component) + } + + if now, ok := c.nowByComponent[node]; ok { + return now + } + + now := c.root.Now(component) + if c.nowByComponent == nil { + c.nowByComponent = make(map[*Node]time.Time) } + c.nowByComponent[node] = now + return now } func (c *mutableCtx) AddTask( @@ -223,7 +244,8 @@ func (c *mutableCtx) AddTask( func (c *mutableCtx) withValue(key any, value any) Context { return &mutableCtx{ - immutableCtx: ContextWithValue(c.immutableCtx, key, value), + immutableCtx: ContextWithValue(c.immutableCtx, key, value), + nowByComponent: c.nowByComponent, } } diff --git a/chasm/tree_test.go b/chasm/tree_test.go index c47fbe7f8c6..13b3d91eff4 100644 --- a/chasm/tree_test.go +++ b/chasm/tree_test.go @@ -3127,6 +3127,37 @@ func (s *nodeSuite) testComponentTree() *Node { return node } +func (s *nodeSuite) TestMutableContextNowStableWithinContext() { + root := s.testComponentTree() + + startTime := time.Date(2026, 1, 1, 1, 0, 0, 0, time.UTC) + updatedTime := startTime.Add(time.Minute) + laterTime := updatedTime.Add(time.Minute) + + s.timeSource.Update(startTime) + + mutableContext := NewMutableContext(context.Background(), root) + component, err := root.Component(mutableContext, ComponentRef{}) + s.NoError(err) + + s.Equal(startTime, mutableContext.Now(component)) + + s.timeSource.Update(updatedTime) + s.Equal(startTime, mutableContext.Now(component)) + + contextWithValue := ContextWithValue(mutableContext, "test-key", "test-value") + s.Equal(startTime, contextWithValue.Now(component)) + + newMutableContext := NewMutableContext(context.Background(), root) + s.Equal(updatedTime, newMutableContext.Now(component)) + + immutableContext := NewContext(context.Background(), root) + s.Equal(updatedTime, immutableContext.Now(component)) + + s.timeSource.Update(laterTime) + s.Equal(laterTime, immutableContext.Now(component)) +} + func (s *nodeSuite) TestExecuteImmediatePureTask() { root := s.testComponentTree() @@ -3182,6 +3213,68 @@ func (s *nodeSuite) TestExecuteImmediatePureTask() { s.Equal(tasks.MaximumKey.FireTime, s.nodeBackend.LastDeletePureTaskCall()) } +func (s *nodeSuite) TestImmediatePureTaskNowStableWithinTaskOnly() { + root := s.testComponentTree() + + _, err := root.CloseTransaction() + s.NoError(err) + + taskStartTime := time.Date(2026, 1, 1, 2, 0, 0, 0, time.UTC) + nextTaskTime := taskStartTime.Add(time.Minute) + s.timeSource.Update(taskStartTime) + + mutableContext := NewMutableContext(context.Background(), root) + component, err := root.Component(mutableContext, ComponentRef{}) + s.NoError(err) + + taskAttributes := TaskAttributes{ScheduledTime: TaskScheduledTimeImmediate} + mutableContext.AddTask( + component, + taskAttributes, + &TestPureTask{ + Payload: &commonpb.Payload{Data: []byte("task-1")}, + }, + ) + mutableContext.AddTask( + component, + taskAttributes, + &TestPureTask{ + Payload: &commonpb.Payload{Data: []byte("task-2")}, + }, + ) + + s.testLibrary.mockPureTaskHandler.EXPECT(). + Validate(gomock.Any(), gomock.Any(), gomock.Eq(taskAttributes), gomock.Any()).Return(true, nil).Times(2) + + var observedTimes []time.Time + s.testLibrary.mockPureTaskHandler.EXPECT(). + Execute( + gomock.AssignableToTypeOf(&mutableCtx{}), + gomock.AssignableToTypeOf(&TestComponent{}), + gomock.Eq(taskAttributes), + gomock.Any(), + ). + DoAndReturn(func(ctx MutableContext, component any, _ TaskAttributes, _ *TestPureTask) error { + chasmComponent := component.(Component) + firstNow := ctx.Now(chasmComponent) + secondNow := ctx.Now(chasmComponent) + s.Equal(firstNow, secondNow) + + observedTimes = append(observedTimes, firstNow) + if len(observedTimes) == 1 { + s.timeSource.Update(nextTaskTime) + } + return nil + }). + Times(2) + + mutations, err := root.CloseTransaction() + s.NoError(err) + s.Empty(mutations.UpdatedNodes) + s.Empty(mutations.DeletedNodes) + s.Equal([]time.Time{taskStartTime, nextTaskTime}, observedTimes) +} + func (s *nodeSuite) TestEachPureTask() { now := s.timeSource.Now() From f8d1183420f7459126d683e85d66e4f40122886c Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Fri, 5 Jun 2026 11:10:19 -0400 Subject: [PATCH 2/3] Fix linter issues --- chasm/tree_test.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/chasm/tree_test.go b/chasm/tree_test.go index 13b3d91eff4..67e7a434c39 100644 --- a/chasm/tree_test.go +++ b/chasm/tree_test.go @@ -3231,16 +3231,12 @@ func (s *nodeSuite) TestImmediatePureTaskNowStableWithinTaskOnly() { mutableContext.AddTask( component, taskAttributes, - &TestPureTask{ - Payload: &commonpb.Payload{Data: []byte("task-1")}, - }, + &TestPureTask{}, ) mutableContext.AddTask( component, taskAttributes, - &TestPureTask{ - Payload: &commonpb.Payload{Data: []byte("task-2")}, - }, + &TestPureTask{}, ) s.testLibrary.mockPureTaskHandler.EXPECT(). From d75dd38a542c2c7034160f7af6fd7384638664bf Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Mon, 8 Jun 2026 17:23:28 -0400 Subject: [PATCH 3/3] Remove per component map --- chasm/context.go | 31 +++++++++++++++++-------------- chasm/tree_test.go | 8 +++++--- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/chasm/context.go b/chasm/context.go index bf1ca11c36e..27a2bdd6c34 100644 --- a/chasm/context.go +++ b/chasm/context.go @@ -99,7 +99,7 @@ type immutableCtx struct { type mutableCtx struct { *immutableCtx - nowByComponent map[*Node]time.Time + now **time.Time } // NewContext creates a new Context from an existing Context and root Node. @@ -210,27 +210,25 @@ func NewMutableContext( ctx context.Context, node *Node, ) MutableContext { + var now *time.Time return &mutableCtx{ - immutableCtx: newContext(ctx, node), - nowByComponent: make(map[*Node]time.Time), + immutableCtx: newContext(ctx, node), + now: &now, } } func (c *mutableCtx) Now(component Component) time.Time { - node, ok := c.root.valueToNode[component] - if !ok || !node.isComponent() { - return c.root.Now(component) + if c.now == nil { + var now *time.Time + c.now = &now } - if now, ok := c.nowByComponent[node]; ok { - return now + if *c.now != nil { + return **c.now } now := c.root.Now(component) - if c.nowByComponent == nil { - c.nowByComponent = make(map[*Node]time.Time) - } - c.nowByComponent[node] = now + *c.now = &now return now } @@ -243,9 +241,14 @@ func (c *mutableCtx) AddTask( } func (c *mutableCtx) withValue(key any, value any) Context { + if c.now == nil { + var now *time.Time + c.now = &now + } + return &mutableCtx{ - immutableCtx: ContextWithValue(c.immutableCtx, key, value), - nowByComponent: c.nowByComponent, + immutableCtx: ContextWithValue(c.immutableCtx, key, value), + now: c.now, } } diff --git a/chasm/tree_test.go b/chasm/tree_test.go index 67e7a434c39..68f25947ee1 100644 --- a/chasm/tree_test.go +++ b/chasm/tree_test.go @@ -3139,14 +3139,16 @@ func (s *nodeSuite) TestMutableContextNowStableWithinContext() { mutableContext := NewMutableContext(context.Background(), root) component, err := root.Component(mutableContext, ComponentRef{}) s.NoError(err) + testComponent := component.(*TestComponent) - s.Equal(startTime, mutableContext.Now(component)) + contextWithValue := ContextWithValue(mutableContext, "test-key", "test-value") + s.Equal(startTime, contextWithValue.Now(component)) s.timeSource.Update(updatedTime) s.Equal(startTime, mutableContext.Now(component)) - contextWithValue := ContextWithValue(mutableContext, "test-key", "test-value") - s.Equal(startTime, contextWithValue.Now(component)) + childComponent := testComponent.SubComponent1.Get(mutableContext) + s.Equal(startTime, mutableContext.Now(childComponent)) newMutableContext := NewMutableContext(context.Background(), root) s.Equal(updatedTime, newMutableContext.Now(component))