copier: retry a chunk whose insert collides with a unique index

docs/invariants.md

@@ -269,12 +269,20 @@ writer holds its locks until its commit is visible and the swap's `ACCESS EXCLUS
 them all out (sync-commit test: a writer under `synchronous_commit = on` against a standby that
 never answers has its UPDATE delivered and left buffered with `Invisible` one and nothing
 applied, the confirmed position below the delivered one, until the backend is cancelled, after
-which the change lands and the catch-up converges). Also open: a primary-key move or a
-unique-value swap that straddles a chunk the copier is reading collides, on the copier's own
-insert of the moved row, with the stale shadow row the flush has not yet deleted (SQLSTATE
-`23505`); the catch-up's half of the exchange is correct and the resolution is the copier's, so
-the convergence tests for key moves and unique-key swaps start their load after the copy has
-landed. *Test obligation:* a
+which the change lands and the catch-up converges). *Enforced* for the copier's side
+of a unique value that moves between rows while its chunk is in flight: the chunk's insert
+collides (SQLSTATE `23505`, or `23P01` for an exclusion constraint) with the shadow row that
+still holds the value — a landed row the flush has yet to rewrite, or a row of the same chunk
+the copier itself read before the swap — and the copier runs the chunk again after a bounded
+wait (`copier.Options.CollisionAttempts` × `CollisionBackoff`, linear), counting each
+collision, until the flush has moved the value and the chunk lands or the attempts run out and
+the server's error is returned with no cause of its own (a copier integration test holds a
+colliding row across several attempts and removes it; a catch-up convergence test runs
+unique-key swaps throughout the copy). *Open:* a primary-key move whose row carries a unique
+secondary value, from a landed key to a key in flight: the drain's pair rule defers the old
+key's delete until the chunk lands, and the chunk cannot land while the old key's row holds the
+value, so the two wait on each other until the copier's attempts run out; the convergence test
+for key moves starts its load after the copy has landed. *Test obligation:* a
 marker-bearing UPDATE for a key inside an in-flight chunk asserts the flush waits for the chunk
 and the row is then completed from the copied shadow row, never an absent-row abort; a
 marker-bearing UPDATE for a key whose row the source then deleted before its chunk was read,

pkg/applier/catchup_collision_integration_test.go

@@ -0,0 +1,133 @@
+package applier_test
+
+import (
+	"fmt"
+	"math"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/block/pg-sprite/pkg/applier"
+	"github.com/block/pg-sprite/pkg/copier"
+)
+
+// A source change moves a unique value from a landed key to a key the copier
+// has yet to copy, and the catch-up's flushes are held, so the landed row's
+// shadow copy keeps the value when the copier reaches the other key: the
+// chunk carrying it collides on the shadow's unique index. The copier runs
+// the chunk again after a wait, the resumed catch-up flushes the landed
+// row's half of the move, the next attempt lands, and the shadow converges
+// (CO-4). Holding the flushes until the collision is seen makes the race
+// happen on every run rather than when the timing allows it.
+func TestCatchupConvergesWhenTheCopierCollidesWithAnUnflushedMove(t *testing.T) {
+	f := newCatchupFixture(t, 1500)
+	// The seeded rows hold uniq = id. Once the copier has landed key 10,
+	// uniq 10 moves to key 1000, far past the chunks in flight then.
+	move := &moveOnLanding{
+		t:      t,
+		f:      f,
+		landed: 10,
+		statements: []string{
+			`UPDATE %s SET uniq = 1000000 WHERE id = 10`,
+			`UPDATE %s SET uniq = 10 WHERE id = 1000`,
+		},
+	}
+	run := startCatchupWith(t, f, catchupSetup{copierClock: move})
+	move.shadow = pgx.Identifier{f.table.Schema, run.shadow.ShadowTable()}.Sanitize()
+
+	resume, err := run.catchup.Pause(t.Context())
+	require.NoError(t, err)
+	t.Cleanup(resume)
+
+	copied := make(chan error, 1)
+	var wg sync.WaitGroup
+	wg.Go(func() { copied <- run.copier.Run(t.Context(), f.pool) })
+	t.Cleanup(wg.Wait)
+
+	const collisionDeadline = 15 * time.Second
+	require.Eventually(t, func() bool { return run.copier.Collisions() >= 1 }, collisionDeadline, 10*time.Millisecond,
+		"the chunk carrying key 1000 should collide on uniq 10, which landed key 10 still holds in the shadow")
+	require.NoError(t, move.result())
+
+	resume()
+	const copyDeadline = 30 * time.Second
+	select {
+	case err := <-copied:
+		require.NoError(t, err, "the collided chunk lands once the flush has moved uniq 10 off key 10")
+	case <-time.After(copyDeadline):
+		require.FailNow(t, "the copy did not finish after the flushes resumed")
+	}
+	assert.Equal(t, copier.NewWatermark(math.MaxInt64), run.copier.Position().Cut, "every key is landed")
+	run.awaitStatus(t, "flush the move's half on landed key 10", func(s applier.Status) bool {
+		return s.Applied > 0 && s.Buffered == 0
+	})
+
+	run.stop(t)
+	f.assertConverged(t, run.shadow)
+}
+
+// moveOnLanding is the copier's clock. On the first reading after the shadow
+// holds key landed, it commits statements against the source in one
+// transaction, once; each statement's %s is the source table. Every reading
+// returns the wall clock.
+type moveOnLanding struct {
+	t          *testing.T
+	f          catchupFixture
+	landed     int64
+	statements []string
+	// shadow is the quoted shadow table, set before the copier runs.
+	shadow string
+
+	mu    sync.Mutex
+	moved bool
+	err   error
+}
+
+func (m *moveOnLanding) Now() time.Time {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if !m.moved && m.err == nil {
+		m.moved, m.err = m.tryMove()
+	}
+	return time.Now()
+}
+
+func (m *moveOnLanding) tryMove() (bool, error) {
+	ctx := m.t.Context()
+	var landed bool
+	if err := m.f.pool.QueryRow(ctx, "SELECT EXISTS (SELECT 1 FROM "+m.shadow+" WHERE id = $1)", m.landed).Scan(&landed); err != nil {
+		return false, fmt.Errorf("read whether key %d has landed: %w", m.landed, err)
+	}
+	if !landed {
+		return false, nil
+	}
+	err := pgx.BeginFunc(ctx, m.f.pool, func(tx pgx.Tx) error {
+		for _, statement := range m.statements {
+			if _, err := tx.Exec(ctx, fmt.Sprintf(statement, m.f.table.Qualified())); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return false, fmt.Errorf("move a unique value off landed key %d: %w", m.landed, err)
+	}
+	return true, nil
+}
+
+// result reports whether the move was committed.
+func (m *moveOnLanding) result() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.err != nil {
+		return m.err
+	}
+	if !m.moved {
+		return fmt.Errorf("key %d never landed, so the move was not made", m.landed)
+	}
+	return nil
+}

pkg/applier/catchup_integration_test.go

@@ -127,6 +127,9 @@ type catchupSetup struct {
 	// keys. The copy then runs to completion before the catch-up starts, so
 	// the shadow holds every row, and the test moves the position by hand.
 	positions applier.PositionSource
+	// copierClock, when set, is the copier's clock, which it reads as each
+	// chunk attempt starts and as each one commits.
+	copierClock progress.Clock
 }
 
 func startCatchup(t *testing.T, f catchupFixture, opts applier.CatchupOptions) *catchupRun {
@@ -178,6 +181,7 @@ func startCatchupWith(t *testing.T, f catchupFixture, setup catchupSetup) *catch
 	c, err := copier.NewCopier(f.target, shadow, lock, copier.Watermark{}, copier.Options{
 		Workers: 2,
 		Chunker: copier.ChunkerOptions{InitialRows: 50, MaxRows: 50},
+		Clock:   setup.copierClock,
 	})
 	require.NoError(t, err)
 	run := &catchupRun{f: f, shadow: shadow, lock: lock, copier: c, done: make(chan error, 1)}
@@ -348,9 +352,10 @@ func TestCatchupConvergesUnderMixedLoad(t *testing.T) {
 // travels together through the drain, and the flush completes the image's
 // unchanged blob from the old key's row before it deletes that row (CO-4,
 // D13). The workload starts after the copy has landed every key: the moved
-// row keeps its unique key, so mid copy the copier's own insert of the row
-// under its new key collides with the stale row the flush has not deleted
-// yet, which is the copier's race to resolve, not the catch-up's.
+// row keeps its unique key, so a move from a landed key to a key in flight
+// collides on the copier's insert with the row still under the old key, and
+// the pair rule defers that key's delete until the chunk lands — the two
+// would wait on each other until the copier's collision attempts ran out.
 func TestCatchupConvergesUnderKeyMoves(t *testing.T) {
 	f := newCatchupFixture(t, 1500)
 	run := startCatchup(t, f, applier.CatchupOptions{})
@@ -449,6 +454,38 @@ func TestCatchupConvergesUnderUniqueKeyMoves(t *testing.T) {
 	f.assertConverged(t, run.shadow)
 }
 
+// Unique key swaps run from before the copy claims its first chunk until
+// after it has landed every key, so a swap can straddle a chunk the copier
+// is reading: one row's new unique key is in a chunk in flight while the
+// other row, already landed, has had its half of the swap flushed — or the
+// chunk read both rows before the swap and the flush has since given a
+// landed row the value the chunk carries. Either way the chunk's insert
+// collides with a shadow row the applier still has to rewrite, the copier
+// runs the chunk again once the flush has committed, and the shadow
+// converges (CO-4).
+func TestCatchupConvergesUnderUniqueKeyMovesDuringTheCopy(t *testing.T) {
+	f := newCatchupFixture(t, 1500)
+	run := startCatchup(t, f, applier.CatchupOptions{})
+	load := testutil.StartLoad(t, f.pool, f.table, testutil.LoadSpec{
+		Seed: 6, Workers: 2, RatePerSecond: 40,
+		Mix: testutil.Mix{UniqueMove: 1},
+	})
+
+	run.awaitStatus(t, "discard changes to keys the copier has not claimed", func(s applier.Status) bool {
+		return s.Discarded > 0
+	})
+	run.copy(t)
+	t.Logf("chunk attempts the copier ran again after a collision: %d", run.copier.Collisions())
+	run.awaitStatus(t, "apply unique key swaps to landed keys", func(s applier.Status) bool {
+		return s.Applied > 0
+	})
+	summary := run.quiesce(t, load)
+	assert.Positive(t, summary.UniqueMoves)
+
+	run.stop(t)
+	f.assertConverged(t, run.shadow)
+}
+
 // A second Run on the same catch-up is refused: its stream is single-use.
 func TestCatchupRunsOnce(t *testing.T) {
 	f := newCatchupFixture(t, 10)

pkg/copier/copier.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -21,8 +22,9 @@ const DefaultWorkers = 4
 
 var (
 	// ErrInvalidOptions reports copier options that cannot bound the copy: a
-	// negative worker count, or a timeout below PostgreSQL's one-millisecond
-	// resolution, which the server would read as no timeout at all.
+	// negative worker count, a timeout below PostgreSQL's one-millisecond
+	// resolution, which the server would read as no timeout at all, a
+	// collision attempt count below one, or a negative collision backoff.
 	ErrInvalidOptions = errors.New("invalid copier options")
 	// ErrAlreadyRun reports a second Run on the same Copier; a Copier drives
 	// one copy and is then only a record of where it stopped.
@@ -31,7 +33,8 @@ var (
 
 // Options bounds the copy. Zero values take the defaults: DefaultWorkers
 // workers, the dbconn session timeouts for every chunk transaction, the
-// ChunkerOptions defaults, and the wall clock.
+// ChunkerOptions defaults, DefaultCollisionAttempts attempts spaced by
+// DefaultCollisionBackoff, and the wall clock.
 type Options struct {
 	// Workers is the number of chunk transactions in flight at once. The
 	// pool's connection limit bounds it from above.
@@ -42,6 +45,15 @@ type Options struct {
 	StatementTimeout time.Duration
 	// Chunker sizes the chunks.
 	Chunker ChunkerOptions
+	// CollisionAttempts is the number of times a chunk whose insert the
+	// server refused on a unique index or an exclusion constraint is run
+	// before the copy stops with that refusal. Each attempt reads the source
+	// afresh, so a value a source change moved between keys lands once the
+	// applier has flushed the move's other half into the shadow.
+	CollisionAttempts int
+	// CollisionBackoff is the wait before a collided chunk's second attempt;
+	// each later wait is one backoff longer.
+	CollisionBackoff time.Duration
 	// Clock times each chunk for the chunker's chunk-time throttling
 	// feedback (docs/copy-and-swap-design.md#d12--throttle-by-chunk-time-and-slot-lag).
 	Clock progress.Clock
@@ -62,6 +74,12 @@ func (o Options) withDefaults() Options {
 	if o.StatementTimeout == 0 {
 		o.StatementTimeout = dbconn.DefaultStatementTimeout
 	}
+	if o.CollisionAttempts == 0 {
+		o.CollisionAttempts = DefaultCollisionAttempts
+	}
+	if o.CollisionBackoff == 0 {
+		o.CollisionBackoff = DefaultCollisionBackoff
+	}
 	if o.Clock == nil {
 		o.Clock = progress.WallClock{}
 	}
@@ -95,6 +113,12 @@ func (o Options) validate() error {
 			return fmt.Errorf("%w: %s %s is below PostgreSQL's one-millisecond resolution; use zero for the default", ErrInvalidOptions, timeout.name, timeout.value)
 		}
 	}
+	if o.CollisionAttempts < 1 {
+		return fmt.Errorf("%w: collision attempts %d is below one; use zero for the default", ErrInvalidOptions, o.CollisionAttempts)
+	}
+	if o.CollisionBackoff < 0 {
+		return fmt.Errorf("%w: collision backoff %s is negative; use zero for the default", ErrInvalidOptions, o.CollisionBackoff)
+	}
 	return nil
 }
 
@@ -130,6 +154,9 @@ type Copier struct {
 	// rowsTotal is the source's catalog row count read once at the start of
 	// Run.
 	rowsTotal uint64
+	// collisions counts chunk attempts the server refused on a unique index
+	// or an exclusion constraint; Collisions reads it from any goroutine.
+	collisions atomic.Int64
 }
 
 // NewCopier prepares a copy of target into shadow that resumes after from
@@ -217,6 +244,14 @@ func (c *Copier) Position() Position {
 // Copier from Position().Watermark clears the shadow above the watermark
 // and continues, so a caller can back off and resume behind a busy table.
 //
+// A chunk whose insert the server refuses on a unique index or an
+// exclusion constraint (SQLSTATE 23505 or 23P01) is run again after a wait,
+// up to Options.CollisionAttempts times: the value it collided on was moved
+// between keys by a source change the applier has yet to flush, and a later
+// attempt's snapshot sees the shadow after that flush. A collision no
+// attempt outlasts returns the server's error wrapped, the chunk in flight,
+// and Collisions counts every refused attempt.
+//
 // Cancellation reaches the server as a closed connection, so a statement
 // that was running keeps running until it finishes or hits the transaction's
 // statement_timeout, and only then is its transaction rolled back; the
@@ -285,12 +320,14 @@ func (c *Copier) finish(cause error) error {
 	return nil
 }
 
-// work is one worker's loop: claim a chunk, copy it in its own transaction,
-// land it, and tell the chunker how long it took. It returns nil when the
-// chunker has no chunk left and the first error otherwise; a cancelled
-// context ends the loop with the context's error, which Run resolves to
-// the cancellation's cause. A chunk whose copy failed is left in flight: it
-// did not land, so its keys must not read as landed.
+// work is one worker's loop: claim a chunk, copy it in its own transaction
+// (again, after a wait, when the insert collided with a shadow row the
+// applier has yet to clear), land it, and tell the chunker how long the
+// committed attempt took. It returns nil when the chunker has no chunk left
+// and the first error otherwise; a cancelled context ends the loop with the
+// context's error, which Run resolves to the cancellation's cause. A chunk
+// whose copy failed is left in flight: it did not land, so its keys must
+// not read as landed.
 func (c *Copier) work(ctx context.Context, pool *pgxpool.Pool) error {
 	for ctx.Err() == nil {
 		chunk, ok, err := c.claim(ctx, pool)
@@ -300,15 +337,14 @@ func (c *Copier) work(ctx context.Context, pool *pgxpool.Pool) error {
 		if !ok {
 			return nil
 		}
-		started := c.opts.Clock.Now()
-		inserted, err := c.copyChunk(ctx, pool, chunk)
+		inserted, took, err := c.copyUntilClear(ctx, pool, chunk)
 		if err != nil {
 			return err
 		}
 		if err := c.land(chunk, inserted); err != nil {
 			return err
 		}
-		if err := c.chunker.Feedback(chunk, c.opts.Clock.Now().Sub(started)); err != nil {
+		if err := c.chunker.Feedback(chunk, took); err != nil {
 			return err
 		}
 	}

pkg/copier/doc.go

@@ -3,7 +3,9 @@
 // that tile the whole int64 range, so every key belongs to exactly one chunk
 // (CO-4) and each copy statement stays bounded (LK-3). A Copier copies those
 // chunks with several workers — one guarded, never-overwriting insert per
-// chunk — first clearing the shadow above the watermark when it resumes, and
+// chunk, run again after a bounded wait when it collides with a shadow row
+// the catch-up has yet to rewrite — first clearing the shadow above the
+// watermark when it resumes, and
 // reports a Position whose Classify is the applier's rule for a captured key:
 // uncut, in flight, or landed. A column the Shadow proof names in
 // Conversions is copied through that USING expression; the copier does not