fix(history): retain maintenance writer lock through rollback

.changeset/history-maintenance-rollback.md

@@ -0,0 +1,5 @@
+---
+"ftw": patch
+---
+
+Wait for cancelled history maintenance transactions to finish rollback before allowing another writer to enter.

go/internal/state/archive_write.go

@@ -47,29 +47,44 @@ func (s *Store) tryArchiveBatch(parent context.Context, prepare, write func(cont
 	if err != nil {
 		return err
 	}
-	defer conn.Close()
 	var busyMS int
-	if err := conn.QueryRowContext(parent, `PRAGMA busy_timeout`).Scan(&busyMS); err != nil {
-		return err
-	}
+	var tx *sql.Tx
+	var restoreBusy, viewLocked, writerLocked bool
 	defer func() {
-		restore, done := context.WithTimeout(context.Background(), time.Second)
-		defer done()
-		if _, err := conn.ExecContext(restore, fmt.Sprintf("PRAGMA busy_timeout=%d", busyMS)); err != nil {
-			// Never lend a connection with changed lock behaviour to live work.
-			_ = conn.Raw(func(any) error { return driver.ErrBadConn })
+		if tx != nil {
+			_ = tx.Rollback()
+		}
+		if restoreBusy {
+			restore, done := context.WithTimeout(context.Background(), time.Second)
+			if _, err := conn.ExecContext(restore, fmt.Sprintf("PRAGMA busy_timeout=%d", busyMS)); err != nil {
+				// Never lend changed lock behaviour to live work.
+				_ = conn.Raw(func(any) error { return driver.ErrBadConn })
+			}
+			done()
+		}
+		// ErrTxDone can precede SQLite's automatic rollback. Close waits
+		// for that rollback before another live or archive writer can enter.
+		_ = conn.Close()
+		if writerLocked {
+			s.historyWriteMu.Unlock()
+		}
+		if viewLocked {
+			s.archiveViewMu.Unlock()
 		}
 	}()
+	if err := conn.QueryRowContext(parent, `PRAGMA busy_timeout`).Scan(&busyMS); err != nil {
+		return err
+	}
+	restoreBusy = true
 	if _, err := conn.ExecContext(parent, `PRAGMA busy_timeout=0`); err != nil {
 		return err
 	}
 	txCtx, cancelTx := context.WithCancel(parent)
 	defer cancelTx()
-	tx, err := conn.BeginTx(txCtx, nil)
+	tx, err = conn.BeginTx(txCtx, nil)
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback()
 	if prepare != nil {
 		readCtx, cancel := context.WithTimeout(parent, 30*time.Second)
 		err := prepare(readCtx, tx)
@@ -89,11 +104,11 @@ func (s *Store) tryArchiveBatch(parent context.Context, prepare, write func(cont
 	if err := lockContext(ctx, s.archiveViewMu.TryLock); err != nil {
 		return err
 	}
-	defer s.archiveViewMu.Unlock()
+	viewLocked = true
 	if err := lockContext(ctx, s.historyWriteMu.TryLock); err != nil {
 		return err
 	}
-	defer s.historyWriteMu.Unlock()
+	writerLocked = true
 	if err := write(ctx, tx); err != nil {
 		return err
 	}

go/internal/state/history_rollback_test.go

@@ -125,7 +125,7 @@ func (c *slowHistoryRollbackConn) BeginTx(ctx context.Context, opts driver.TxOpt
 
 func (c *slowHistoryRollbackConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	r, err := c.Conn.(driver.ExecerContext).ExecContext(ctx, query, args)
-	if err == nil && ctx.Done() != nil && c.d.armed.CompareAndSwap(true, false) {
+	if err == nil && ctx.Done() != nil && isHistoryTestWrite(query) && c.d.armed.CompareAndSwap(true, false) {
 		<-ctx.Done()
 		// Let database/sql's cancellation goroutine claim the Tx before the
 		// driver returns, while this operation still owns the connection lock.
@@ -161,3 +161,104 @@ func (tx *slowHistoryRollbackTx) Rollback() error {
 	}
 	return tx.Tx.Rollback()
 }
+
+func isHistoryTestWrite(query string) bool {
+	q := strings.ToUpper(strings.TrimSpace(query))
+	return strings.HasPrefix(q, "INSERT ") || strings.HasPrefix(q, "UPDATE ") || strings.HasPrefix(q, "DELETE ")
+}
+
+func TestHistoryMaintenanceWaitsForCancelledRollback(t *testing.T) {
+	for _, kind := range []string{"archive", "backfill", "completion"} {
+		t.Run(kind, func(t *testing.T) {
+			s := freshStore(t)
+			if kind == "backfill" {
+				if err := s.RecordSamples([]Sample{{Driver: "meter", Metric: "pv_w", TsMs: seriesHourMs, Value: 1200}}); err != nil {
+					t.Fatal(err)
+				}
+			}
+			for _, query := range []string{`DELETE FROM ts_series_hour`, `DELETE FROM history_migrations WHERE name='ts-series-hour-v1'`} {
+				if _, err := s.history.Exec(query); err != nil {
+					t.Fatal(err)
+				}
+			}
+			if err := s.history.Close(); err != nil {
+				t.Fatal(err)
+			}
+			d := &slowHistoryRollbackDriver{started: make(chan struct{}), release: make(chan struct{})}
+			name := fmt.Sprintf("maintenance-rollback-%s-%d", kind, time.Now().UnixNano())
+			sql.Register(name, d)
+			var err error
+			s.history, err = sql.Open(name, s.historyPath+strings.Replace(sqlitePragmas, "synchronous(NORMAL)", "synchronous(FULL)", 1))
+			if err != nil {
+				t.Fatal(err)
+			}
+			var release sync.Once
+			defer release.Do(func() { close(d.release) })
+			run := func(ctx context.Context) error {
+				switch kind {
+				case "archive":
+					return s.tryArchiveBatch(ctx, nil, func(ctx context.Context, tx *sql.Tx) error {
+						_, err := tx.ExecContext(ctx, `INSERT INTO history_site_cursor(key,value) VALUES ('rollback-test',1)`)
+						return err
+					})
+				case "backfill":
+					_, err := s.backfillSeriesHoursSnapshot(ctx, 64)
+					return err
+				default:
+					return s.ensureSeriesHours(ctx)
+				}
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+			defer cancel()
+			d.armed.Store(true)
+			done := make(chan error, 1)
+			go func() { done <- run(ctx) }()
+			select {
+			case <-d.started:
+			case err := <-done:
+				t.Fatalf("maintenance ended before rollback: %v", err)
+			case <-time.After(5 * time.Second):
+				t.Fatal("maintenance did not start rollback")
+			}
+			waitCtx, stop := context.WithTimeout(context.Background(), 50*time.Millisecond)
+			lockErr := lockContext(waitCtx, s.historyWriteMu.TryLock)
+			stop()
+			if lockErr == nil {
+				s.historyWriteMu.Unlock()
+			}
+			release.Do(func() { close(d.release) })
+			select {
+			case err = <-done:
+			case <-time.After(5 * time.Second):
+				t.Fatal("maintenance did not finish after rollback")
+			}
+			if lockErr == nil {
+				t.Fatal("maintenance released the writer mutex before SQLite rollback finished")
+			}
+			if !errors.Is(err, context.DeadlineExceeded) {
+				t.Fatalf("cancelled maintenance: %v", err)
+			}
+			if err := run(context.Background()); err != nil {
+				t.Fatalf("retry after rollback: %v", err)
+			}
+			// Check every pooled connection: archive must restore its zero busy timeout.
+			var conns []*sql.Conn
+			defer func() {
+				for _, c := range conns {
+					c.Close()
+				}
+			}()
+			for i := 0; i < 4; i++ {
+				c, err := s.history.Conn(context.Background())
+				if err != nil {
+					t.Fatal(err)
+				}
+				conns = append(conns, c)
+				var busy int
+				if err := c.QueryRowContext(context.Background(), `PRAGMA busy_timeout`).Scan(&busy); err != nil || busy != 5000 {
+					t.Fatalf("busy timeout leaked: %d %v", busy, err)
+				}
+			}
+		})
+	}
+}

go/internal/state/history_series_hour.go

@@ -135,7 +135,12 @@ func (s *Store) ensureSeriesHours(ctx context.Context) error {
 		return err
 	}
 	defer s.historyWriteMu.Unlock()
-	tx, err := s.history.BeginTx(ctx, nil)
+	conn, err := s.history.Conn(ctx)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	tx, err := conn.BeginTx(ctx, nil)
 	if err != nil {
 		return err
 	}
@@ -177,7 +182,12 @@ func (s *Store) backfillSeriesHoursSnapshot(ctx context.Context, maxHours int) (
 	defer cancelTx()
 	// Deferred: reading never owns SQLite's writer. A concurrent write makes
 	// the later upgrade fail with BUSY_SNAPSHOT, including the cursor write.
-	tx, err := s.history.BeginTx(txCtx, nil)
+	conn, err := s.history.Conn(txCtx)
+	if err != nil {
+		return false, err
+	}
+	defer conn.Close()
+	tx, err := conn.BeginTx(txCtx, nil)
 	if err != nil {
 		return false, err
 	}
@@ -235,7 +245,11 @@ func (s *Store) backfillSeriesHoursSnapshot(ctx context.Context, maxHours int) (
 	if err := lockContext(writeCtx, s.historyWriteMu.TryLock); err != nil {
 		return false, err
 	}
-	defer s.historyWriteMu.Unlock()
+	defer func() {
+		_ = tx.Rollback()
+		_ = conn.Close()
+		s.historyWriteMu.Unlock()
+	}()
 	stmt, err := tx.PrepareContext(writeCtx, `INSERT INTO ts_series_hour
 		(driver_id,metric_id,hour_ms,sum_value,min_value,max_value,n,last_ts_ms)
 		VALUES (?,?,?,?,?,?,?,?) ON CONFLICT DO NOTHING`)