fix(message): retry message writes on MySQL deadlock (1213/1205)

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<<attempt, max)]. Full
+// jitter keeps two transactions that just deadlocked from retrying in lockstep
+// and re-colliding.
+func txRetryBackoff(attempt int) time.Duration {
+	d := txRetryBaseDelay << attempt
+	if d <= 0 || d > txRetryMaxDelay {
+		d = txRetryMaxDelay
+	}
+	return time.Duration(rand.Int64N(int64(d) + 1))
+}

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)
+		}
+	})
+}

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)