refactor(store): fix goroutine safety issues across cron, sessions, skills, and gateway

- Replace 16x bare db.Query/db.Exec with context-aware variants to prevent
  orphan DB operations when parent context is cancelled
- Thread baseCtx through PGCronStore lifecycle (Start/Stop) for all internal
  scheduler, executor, and job cache operations
- Add ctx parameter to StoreMissingDeps and ListSystemSkillDirs, update all callers
- Add sync.WaitGroup tracking to fire-and-forget goroutines in gateway consumer
  handlers (subagent announce, teammate message) with safego.Recover panic protection
- Add safego.Recover to 3 Slack channel goroutines to prevent WaitGroup hang on panic
This commit is contained in:
viettranx committed 2026-03-25 08:20:10 +07:00
1 parent d63a7d4ced
commit 81c8414423
15 files changed
+62 -34

No files matched your search

+10 -2
View File
@@ -79,6 +79,14 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
slog.Info("inbound debounce configured", "debounce_ms", debounceMs)
// Track background goroutines (subagent announces, teammate messages)
// so shutdown can wait for in-flight work to complete.
var bgWg sync.WaitGroup
defer func() {
bgWg.Wait()
slog.Info("inbound consumer: all background goroutines drained")
}()
for {
msg, ok := msgBus.ConsumeInbound(ctx)
if !ok {
@@ -95,10 +103,10 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
}
}
if handleSubagentAnnounce(ctx, msg, cfg, sched, channelMgr, msgBus, getAnnounceMu) {
if handleSubagentAnnounce(ctx, msg, cfg, sched, channelMgr, msgBus, getAnnounceMu, &bgWg) {
continue
}
if handleTeammateMessage(ctx, msg, cfg, sched, channelMgr, teamStore, agentStore, msgBus, postTurn, &taskRunSessions) {
if handleTeammateMessage(ctx, msg, cfg, sched, channelMgr, teamStore, agentStore, msgBus, postTurn, &taskRunSessions, &bgWg) {
continue
}
if handleResetCommand(msg, cfg, sessStore) {
+9
View File
@@ -17,6 +17,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/channels"
"github.com/nextlevelbuilder/goclaw/internal/config"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/safego"
"github.com/nextlevelbuilder/goclaw/internal/scheduler"
"github.com/nextlevelbuilder/goclaw/internal/sessions"
"github.com/nextlevelbuilder/goclaw/internal/store"
@@ -35,6 +36,7 @@ func handleSubagentAnnounce(
channelMgr *channels.Manager,
msgBus *bus.MessageBus,
getAnnounceMu func(string) *sync.Mutex,
bgWg *sync.WaitGroup,
) bool {
if !(msg.Channel == tools.ChannelSystem && strings.HasPrefix(msg.SenderID, "subagent:")) {
return false
@@ -128,7 +130,10 @@ func handleSubagentAnnounce(
// Handle announce asynchronously with per-session serialization.
// The mutex ensures concurrent announces for the same session wait for
// each other, so each reads up-to-date session history.
bgWg.Add(1)
go func(sessionKey, origCh, chatID, senderID, label string, meta map[string]string, req agent.RunRequest) {
defer bgWg.Done()
defer safego.Recover(nil, "component", "subagent_announce", "session", sessionKey)
mu := getAnnounceMu(sessionKey)
mu.Lock()
defer mu.Unlock()
@@ -193,6 +198,7 @@ func handleTeammateMessage(
msgBus *bus.MessageBus,
postTurn tools.PostTurnProcessor,
taskRunSessions *sync.Map,
bgWg *sync.WaitGroup,
) bool {
if !(msg.Channel == tools.ChannelSystem && strings.HasPrefix(msg.SenderID, "teammate:")) {
return false
@@ -286,7 +292,10 @@ func handleTeammateMessage(
LinkedTraceID: linkedTraceID,
})
bgWg.Add(1)
go func(origCh, origChatID, senderID, taskID string, meta, inMeta map[string]string) {
defer bgWg.Done()
defer safego.Recover(nil, "component", "teammate_message", "task_id", taskID)
// Lock renewal heartbeat: extend task lock every 10 min to prevent
// the ticker from recovering long-running tasks as stale.
var lockStop chan struct{}
+4
View File
@@ -15,6 +15,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/bus"
"github.com/nextlevelbuilder/goclaw/internal/channels"
"github.com/nextlevelbuilder/goclaw/internal/config"
"github.com/nextlevelbuilder/goclaw/internal/safego"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
@@ -167,12 +168,14 @@ func (c *Channel) Start(ctx context.Context) error {
// Goroutine 1: Event loop
go func() {
defer c.wg.Done()
defer safego.Recover(nil, "component", "slack_event_loop")
c.eventLoop(smCtx)
}()
// Goroutine 2: Socket Mode connection with dead socket error classification
go func() {
defer c.wg.Done()
defer safego.Recover(nil, "component", "slack_socket_mode")
for {
if err := c.sm.RunContext(smCtx); err != nil {
if smCtx.Err() != nil {
@@ -192,6 +195,7 @@ func (c *Channel) Start(ctx context.Context) error {
// Goroutine 3: Periodic sweep (every 2 minutes) for TTL-based map eviction
go func() {
defer c.wg.Done()
defer safego.Recover(nil, "component", "slack_sweep")
ticker := time.NewTicker(2 * time.Minute)
defer ticker.Stop()
for {
+4 -4
View File
@@ -203,7 +203,7 @@ func (h *SkillsHandler) handleInstallDeps(w http.ResponseWriter, r *http.Request
if !h.requireMasterTenant(w, r) {
return
}
dirs := h.skills.ListSystemSkillDirs()
dirs := h.skills.ListSystemSkillDirs(r.Context())
if len(dirs) == 0 {
writeJSON(w, http.StatusOK, map[string]string{"message": "no system skills"})
return
@@ -342,7 +342,7 @@ func (h *SkillsHandler) rescanAndUpdate() (updated int, results []depResult) {
if manifest == nil || manifest.IsEmpty() {
// No deps needed — if archived, recover to active and clear stale deps.
if sk.Status == "archived" {
_ = h.skills.StoreMissingDeps(id, nil)
_ = h.skills.StoreMissingDeps(masterCtx, id, nil)
_ = h.skills.UpdateSkill(masterCtx, id, map[string]any{"status": "active"})
results = append(results, depResult{Slug: sk.Slug, Status: "active"})
updated++
@@ -354,7 +354,7 @@ func (h *SkillsHandler) rescanAndUpdate() (updated int, results []depResult) {
}
ok, missing := skills.CheckSkillDeps(manifest)
_ = h.skills.StoreMissingDeps(id, missing)
_ = h.skills.StoreMissingDeps(masterCtx, id, missing)
switch {
case ok && sk.Status == "archived":
@@ -471,7 +471,7 @@ func (h *SkillsHandler) handleToggle(w http.ResponseWriter, r *http.Request) {
manifest := h.scanWithFallback(sk)
if manifest != nil && !manifest.IsEmpty() {
depOk, missing := skills.CheckSkillDeps(manifest)
_ = h.skills.StoreMissingDeps(id, missing)
_ = h.skills.StoreMissingDeps(r.Context(), id, missing)
if depOk {
newStatus = "active"
} else {
+2 -2
View File
@@ -22,7 +22,7 @@ type SystemSkillStore interface {
GetNextVersion(slug string) int
BumpVersion()
UpdateSkill(ctx context.Context, id uuid.UUID, updates map[string]interface{}) error
StoreMissingDeps(id uuid.UUID, missing []string) error
StoreMissingDeps(ctx context.Context, id uuid.UUID, missing []string) error
}
// seededSkill tracks a skill that was seeded and needs async dep checking.
@@ -173,7 +173,7 @@ func (s *Seeder) CheckDepsAsync(skills []seededSkill, msgBus *bus.MessageBus) {
ok, missing := CheckSkillDeps(manifest)
// Always persist missing deps so UI can display them per-skill
_ = s.store.StoreMissingDeps(sk.id, missing)
_ = s.store.StoreMissingDeps(context.Background(), sk.id, missing)
status := "active"
if !ok {
status = "archived"
+13 -6
View File
@@ -1,6 +1,7 @@
package pg
import (
"context"
"database/sql"
"log/slog"
"sync"
@@ -15,12 +16,14 @@ const defaultCronCacheTTL = 2 * time.Minute
// PGCronStore implements store.CronStore backed by Postgres.
// GetDueJobs() uses an in-memory cache with TTL to reduce DB polling (1s interval).
type PGCronStore struct {
db *sql.DB
mu sync.Mutex
onJob func(job *store.CronJob) (*store.CronJobResult, error)
onEvent func(event store.CronEvent)
running bool
stop chan struct{}
db *sql.DB
mu sync.Mutex
baseCtx context.Context // lifecycle context for background DB operations
cancelCtx context.CancelFunc // cancelled on Stop()
onJob func(job *store.CronJob) (*store.CronJobResult, error)
onEvent func(event store.CronEvent)
running bool
stop chan struct{}
// Job cache: reduces GetDueJobs polling from 86,400 queries/day to ~720/day
jobCache []store.CronJob
@@ -57,6 +60,7 @@ func (s *PGCronStore) Start() error {
if s.running {
return nil
}
s.baseCtx, s.cancelCtx = context.WithCancel(context.Background())
s.stop = make(chan struct{})
s.running = true
s.recomputeStaleJobs()
@@ -72,6 +76,9 @@ func (s *PGCronStore) Stop() {
return
}
close(s.stop)
if s.cancelCtx != nil {
s.cancelCtx()
}
s.running = false
}
+2 -2
View File
@@ -66,7 +66,7 @@ func (s *PGCronStore) AddJob(ctx context.Context, name string, schedule store.Cr
nextRun := computeNextRun(&schedule, now, s.defaultTZ)
_, err := s.db.Exec(
_, err := s.db.ExecContext(ctx,
`INSERT INTO cron_jobs (id, tenant_id, agent_id, user_id, name, enabled, schedule_kind, cron_expression, run_at, timezone,
interval_ms, payload, delete_after_run, next_run_at, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, true, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)`,
@@ -133,7 +133,7 @@ func (s *PGCronStore) ListJobs(ctx context.Context, includeDisabled bool, agentI
q += " ORDER BY created_at DESC"
rows, err := s.db.Query(q, args...)
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil
}
+1 -1
View File
@@ -27,7 +27,7 @@ func (s *PGCronStore) RunJob(ctx context.Context, jobID string, force bool) (boo
// Mark job as running before execution
if id, parseErr := uuid.Parse(jobID); parseErr == nil {
s.db.Exec("UPDATE cron_jobs SET last_status = 'running', updated_at = $1 WHERE id = $2", time.Now(), id)
s.db.ExecContext(ctx, "UPDATE cron_jobs SET last_status = 'running', updated_at = $1 WHERE id = $2", time.Now(), id)
}
s.mu.Lock()
s.cacheLoaded = false
+8 -8
View File
@@ -33,7 +33,7 @@ func (s *PGCronStore) GetDueJobs(now time.Time) []store.CronJob {
// refreshJobCache reloads all enabled jobs from DB. Must be called with mu held.
func (s *PGCronStore) refreshJobCache() {
rows, err := s.db.Query(
rows, err := s.db.QueryContext(s.baseCtx,
`SELECT id, tenant_id, agent_id, user_id, name, enabled, schedule_kind, cron_expression, run_at, timezone,
interval_ms, payload, delete_after_run, next_run_at, last_run_at, last_status, last_error,
created_at, updated_at FROM cron_jobs WHERE enabled = true`)
@@ -65,7 +65,7 @@ func (s *PGCronStore) InvalidateCache() {
// This happens when the gateway was stopped/crashed while a job was executing,
// or when the previously computed next_run_at was consumed but never recomputed.
func (s *PGCronStore) recomputeStaleJobs() {
rows, err := s.db.Query(
rows, err := s.db.QueryContext(s.baseCtx,
`SELECT id, schedule_kind, cron_expression, run_at, timezone, interval_ms
FROM cron_jobs WHERE enabled = true AND next_run_at IS NULL`)
if err != nil {
@@ -105,12 +105,12 @@ func (s *PGCronStore) recomputeStaleJobs() {
next := computeNextRun(&schedule, now, s.defaultTZ)
if next == nil {
if scheduleKind == "at" {
s.db.Exec("UPDATE cron_jobs SET enabled = false, updated_at = $1 WHERE id = $2", now, id)
s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET enabled = false, updated_at = $1 WHERE id = $2", now, id)
}
continue
}
s.db.Exec("UPDATE cron_jobs SET next_run_at = $1, updated_at = $2 WHERE id = $3", *next, now, id)
s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET next_run_at = $1, updated_at = $2 WHERE id = $3", *next, now, id)
fixed++
}
@@ -149,7 +149,7 @@ func (s *PGCronStore) checkAndRunDueJobs() {
// Clear next_run for all due jobs first to prevent duplicate fires
for _, job := range dueJobs {
if id, parseErr := uuid.Parse(job.ID); parseErr == nil {
s.db.Exec("UPDATE cron_jobs SET next_run_at = NULL WHERE id = $1", id)
s.db.ExecContext(s.baseCtx, "UPDATE cron_jobs SET next_run_at = NULL WHERE id = $1", id)
}
}
@@ -222,7 +222,7 @@ func (s *PGCronStore) executeOneJob(job store.CronJob, handler func(job *store.C
if aid, aidErr := uuid.Parse(job.AgentID); aidErr == nil {
agentUUID = &aid
}
s.db.Exec(
s.db.ExecContext(s.baseCtx,
`INSERT INTO cron_run_logs (id, job_id, agent_id, status, error, summary, duration_ms, input_tokens, output_tokens, ran_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
logID, id, agentUUID, status, lastError, summary, durationMS, inputTokens, outputTokens, now,
@@ -232,12 +232,12 @@ func (s *PGCronStore) executeOneJob(job store.CronJob, handler func(job *store.C
// Recompute next run or delete
if job.DeleteAfterRun {
if id, parseErr := uuid.Parse(job.ID); parseErr == nil {
s.db.Exec("DELETE FROM cron_jobs WHERE id = $1", id)
s.db.ExecContext(s.baseCtx, "DELETE FROM cron_jobs WHERE id = $1", id)
}
} else if id, parseErr := uuid.Parse(job.ID); parseErr == nil {
schedule := job.Schedule
next := computeNextRun(&schedule, now, s.defaultTZ)
s.db.Exec(
s.db.ExecContext(s.baseCtx,
"UPDATE cron_jobs SET last_run_at = $1, last_status = $2, last_error = $3, next_run_at = $4, updated_at = $5 WHERE id = $6",
now, status, lastError, next, now, id,
)
+2 -2
View File
@@ -41,7 +41,7 @@ func NewPGSessionStore(db *sql.DB) *PGSessionStore {
// The last hyphen-delimited segment is the base36 timestamp used as convId.
// Idempotent — no-op if no legacy keys exist.
func (s *PGSessionStore) migrateLegacyWSKeys() {
res, err := s.db.Exec(`
res, err := s.db.ExecContext(context.Background(), `
UPDATE sessions
SET session_key = regexp_replace(
session_key,
@@ -103,7 +103,7 @@ func (s *PGSessionStore) GetOrCreate(ctx context.Context, key string) *store.Ses
s.cache[sessionCacheKey(ctx, key)] = data
msgsJSON, _ := json.Marshal([]providers.Message{})
s.db.Exec(
s.db.ExecContext(ctx,
`INSERT INTO sessions (id, session_key, messages, created_at, updated_at, team_id, tenant_id)
VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (tenant_id, session_key) DO NOTHING`,
uuid.Must(uuid.NewV7()), key, msgsJSON, now, now, teamID, tenantIDForInsert(ctx),
+1 -1
View File
@@ -352,7 +352,7 @@ func (s *PGSessionStore) getOrInit(ctx context.Context, key string) *store.Sessi
s.cache[sessionCacheKey(ctx, key)] = data
msgsJSON, _ := json.Marshal([]providers.Message{})
s.db.Exec(
s.db.ExecContext(ctx,
`INSERT INTO sessions (id, session_key, messages, created_at, updated_at, tenant_id)
VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (tenant_id, session_key) DO NOTHING`,
uuid.Must(uuid.NewV7()), key, msgsJSON, now, now, tenantIDForInsert(ctx),
+2 -2
View File
@@ -188,7 +188,7 @@ func (s *PGSkillStore) scanSkillInfoList(rows *sql.Rows) []store.SkillInfo {
// StoreMissingDeps persists the missing_deps list for a skill into the deps JSONB column.
// Only updates system skills unscoped; custom skills require tenant match.
func (s *PGSkillStore) StoreMissingDeps(id uuid.UUID, missing []string) error {
func (s *PGSkillStore) StoreMissingDeps(ctx context.Context, id uuid.UUID, missing []string) error {
if missing == nil {
missing = []string{}
}
@@ -198,7 +198,7 @@ func (s *PGSkillStore) StoreMissingDeps(id uuid.UUID, missing []string) error {
}
// System skills can be updated without tenant; custom skills need tenant scope.
// Use is_system check to ensure cross-tenant safety for custom skills.
_, err = s.db.Exec(
_, err = s.db.ExecContext(ctx,
`UPDATE skills SET deps = $1, updated_at = NOW() WHERE id = $2 AND is_system = true`,
encoded, id,
)
+2 -2
View File
@@ -77,8 +77,8 @@ func (s *PGSkillStore) UpsertSystemSkill(ctx context.Context, p SkillCreateParam
// ListSystemSkillDirs returns slug->file_path map for all enabled system skills.
// Disabled system skills are excluded — dep checking and injection are skipped for them.
func (s *PGSkillStore) ListSystemSkillDirs() map[string]string {
rows, err := s.db.Query(
func (s *PGSkillStore) ListSystemSkillDirs(ctx context.Context) map[string]string {
rows, err := s.db.QueryContext(ctx,
`SELECT slug, file_path FROM skills WHERE is_system = true AND enabled = true`)
if err != nil {
return nil
+1 -1
View File
@@ -175,7 +175,7 @@ func (t *PublishSkillTool) Execute(ctx context.Context, args map[string]any) *Re
if manifest != nil && !manifest.IsEmpty() {
ok, missing := skills.CheckSkillDeps(manifest)
if !ok {
_ = t.skills.StoreMissingDeps(id, missing)
_ = t.skills.StoreMissingDeps(ctx, id, missing)
depsWarning = skills.FormatMissing(missing)
}
}
+1 -1
View File
@@ -198,7 +198,7 @@ func (t *SkillManageTool) executeCreate(ctx context.Context, args map[string]any
if manifest != nil && !manifest.IsEmpty() {
ok, missing := skills.CheckSkillDeps(manifest)
if !ok {
_ = t.skills.StoreMissingDeps(id, missing)
_ = t.skills.StoreMissingDeps(ctx, id, missing)
depsWarning = skills.FormatMissing(missing)
}
}