mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-15 02:16:56 +00:00
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:
1 parent
d63a7d4ced
commit
81c8414423
15 files changed
+62
-34
No files matched your search
+10
-2
@@ -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) {
|
||||
|
||||
@@ -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{}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user