diff --git a/workflow/checkpoint/manager.go b/workflow/checkpoint/manager.go index b43bcc6d..dd866c9d 100644 --- a/workflow/checkpoint/manager.go +++ b/workflow/checkpoint/manager.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "slices" + "sync" "github.com/microsoft/agent-framework-go/workflow" "github.com/microsoft/agent-framework-go/workflow/internal/checkpoint" @@ -33,6 +34,10 @@ func NewJSONManager(store Store[json.RawMessage]) Manager { } type inMemoryManager struct { + // mu guards Store (and the per-session caches within it) so a single + // manager can be shared across concurrent workflow runs — which its + // session-keyed design implies — without racing on the map or the caches. + mu sync.Mutex Store map[string]*checkpoint.SessionCache[*checkpoint.Checkpoint] } @@ -61,6 +66,8 @@ func (s *inMemoryManager) Commit(_ context.Context, sessionID string, checkpoint return workflow.CheckpointInfo{}, fmt.Errorf("checkpoint: parent sessionID %q does not match sessionID %q", checkpoint.Parent.SessionID, sessionID) } + s.mu.Lock() + defer s.mu.Unlock() store := s.sessionStore(sessionID) return store.Add(sessionID, checkpoint), nil } @@ -73,6 +80,8 @@ func (s *inMemoryManager) Lookup(_ context.Context, sessionID string, checkpoint return nil, fmt.Errorf("checkpoint: checkpoint sessionID %q does not match sessionID %q", checkpointInfo.SessionID, sessionID) } + s.mu.Lock() + defer s.mu.Unlock() store := s.sessionStore(sessionID) v, ok := store.Get(checkpointInfo) if !ok { @@ -89,6 +98,8 @@ func (s *inMemoryManager) RetrieveIndex(_ context.Context, sessionID string, wit return nil, fmt.Errorf("checkpoint: parent sessionID %q does not match sessionID %q", withParent.SessionID, sessionID) } + s.mu.Lock() + defer s.mu.Unlock() store := s.sessionStore(sessionID) if withParent == nil { return slices.Clone(store.CheckpointIndex), nil diff --git a/workflow/checkpoint/manager_race_internal_test.go b/workflow/checkpoint/manager_race_internal_test.go new file mode 100644 index 00000000..859294e8 --- /dev/null +++ b/workflow/checkpoint/manager_race_internal_test.go @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft. All rights reserved. + +package checkpoint + +import ( + "context" + "fmt" + "sync" + "testing" + + icheckpoint "github.com/microsoft/agent-framework-go/workflow/internal/checkpoint" +) + +// The in-memory manager is keyed by session ID, so a single manager can be +// shared across concurrent workflow runs (distinct sessions). Its Store map and +// per-session caches must be synchronized: without a lock, concurrent Commit +// calls race on the map and can crash with "fatal: concurrent map writes". +func TestInMemoryManager_ConcurrentSessions_NoRace(t *testing.T) { + mgr := &inMemoryManager{} + ctx := context.Background() + + const n = 64 + var wg sync.WaitGroup + wg.Add(n) + for i := range n { + go func(i int) { + defer wg.Done() + sid := fmt.Sprintf("session-%d", i) + if _, err := mgr.Commit(ctx, sid, &icheckpoint.Checkpoint{}); err != nil { + t.Errorf("Commit(%s): %v", sid, err) + } + if _, err := mgr.RetrieveIndex(ctx, sid, nil); err != nil { + t.Errorf("RetrieveIndex(%s): %v", sid, err) + } + }(i) + } + wg.Wait() +}