From f12e2ed15524e700973c8493aec265ab50842f55 Mon Sep 17 00:00:00 2001 From: WeiCHEN Date: Tue, 22 Sep 2026 21:24:20 +0800 Subject: [PATCH] fix: reclaim cancelFns, TTL idemLocks, cache scope probe, invalidate C1, SSE subscribe order Close T20-B2/B6/B8/B9/B10 low-risk leaks: map-backed cancel reclaim, idempotency lock TTL in cleanupLoop, sync.Once information_schema probe, InvalidateLayer after chat memory upsert, and Subscribe before SSE headers. --- internal/ctxmgr/chat_memory.go | 3 ++ internal/ctxmgr/manager.go | 4 +++ internal/ctxmgr/message_scope_io.go | 19 +++++++------ internal/events/worker.go | 12 ++++++-- internal/server/task_run_handlers.go | 16 ++++++----- internal/taskrun/manager.go | 42 ++++++++++++++++++++++++---- 6 files changed, 73 insertions(+), 23 deletions(-) diff --git a/internal/ctxmgr/chat_memory.go b/internal/ctxmgr/chat_memory.go index e75d70d8..b5d72a21 100644 --- a/internal/ctxmgr/chat_memory.go +++ b/internal/ctxmgr/chat_memory.go @@ -147,6 +147,8 @@ func (m *ContextManager) RecordTurnMemoryWithVisibility( if err != nil { return 0, fmt.Errorf("insert chat memory chunk to db failed: %w", err) } + // C1 search cache keys include query text; drop the layer so new chunks are visible. + m.Cache().InvalidateLayer(LayerC1Search) return chunkID, nil } @@ -184,6 +186,7 @@ func (m *ContextManager) RecordTurnMemoryWithVisibility( m.memoryChunks = append(m.memoryChunks, item) } + m.Cache().InvalidateLayer(LayerC1Search) return chunkID, nil } diff --git a/internal/ctxmgr/manager.go b/internal/ctxmgr/manager.go index 864ec7a3..8b2bfa1e 100644 --- a/internal/ctxmgr/manager.go +++ b/internal/ctxmgr/manager.go @@ -39,6 +39,10 @@ type ContextManager struct { nextActivityID int64 nextMessageID int64 nextChunkID int64 + + // message scope column probe (information_schema); cached once per manager + messageScopeOnce sync.Once + messageScopeOK bool } // NewContextManager 创建上下文管理器 diff --git a/internal/ctxmgr/message_scope_io.go b/internal/ctxmgr/message_scope_io.go index 568d1624..e824e023 100644 --- a/internal/ctxmgr/message_scope_io.go +++ b/internal/ctxmgr/message_scope_io.go @@ -25,14 +25,17 @@ func (m *ContextManager) hasMessageScopeColumns(ctx context.Context) bool { if m == nil || m.pool == nil { return false } - var ok bool - _ = m.pool.QueryRow(ctx, ` - SELECT EXISTS ( - SELECT 1 FROM information_schema.columns - WHERE table_schema = 'public' AND table_name = 'chat_messages' AND column_name = 'scope_kind' - ) - `).Scan(&ok) - return ok + m.messageScopeOnce.Do(func() { + var ok bool + _ = m.pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'chat_messages' AND column_name = 'scope_kind' + ) + `).Scan(&ok) + m.messageScopeOK = ok + }) + return m.messageScopeOK } func (m *ContextManager) chatMessageSelectColumns(ctx context.Context) string { diff --git a/internal/events/worker.go b/internal/events/worker.go index 56df8ad3..3b7a2fac 100644 --- a/internal/events/worker.go +++ b/internal/events/worker.go @@ -30,7 +30,7 @@ type Dispatcher struct { mu sync.RWMutex claiming bool inflight map[string]struct{} - cancelFns []context.CancelFunc + cancelFns map[string]context.CancelFunc // keyed by event_id:consumer; reclaimed after handle lastBeat atomic.Int64 } @@ -56,6 +56,7 @@ func NewDispatcher(pool *pgxpool.Pool, environment, owner string) *Dispatcher { Metrics: &Metrics{}, claiming: true, inflight: map[string]struct{}{}, + cancelFns: map[string]context.CancelFunc{}, } d.Touch() return d @@ -73,7 +74,10 @@ func (d *Dispatcher) Alive(maxAge time.Duration) bool { func (d *Dispatcher) StopClaim() { d.mu.Lock() d.claiming = false - cancels := append([]context.CancelFunc(nil), d.cancelFns...) + cancels := make([]context.CancelFunc, 0, len(d.cancelFns)) + for _, c := range d.cancelFns { + cancels = append(cancels, c) + } d.mu.Unlock() for _, c := range cancels { c() @@ -193,15 +197,17 @@ func (d *Dispatcher) shadowScan(ctx context.Context, consumer string) (int, erro func (d *Dispatcher) handle(ctx context.Context, lease Lease) { key := lease.Consumer + ":" + lease.AggregateType + ":" + lease.AggregateID + cancelKey := lease.EventID + ":" + lease.Consumer d.mu.Lock() d.inflight[key] = struct{}{} hctx, cancel := context.WithCancel(ctx) - d.cancelFns = append(d.cancelFns, cancel) + d.cancelFns[cancelKey] = cancel d.mu.Unlock() defer func() { cancel() d.mu.Lock() delete(d.inflight, key) + delete(d.cancelFns, cancelKey) d.mu.Unlock() }() diff --git a/internal/server/task_run_handlers.go b/internal/server/task_run_handlers.go index 435616c2..654a3acf 100644 --- a/internal/server/task_run_handlers.go +++ b/internal/server/task_run_handlers.go @@ -463,6 +463,15 @@ func (s *Server) handleTaskRunEvents(w http.ResponseWriter, r *http.Request) { afterSeq, _ = strconv.ParseInt(raw, 10, 64) } + // Subscribe before declaring text/event-stream so NotFound stays a normal JSON error + // (same order as handleChatStream / CreateAndStart in sse.go). + ch, replay, err := s.taskRuns.Subscribe(runID, afterSeq) + if err != nil { + writeSafeError(w, r, http.StatusNotFound, err.Error()) + return + } + defer s.taskRuns.Unsubscribe(runID, ch) + w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") @@ -471,13 +480,6 @@ func (s *Server) handleTaskRunEvents(w http.ResponseWriter, r *http.Request) { flusher.Flush() extendSSEWriteDeadline(w) - ch, replay, err := s.taskRuns.Subscribe(runID, afterSeq) - if err != nil { - writeSafeError(w, r, http.StatusNotFound, err.Error()) - return - } - defer s.taskRuns.Unsubscribe(runID, ch) - writeStoredEvents := func(events []taskrun.StoredEvent) error { for _, stored := range events { data, err := json.Marshal(stored.Event) diff --git a/internal/taskrun/manager.go b/internal/taskrun/manager.go index 58fa7392..ccf9c183 100644 --- a/internal/taskrun/manager.go +++ b/internal/taskrun/manager.go @@ -41,10 +41,16 @@ type Manager struct { ctxMgr *ctxmgr.ContextManager idemMu sync.Mutex - idemLocks map[string]*sync.Mutex // keyed by idempotency_key + idemLocks map[string]*idemLockEntry // keyed by idempotency_key; TTL-evicted in cleanupLoop +} + +type idemLockEntry struct { + mu sync.Mutex + lastUsed time.Time } const runRetention = time.Hour +const idemLockRetention = time.Hour var ( ErrRunNotResumable = errors.New("taskrun: run is not resumable") @@ -58,7 +64,7 @@ func NewManager(pwSvc *projectworkflow.Service, ctxMgr *ctxmgr.ContextManager) * workerID: newWorkerID(), pwSvc: pwSvc, ctxMgr: ctxMgr, - idemLocks: make(map[string]*sync.Mutex), + idemLocks: make(map[string]*idemLockEntry), } go m.cleanupLoop() return m @@ -71,12 +77,18 @@ func (m *Manager) lockIdempotencyKey(key string) func() { m.idemMu.Lock() lk, ok := m.idemLocks[key] if !ok { - lk = &sync.Mutex{} + lk = &idemLockEntry{} m.idemLocks[key] = lk } + lk.lastUsed = time.Now() m.idemMu.Unlock() - lk.Lock() - return lk.Unlock + lk.mu.Lock() + return func() { + lk.mu.Unlock() + m.idemMu.Lock() + lk.lastUsed = time.Now() + m.idemMu.Unlock() + } } // SetStore 注入持久化存储（PG）；启动时标记 lease 过期的 running 任务为 interrupted。 @@ -102,6 +114,7 @@ func (m *Manager) cleanupLoop() { defer ticker.Stop() for range ticker.C { m.evictStaleRuns() + m.evictStaleIdemLocks() } } @@ -120,6 +133,25 @@ func (m *Manager) evictStaleRuns() { } } +func (m *Manager) evictStaleIdemLocks() { + cutoff := time.Now().Add(-idemLockRetention) + m.idemMu.Lock() + defer m.idemMu.Unlock() + for key, lk := range m.idemLocks { + if lk.lastUsed.After(cutoff) { + continue + } + // Skip locks currently held; TryLock avoids blocking cleanup on a live claim. + if !lk.mu.TryLock() { + continue + } + if !lk.lastUsed.After(cutoff) { + delete(m.idemLocks, key) + } + lk.mu.Unlock() + } +} + func (m *Manager) SetExecutor(exec Executor, resume ResumeExecutor, stage StageExecutor) { m.exec = exec m.resume = resume