diff --git a/internal/db/retry.go b/internal/db/retry.go new file mode 100644 index 0000000000..bf0bf10f7a --- /dev/null +++ b/internal/db/retry.go @@ -0,0 +1,86 @@ +package db + +import ( + "context" + "errors" + "math/rand/v2" + "time" + + "github.com/go-sql-driver/mysql" +) + +// MySQL server error numbers that InnoDB expects the caller to recover from by +// restarting the transaction (the 1213 message literally says "try restarting +// transaction"). +const ( + mysqlErrLockDeadlock = 1213 // ER_LOCK_DEADLOCK + mysqlErrLockWaitTimeout = 1205 // ER_LOCK_WAIT_TIMEOUT +) + +// Transaction-retry policy. A deadlock clears the instant InnoDB rolls the +// victim back, so a few quick jittered attempts absorb the contention without a +// human-perceptible stall. +const ( + txRetryMaxAttempts = 5 + txRetryBaseDelay = 2 * time.Millisecond + txRetryMaxDelay = 100 * time.Millisecond +) + +// IsRetryableTxError reports whether err is a transient transaction conflict +// worth retrying: a MySQL deadlock (1213) or lock-wait timeout (1205). It +// unwraps err, so wrapped errors are detected. Any non-MySQL error — including +// every SQLite error — returns false. +func IsRetryableTxError(err error) bool { + var myErr *mysql.MySQLError + if errors.As(err, &myErr) { + return myErr.Number == mysqlErrLockDeadlock || myErr.Number == mysqlErrLockWaitTimeout + } + return false +} + +// WithTxRetry runs fn, retrying while it returns a retryable transaction +// conflict (see IsRetryableTxError), up to txRetryMaxAttempts times with +// full-jitter exponential backoff between attempts. +// +// fn MUST be self-contained and safe to run more than once: it opens its own +// transaction and rolls back on any error, because a retry re-runs it from +// scratch. On the message store's concurrent SELECT MAX(seq) / INSERT path this +// turns an otherwise fatal InnoDB deadlock — two writes to the same session +// racing for the session_id index range — into a transparent retry, matching +// MySQL's own "try restarting transaction" guidance. +// +// A non-retryable error is returned immediately. If ctx is cancelled, the +// context error is returned without another attempt. +func WithTxRetry(ctx context.Context, fn func() error) error { + var err error + for attempt := 0; attempt < txRetryMaxAttempts; attempt++ { + if ctx.Err() != nil { + return ctx.Err() + } + err = fn() + if err == nil || !IsRetryableTxError(err) { + return err + } + if attempt == txRetryMaxAttempts-1 { + break // exhausted: surface the last error below rather than sleeping + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(txRetryBackoff(attempt)): + } + } + return err +} + +// txRetryBackoff returns a full-jitter exponential delay for the given +// zero-based attempt: a random duration in [0, min(base< txRetryMaxDelay { + d = txRetryMaxDelay + } + return time.Duration(rand.Int64N(int64(d) + 1)) +} diff --git a/internal/db/retry_test.go b/internal/db/retry_test.go new file mode 100644 index 0000000000..21924064a4 --- /dev/null +++ b/internal/db/retry_test.go @@ -0,0 +1,112 @@ +package db + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/go-sql-driver/mysql" +) + +func TestIsRetryableTxError(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"deadlock 1213", &mysql.MySQLError{Number: mysqlErrLockDeadlock, Message: "Deadlock found"}, true}, + {"lock wait timeout 1205", &mysql.MySQLError{Number: mysqlErrLockWaitTimeout, Message: "Lock wait timeout"}, true}, + {"wrapped deadlock", fmt.Errorf("create message: %w", &mysql.MySQLError{Number: mysqlErrLockDeadlock}), true}, + {"other mysql error (dup key 1062)", &mysql.MySQLError{Number: 1062, Message: "Duplicate entry"}, false}, + {"non-mysql error", errors.New("boom"), false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsRetryableTxError(tt.err); got != tt.want { + t.Fatalf("IsRetryableTxError(%v) = %v, want %v", tt.err, got, tt.want) + } + }) + } +} + +func TestWithTxRetry(t *testing.T) { + deadlock := &mysql.MySQLError{Number: mysqlErrLockDeadlock, Message: "Deadlock found"} + + t.Run("succeeds on first attempt", func(t *testing.T) { + calls := 0 + err := WithTxRetry(context.Background(), func() error { + calls++ + return nil + }) + if err != nil { + t.Fatalf("err = %v, want nil", err) + } + if calls != 1 { + t.Fatalf("calls = %d, want 1", calls) + } + }) + + t.Run("retries retryable error then succeeds", func(t *testing.T) { + calls := 0 + err := WithTxRetry(context.Background(), func() error { + calls++ + if calls < 3 { + return deadlock + } + return nil + }) + if err != nil { + t.Fatalf("err = %v, want nil", err) + } + if calls != 3 { + t.Fatalf("calls = %d, want 3", calls) + } + }) + + t.Run("returns non-retryable error immediately", func(t *testing.T) { + calls := 0 + sentinel := errors.New("not retryable") + err := WithTxRetry(context.Background(), func() error { + calls++ + return sentinel + }) + if !errors.Is(err, sentinel) { + t.Fatalf("err = %v, want %v", err, sentinel) + } + if calls != 1 { + t.Fatalf("calls = %d, want 1", calls) + } + }) + + t.Run("gives up after max attempts and returns last error", func(t *testing.T) { + calls := 0 + err := WithTxRetry(context.Background(), func() error { + calls++ + return deadlock + }) + if !errors.Is(err, deadlock) { + t.Fatalf("err = %v, want deadlock", err) + } + if calls != txRetryMaxAttempts { + t.Fatalf("calls = %d, want %d", calls, txRetryMaxAttempts) + } + }) + + t.Run("returns context error without calling fn when already cancelled", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + calls := 0 + err := WithTxRetry(ctx, func() error { + calls++ + return deadlock + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want context.Canceled", err) + } + if calls != 0 { + t.Fatalf("calls = %d, want 0", calls) + } + }) +} diff --git a/internal/message/message.go b/internal/message/message.go index 7706d95c10..539d0a11e4 100644 --- a/internal/message/message.go +++ b/internal/message/message.go @@ -141,35 +141,47 @@ func (s *service) Create(ctx context.Context, sessionID string, params CreateMes return message, nil } - tx, err := s.db.Begin() - if err != nil { - return Message{}, fmt.Errorf("failed to begin transaction: %w", err) - } - defer tx.Rollback() + // seq is allocated inside the transaction as MAX(seq)+1 for the session. + // That SELECT-then-INSERT races with any concurrent write to the same + // session (e.g. a foreground tool result and a background task's synthetic + // result landing together), which InnoDB resolves as a deadlock; retry the + // whole transaction on that transient conflict. See db.WithTxRetry. + var dbMessage db.Message + if err := db.WithTxRetry(ctx, func() error { + tx, err := s.db.Begin() + if err != nil { + return fmt.Errorf("failed to begin transaction: %w", err) + } + defer tx.Rollback() - qtx := s.q.WithTx(tx) + qtx := s.q.WithTx(tx) - seq, err = s.nextSeqTx(ctx, qtx, sessionID) - if err != nil { - return Message{}, err - } + seq, err := s.nextSeqTx(ctx, qtx, sessionID) + if err != nil { + return err + } - dbMessage, err := qtx.CreateMessage(ctx, db.CreateMessageParams{ - ID: uuid.New().String(), - SessionID: sessionID, - Role: string(params.Role), - Parts: string(partsJSON), - Model: sql.NullString{String: string(params.Model), Valid: true}, - Seq: sql.NullInt64{Int64: seq, Valid: true}, - Synthetic: params.Synthetic, - }) - if err != nil { + dbMessage, err = qtx.CreateMessage(ctx, db.CreateMessageParams{ + ID: uuid.New().String(), + SessionID: sessionID, + Role: string(params.Role), + Parts: string(partsJSON), + Model: sql.NullString{String: string(params.Model), Valid: true}, + Seq: sql.NullInt64{Int64: seq, Valid: true}, + Synthetic: params.Synthetic, + }) + if err != nil { + return err + } + + if err = tx.Commit(); err != nil { + return fmt.Errorf("failed to commit transaction: %w", err) + } + return nil + }); err != nil { return Message{}, err } - if err = tx.Commit(); err != nil { - return Message{}, fmt.Errorf("failed to commit transaction: %w", err) - } message, err := s.fromDBItem(dbMessage) if err != nil { return Message{}, err @@ -204,52 +216,61 @@ func (s *service) CreatePair(ctx context.Context, sessionID string, first, secon return Message{}, Message{}, err } - tx, err := s.db.Begin() - if err != nil { - return Message{}, Message{}, fmt.Errorf("failed to begin transaction: %w", err) - } - defer tx.Rollback() + // Same concurrent SELECT MAX(seq) / INSERT deadlock window as Create when + // the pair's sequence is allocated rather than supplied — retry the whole + // transaction on a transient conflict. See db.WithTxRetry. + var dbMsg1, dbMsg2 db.Message + if err := db.WithTxRetry(ctx, func() error { + tx, err := s.db.Begin() + if err != nil { + return fmt.Errorf("failed to begin transaction: %w", err) + } + defer tx.Rollback() - qtx := s.q.WithTx(tx) + qtx := s.q.WithTx(tx) - seq1 := first.Seq - seq2 := second.Seq - if seq1 == 0 { - seq1, err = s.nextSeqTx(ctx, qtx, sessionID) - if err != nil { - return Message{}, Message{}, err + seq1 := first.Seq + seq2 := second.Seq + if seq1 == 0 { + seq1, err = s.nextSeqTx(ctx, qtx, sessionID) + if err != nil { + return err + } + seq2 = seq1 + 1 } - seq2 = seq1 + 1 - } - dbMsg1, err := qtx.CreateMessage(ctx, db.CreateMessageParams{ - ID: uuid.New().String(), - SessionID: sessionID, - Role: string(first.Role), - Parts: string(firstJSON), - Model: sql.NullString{String: string(first.Model), Valid: true}, - Seq: sql.NullInt64{Int64: seq1, Valid: true}, - Synthetic: first.Synthetic, - }) - if err != nil { - return Message{}, Message{}, fmt.Errorf("failed to create first message: %w", err) - } + dbMsg1, err = qtx.CreateMessage(ctx, db.CreateMessageParams{ + ID: uuid.New().String(), + SessionID: sessionID, + Role: string(first.Role), + Parts: string(firstJSON), + Model: sql.NullString{String: string(first.Model), Valid: true}, + Seq: sql.NullInt64{Int64: seq1, Valid: true}, + Synthetic: first.Synthetic, + }) + if err != nil { + return fmt.Errorf("failed to create first message: %w", err) + } - dbMsg2, err := qtx.CreateMessage(ctx, db.CreateMessageParams{ - ID: uuid.New().String(), - SessionID: sessionID, - Role: string(second.Role), - Parts: string(secondJSON), - Model: sql.NullString{String: string(second.Model), Valid: true}, - Seq: sql.NullInt64{Int64: seq2, Valid: true}, - Synthetic: second.Synthetic, - }) - if err != nil { - return Message{}, Message{}, fmt.Errorf("failed to create second message: %w", err) - } + dbMsg2, err = qtx.CreateMessage(ctx, db.CreateMessageParams{ + ID: uuid.New().String(), + SessionID: sessionID, + Role: string(second.Role), + Parts: string(secondJSON), + Model: sql.NullString{String: string(second.Model), Valid: true}, + Seq: sql.NullInt64{Int64: seq2, Valid: true}, + Synthetic: second.Synthetic, + }) + if err != nil { + return fmt.Errorf("failed to create second message: %w", err) + } - if err = tx.Commit(); err != nil { - return Message{}, Message{}, fmt.Errorf("failed to commit transaction: %w", err) + if err = tx.Commit(); err != nil { + return fmt.Errorf("failed to commit transaction: %w", err) + } + return nil + }); err != nil { + return Message{}, Message{}, err } msg1, err := s.fromDBItem(dbMsg1)