perf(storage): allow concurrent reads; serialize writes via a mutex

SetMaxOpenConns(1) forced every read (List/Get/GetPR/...) to queue behind
the single connection, so the dashboard's reads contended with the LCM's
writes. WAL already supports many concurrent readers, so raise the pool to 8
and instead serialize *writes* with a Store.writeMu. That keeps WAL's
single-writer rule and the revision-CAS read-then-write atomic regardless of
pool size, while reads now run in parallel across the pool.

Every write method takes writeMu (Upsert, PatchMetadata, UpsertPR/DeletePR,
the pr_check/pr_comment Replace* via inTx, the CDC outbox/offset writes,
project writes, reaction-tracker writes); reads take nothing. Added
TestConcurrentReadsAndWrites (16 writers + 16 readers) which passes under
-race.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
prateek 2026-05-31 00:51:18 +05:30
parent b0e4fffa62
commit ba47212802
8 changed files with 109 additions and 14 deletions

View File

@ -45,6 +45,8 @@ func (s *Store) ListUnsent(ctx context.Context, limit int) ([]OutboxEvent, error
// MarkSent flags an outbox row delivered.
func (s *Store) MarkSent(ctx context.Context, outboxID int64, at time.Time) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.MarkOutboxSent(ctx, gen.MarkOutboxSentParams{
SentAt: sql.NullTime{Time: at, Valid: true},
ID: outboxID,
@ -53,6 +55,8 @@ func (s *Store) MarkSent(ctx context.Context, outboxID int64, at time.Time) erro
// MarkFailed bumps the attempt count and records the last error for an outbox row.
func (s *Store) MarkFailed(ctx context.Context, outboxID int64, errMsg string) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.MarkOutboxFailed(ctx, gen.MarkOutboxFailedParams{LastError: errMsg, ID: outboxID})
}
@ -70,6 +74,8 @@ func (s *Store) GetOffset(ctx context.Context, consumer string) (int64, error) {
// SetOffset durably records a consumer's acknowledged seq.
func (s *Store) SetOffset(ctx context.Context, consumer string, seq int64, at time.Time) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.UpsertConsumerOffset(ctx, gen.UpsertConsumerOffsetParams{
Consumer: consumer,
LastSeq: seq,
@ -100,5 +106,7 @@ func (s *Store) MinConsumerOffset(ctx context.Context) (int64, error) {
// DeleteSentOutboxBelow removes delivered outbox rows whose seq is below the
// watermark, returning the number removed.
func (s *Store) DeleteSentOutboxBelow(ctx context.Context, seq int64) (int64, error) {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.DeleteSentOutboxBelow(ctx, seq)
}

View File

@ -18,17 +18,25 @@ import (
//go:embed migrations/*.sql
var migrationsFS embed.FS
// pragmas are applied on every connection open. WAL + NORMAL gives concurrent
// reads alongside the single writer; busy_timeout absorbs brief writer
// contention; foreign_keys enforces the session_metadata cascade.
// pragmas are applied on every connection open. WAL + NORMAL lets readers run
// concurrently with the writer; busy_timeout absorbs brief writer contention;
// foreign_keys enforces the cascades.
const pragmas = "?_pragma=journal_mode(WAL)" +
"&_pragma=busy_timeout(5000)" +
"&_pragma=foreign_keys(ON)" +
"&_pragma=synchronous(NORMAL)"
// maxConnections caps the pool. WAL allows many concurrent readers, so reads
// (List/Get/GetPR/...) scale across the pool instead of queuing behind a single
// connection. Writes do NOT rely on the pool for serialization — the Store funnels
// every write through its writeMu (see store.go), which keeps WAL's single-writer
// rule and the revision-CAS read-then-write atomic regardless of pool size.
const maxConnections = 8
// Open opens (creating if absent) the SQLite database under dataDir, applies the
// connection pragmas, and runs all goose migrations up. The returned *sql.DB is
// safe for the single-writer / many-reader workload the LCM and readers impose.
// sized for the many-reader / serialized-single-writer workload the LCM and
// readers impose.
func Open(dataDir string) (*sql.DB, error) {
if err := os.MkdirAll(dataDir, 0o755); err != nil {
return nil, fmt.Errorf("create data dir: %w", err)
@ -38,10 +46,8 @@ func Open(dataDir string) (*sql.DB, error) {
if err != nil {
return nil, fmt.Errorf("open sqlite: %w", err)
}
// Single writer: serialize all access through one connection so WAL's
// single-writer rule is never violated by the pool handing out a second
// writable conn mid-transaction.
db.SetMaxOpenConns(1)
db.SetMaxOpenConns(maxConnections)
db.SetMaxIdleConns(maxConnections) // keep reader conns warm; avoid open/close churn
if err := migrate(db); err != nil {
db.Close()

View File

@ -47,6 +47,8 @@ type PRComment struct {
// UpsertPR inserts or replaces the scalar PR facts for one session.
func (s *Store) UpsertPR(ctx context.Context, r PRRow) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.UpsertPR(ctx, gen.UpsertPRParams{
SessionID: r.SessionID,
ReviewDecision: r.ReviewDecision,
@ -87,6 +89,8 @@ func (s *Store) GetPR(ctx context.Context, sessionID string) (PRRow, bool, error
// comments. Normally unnecessary (the chain cascades on session delete); exposed
// for explicit eviction.
func (s *Store) DeletePR(ctx context.Context, sessionID string) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.DeletePR(ctx, sessionID)
}
@ -172,9 +176,12 @@ func (s *Store) ListPRComments(ctx context.Context, sessionID string) ([]PRComme
return out, nil
}
// inTx runs fn inside a single transaction over the store's queries, rolling
// back on error.
// inTx runs fn inside a single write transaction over the store's queries,
// rolling back on error. It holds writeMu for the duration, so callers must not
// already hold it.
func (s *Store) inTx(ctx context.Context, what string, fn func(*gen.Queries) error) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
tx, err := s.db.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin %s: %w", what, err)

View File

@ -34,6 +34,8 @@ type ProjectRow struct {
// UpsertProject inserts or updates one registered project.
func (s *Store) UpsertProject(ctx context.Context, r ProjectRow) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.UpsertProject(ctx, gen.UpsertProjectParams{
ID: r.ID,
Path: r.Path,
@ -53,6 +55,8 @@ func (s *Store) UpsertProject(ctx context.Context, r ProjectRow) error {
// ArchiveProject soft-deletes one project, keeping the row so a session's
// project_id still resolves. Active-only reads (ListProjects) then hide it.
func (s *Store) ArchiveProject(ctx context.Context, id string, t time.Time) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.ArchiveProject(ctx, gen.ArchiveProjectParams{
ArchivedAt: nullTime(t),
ID: id,
@ -86,6 +90,8 @@ func (s *Store) ListProjects(ctx context.Context) ([]ProjectRow, error) {
// DeleteProject removes one project by id.
func (s *Store) DeleteProject(ctx context.Context, id string) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.DeleteProject(ctx, id)
}

View File

@ -48,6 +48,8 @@ func (s *Store) ListReactionTrackers(ctx context.Context) ([]ReactionTrackerRow,
// SaveReactionTracker durably persists one escalation budget (insert or update).
func (s *Store) SaveReactionTracker(ctx context.Context, r ReactionTrackerRow) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
escalated := int64(0)
if r.Escalated {
escalated = 1
@ -68,6 +70,8 @@ func (s *Store) SaveReactionTracker(ctx context.Context, r ReactionTrackerRow) e
// DeleteReactionTracker drops one escalation budget.
func (s *Store) DeleteReactionTracker(ctx context.Context, sessionID, reactionKey string) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.DeleteReactionTracker(ctx, gen.DeleteReactionTrackerParams{
SessionID: sessionID,
ReactionKey: reactionKey,
@ -76,5 +80,7 @@ func (s *Store) DeleteReactionTracker(ctx context.Context, sessionID, reactionKe
// DeleteSessionReactionTrackers drops every escalation budget for a session.
func (s *Store) DeleteSessionReactionTrackers(ctx context.Context, sessionID string) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.DeleteSessionReactionTrackers(ctx, sessionID)
}

View File

@ -5,6 +5,7 @@ import (
"database/sql"
"errors"
"fmt"
"sync"
"time"
"github.com/aoagents/agent-orchestrator/backend/internal/domain"
@ -12,11 +13,17 @@ import (
"github.com/aoagents/agent-orchestrator/backend/internal/storage/sqlite/gen"
)
// Store is the SQLite-backed ports.LifecycleStore. The LCM is its sole logical
// writer (via Upsert); readers (Session Manager, reaper) use Load/Get/List.
// Store is the SQLite-backed ports.LifecycleStore. Reads (Load/Get/List/...) run
// concurrently across the connection pool; every write is funnelled through
// writeMu so there is exactly one writer at a time. That single-writer guarantee
// is load-bearing: it keeps WAL's single-writer rule and makes the revision-CAS
// (read-then-write in Upsert) atomic without depending on the pool size. Hold
// writeMu only around writes — never around a read — and never call one
// write method from inside another (the mutex is not reentrant).
type Store struct {
db *sql.DB
q *gen.Queries
db *sql.DB
q *gen.Queries
writeMu sync.Mutex
}
var _ ports.LifecycleStore = (*Store)(nil)
@ -105,6 +112,8 @@ func (s *Store) PatchMetadata(ctx context.Context, id domain.SessionID, meta dom
if meta.IsZero() {
return nil
}
s.writeMu.Lock()
defer s.writeMu.Unlock()
return s.q.UpsertSessionMetadata(ctx, gen.UpsertSessionMetadataParams{
SessionID: string(id),
Branch: meta.Branch,

View File

@ -2,7 +2,9 @@ package sqlite
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"time"
@ -255,3 +257,52 @@ func assertOutboxCount(t *testing.T, s *Store, ctx context.Context, want int) {
t.Fatalf("outbox count = %d, want %d", len(rows), want)
}
}
// TestConcurrentReadsAndWrites exercises the read-pool + write-mutex model:
// many writers (each its own session) run alongside many readers hammering
// ListAll. Reads must not be serialized behind writes, writes must not corrupt
// or error under the revision-CAS, and the final state must be exact. Run under
// -race this also guards the writeMu discipline.
func TestConcurrentReadsAndWrites(t *testing.T) {
s := newTestStore(t)
ctx := context.Background()
const n = 16
var wg sync.WaitGroup
errc := make(chan error, n*2)
for i := 0; i < n; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
if err := s.Upsert(ctx, sampleRecord(fmt.Sprintf("s%02d", i)), ports.EventSessionCreated); err != nil {
errc <- err
}
}(i)
}
for i := 0; i < n; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 25; j++ {
if _, err := s.ListAll(ctx); err != nil {
errc <- err
return
}
}
}()
}
wg.Wait()
close(errc)
for err := range errc {
t.Fatalf("concurrent op error: %v", err)
}
got, err := s.ListAll(ctx)
if err != nil {
t.Fatal(err)
}
if len(got) != n {
t.Fatalf("after %d concurrent inserts, ListAll returned %d", n, len(got))
}
}

View File

@ -24,6 +24,8 @@ import (
// stored+1.
// - insert: rec.Lifecycle.Revision must be 0, persisted as 1.
func (s *Store) Upsert(ctx context.Context, rec domain.SessionRecord, eventType ports.EventType) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
tx, err := s.db.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin upsert: %w", err)