Migrate source consumers to SourceRuntime (GTID failover 2/3)

cursor.go

@@ -36,6 +36,15 @@ type CursorConfig struct {
 	DB        *sql.DB
 	Throttler Throttler
 
+	// SourceRuntime, when set, is the authoritative source of the current
+	// source connection. NewCursor resolves the cursor's DB from it at cursor
+	// creation time so that cursors started after a source master failover run
+	// against the promoted writer. When nil, the static DB field is used. A
+	// cursor captures the resolved handle for the duration of its scan (a scan
+	// must be consistent against a single server); only newly created cursors
+	// observe a swap.
+	SourceRuntime *SourceRuntime
+
 	ColumnsToSelect []string
 	BuildSelect     func([]string, *TableSchema, PaginationKey, uint64) (squirrel.SelectBuilder, error)
 	// BatchSize is a pointer to the BatchSize in Config.UpdatableConfig which can be independently updated from this code.
@@ -45,10 +54,24 @@ type CursorConfig struct {
 	ReadRetries               int
 }
 
+// resolvedDB returns the source connection to use for a new cursor: the current
+// handle from SourceRuntime when configured, otherwise the static DB.
+func (c *CursorConfig) resolvedDB() *sql.DB {
+	if c.SourceRuntime != nil {
+		if db := c.SourceRuntime.DB(); db != nil {
+			return db
+		}
+	}
+	return c.DB
+}
+
 // returns a new Cursor with an embedded copy of itself
 func (c *CursorConfig) NewCursor(table *TableSchema, startPaginationKey, maxPaginationKey PaginationKey) *Cursor {
+	cfg := *c
+	// Bind the cursor to the current source connection for the whole scan.
+	cfg.DB = c.resolvedDB()
 	return &Cursor{
-		CursorConfig:                *c,
+		CursorConfig:                cfg,
 		Table:                       table,
 		MaxPaginationKey:            maxPaginationKey,
 		RowLock:                     true,

data_iterator.go

@@ -8,7 +8,11 @@ import (
 )
 
 type DataIterator struct {
-	DB                *sql.DB
+	DB *sql.DB
+	// SourceRuntime, when set, provides the current source connection so a data
+	// iteration run started after a source master failover targets the promoted
+	// writer. When nil, DB is used. Resolved once at the start of Run.
+	SourceRuntime     *SourceRuntime
 	Concurrency       int
 	SelectFingerprint bool
 
@@ -38,8 +42,18 @@ func (d *DataIterator) Run(tables []*TableSchema) {
 		d.StateTracker = NewStateTracker(0)
 	}
 
+	// Resolve the source connection for this run. When a SourceRuntime is
+	// configured it is authoritative (so a run started after a failover uses the
+	// promoted writer); otherwise fall back to the static DB.
+	db := d.DB
+	if d.SourceRuntime != nil {
+		if current := d.SourceRuntime.DB(); current != nil {
+			db = current
+		}
+	}
+
 	d.logger.WithField("tablesCount", len(tables)).Info("starting data iterator run")
-	tablesWithData, emptyTables, err := MaxPaginationKeys(d.DB, tables, d.logger)
+	tablesWithData, emptyTables, err := MaxPaginationKeys(db, tables, d.logger)
 	if err != nil {
 		d.ErrorHandler.Fatal("data_iterator", err)
 	}

ferry.go

@@ -109,13 +109,15 @@ func (f *Ferry) NewDataIterator() *DataIterator {
 
 	dataIterator := &DataIterator{
 		DB:                f.SourceDB,
+		SourceRuntime:     f.sourceRuntime,
 		Concurrency:       f.Config.DataIterationConcurrency,
 		SelectFingerprint: f.Config.VerifierType == VerifierTypeInline,
 
 		ErrorHandler: f.ErrorHandler,
 		CursorConfig: &CursorConfig{
-			DB:        f.SourceDB,
-			Throttler: f.Throttler,
+			DB:            f.SourceDB,
+			SourceRuntime: f.sourceRuntime,
+			Throttler:     f.Throttler,
 
 			BatchSize:                 &f.Config.UpdatableConfig.DataIterationBatchSize,
 			BatchSizePerTableOverride: f.Config.DataIterationBatchSizePerTableOverride,
@@ -227,6 +229,7 @@ func (f *Ferry) NewChecksumTableVerifier() *ChecksumTableVerifier {
 
 	return &ChecksumTableVerifier{
 		SourceDB:         f.SourceDB,
+		SourceRuntime:    f.sourceRuntime,
 		TargetDB:         f.TargetDB,
 		DatabaseRewrites: f.Config.DatabaseRewrites,
 		TableRewrites:    f.Config.TableRewrites,
@@ -246,6 +249,7 @@ func (f *Ferry) NewInlineVerifier() *InlineVerifier {
 
 	return &InlineVerifier{
 		SourceDB:                   f.SourceDB,
+		SourceRuntime:              f.sourceRuntime,
 		TargetDB:                   f.TargetDB,
 		DatabaseRewrites:           f.Config.DatabaseRewrites,
 		TableRewrites:              f.Config.TableRewrites,
@@ -319,13 +323,15 @@ func (f *Ferry) NewIterativeVerifier() (*IterativeVerifier, error) {
 	v := &IterativeVerifier{
 		CursorConfig: &CursorConfig{
 			DB:                        f.SourceDB,
+			SourceRuntime:             f.sourceRuntime,
 			BatchSize:                 &f.Config.UpdatableConfig.DataIterationBatchSize,
 			BatchSizePerTableOverride: f.Config.DataIterationBatchSizePerTableOverride,
 			ReadRetries:               f.Config.DBReadRetries,
 		},
 
 		BinlogStreamer:      f.BinlogStreamer,
 		SourceDB:            f.SourceDB,
+		SourceRuntime:       f.sourceRuntime,
 		TargetDB:            f.TargetDB,
 		CompressionVerifier: compressionVerifier,
 

inline_verifier.go

@@ -257,7 +257,11 @@ type InlineVerifierMismatches struct {
 }
 
 type InlineVerifier struct {
-	SourceDB                   *sql.DB
+	SourceDB *sql.DB
+	// SourceRuntime, when set, is the authoritative source of the current
+	// source connection (see currentSourceDB). It lets inline verification
+	// follow a source master failover. When nil, SourceDB is used directly.
+	SourceRuntime              *SourceRuntime
 	TargetDB                   *sql.DB
 	DatabaseRewrites           map[string]string
 	TableRewrites              map[string]string
@@ -274,7 +278,10 @@ type InlineVerifier struct {
 
 	sourceStmtCache *StmtCache
 	targetStmtCache *StmtCache
-	logger          Logger
+	// sourceDBMu guards SourceDB and sourceStmtCache when currentSourceDB
+	// resolves the runtime handle and may reset the cache on a failover swap.
+	sourceDBMu sync.Mutex
+	logger     Logger
 
 	// Used only for the ControlServer initiated VerifyDuringCutover
 	backgroundVerificationResultAndStatus VerificationResultAndStatus
@@ -524,8 +531,32 @@ func (v *InlineVerifier) VerifyDuringCutover() (VerificationResult, error) {
 	}, nil
 }
 
+// currentSourceDB returns the source connection inline verification should use,
+// resolving it from SourceRuntime when configured so a failover is followed. If
+// the source handle changed since the last call, the source statement cache is
+// reset because cached prepared statements are bound to the previous DB.
+// currentSourceDB returns the source connection inline verification should use,
+// resolving it from SourceRuntime when configured so a failover is followed. If
+// the source handle changed since the last call, the source statement cache is
+// reset because cached prepared statements are bound to the previous DB. The
+// resolved DB and its matching statement cache are returned together under a
+// lock so concurrent verification workers stay consistent.
+func (v *InlineVerifier) currentSourceDB() (*sql.DB, *StmtCache) {
+	v.sourceDBMu.Lock()
+	defer v.sourceDBMu.Unlock()
+
+	if v.SourceRuntime != nil {
+		if current := v.SourceRuntime.DB(); current != nil && current != v.SourceDB {
+			v.SourceDB = current
+			v.sourceStmtCache = NewStmtCache()
+		}
+	}
+	return v.SourceDB, v.sourceStmtCache
+}
+
 func (v *InlineVerifier) getFingerprintDataFromSourceDb(schemaName, tableName string, tx *sql.Tx, table *TableSchema, paginationKeys []interface{}) (map[string][]byte, map[string]map[string][]byte, error) {
-	return v.getFingerprintDataFromDb(v.SourceDB, v.sourceStmtCache, schemaName, tableName, tx, table, paginationKeys)
+	db, stmtCache := v.currentSourceDB()
+	return v.getFingerprintDataFromDb(db, stmtCache, schemaName, tableName, tx, table, paginationKeys)
 }
 
 func (v *InlineVerifier) getFingerprintDataFromTargetDb(schemaName, tableName string, tx *sql.Tx, table *TableSchema, paginationKeys []interface{}) (map[string][]byte, map[string]map[string][]byte, error) {

iterative_verifier.go

@@ -120,7 +120,11 @@ type IterativeVerifier struct {
 	BinlogStreamer      *BinlogStreamer
 	TableSchemaCache    TableSchemaCache
 	SourceDB            *sql.DB
-	TargetDB            *sql.DB
+	// SourceRuntime, when set, provides the current source connection so
+	// iterative verification follows a source master failover. When nil,
+	// SourceDB is used. See currentSourceDB.
+	SourceRuntime *SourceRuntime
+	TargetDB      *sql.DB
 
 	Tables              []*TableSchema
 	IgnoredTables       []string
@@ -289,6 +293,19 @@ func (v *IterativeVerifier) Result() (VerificationResultAndStatus, error) {
 	return v.verificationResultAndStatus, v.verificationErr
 }
 
+// currentSourceDB resolves the source connection from SourceRuntime when
+// configured (so verification follows a source master failover), otherwise the
+// static SourceDB. Unlike the inline verifier there is no persistent
+// source-bound statement cache to reset: GetHashes prepares per call.
+func (v *IterativeVerifier) currentSourceDB() *sql.DB {
+	if v.SourceRuntime != nil {
+		if db := v.SourceRuntime.DB(); db != nil {
+			return db
+		}
+	}
+	return v.SourceDB
+}
+
 func (v *IterativeVerifier) GetHashes(db *sql.DB, schemaName, tableName, paginationKeyColumn string, columns []schema.TableColumn, paginationKeys []interface{}) (map[string][]byte, error) {
 	sql, args, err := GetMd5HashesSql(schemaName, tableName, paginationKeyColumn, columns, paginationKeys)
 	if err != nil {
@@ -595,7 +612,7 @@ func (v *IterativeVerifier) compareFingerprints(paginationKeys []interface{}, ta
 	go func() {
 		defer wg.Done()
 		sourceErr = WithRetries(5, 0, v.logger, "get fingerprints from source db", func() (err error) {
-			sourceHashes, err = v.GetHashes(v.SourceDB, table.Schema, table.Name, table.GetPaginationColumn().Name, v.columnsToVerify(table), paginationKeys)
+			sourceHashes, err = v.GetHashes(v.currentSourceDB(), table.Schema, table.Name, table.GetPaginationColumn().Name, v.columnsToVerify(table), paginationKeys)
 			return
 		})
 	}()
@@ -627,7 +644,7 @@ func (v *IterativeVerifier) compareFingerprints(paginationKeys []interface{}, ta
 }
 
 func (v *IterativeVerifier) compareCompressedHashes(targetDb, targetTable string, table *TableSchema, paginationKeys []interface{}) ([]string, error) {
-	sourceHashes, err := v.CompressionVerifier.GetCompressedHashes(v.SourceDB, table.Schema, table.Name, table.GetPaginationColumn().Name, v.columnsToVerify(table), paginationKeys)
+	sourceHashes, err := v.CompressionVerifier.GetCompressedHashes(v.currentSourceDB(), table.Schema, table.Name, table.GetPaginationColumn().Name, v.columnsToVerify(table), paginationKeys)
 	if err != nil {
 		return nil, err
 	}

source_runtime_migration_test.go

@@ -0,0 +1,106 @@
+package ghostferry
+
+import (
+	"testing"
+
+	sql "github.com/Shopify/ghostferry/sqlwrapper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// These tests verify that source consumers resolve the current source
+// connection from a SourceRuntime when one is configured, and fall back to the
+// static DB otherwise. sql.DB values are sentinels; no query is executed.
+
+func TestCursorConfigResolvedDB(t *testing.T) {
+	staticDB := &sql.DB{Marginalia: "static"}
+
+	// No runtime: uses static DB.
+	c := &CursorConfig{DB: staticDB}
+	assert.Same(t, staticDB, c.resolvedDB())
+
+	// Runtime with a current DB: uses runtime DB.
+	runtimeDB := &sql.DB{Marginalia: "runtime"}
+	c.SourceRuntime = NewSourceRuntime(runtimeDB, runtimeTestConfig("h"))
+	assert.Same(t, runtimeDB, c.resolvedDB())
+
+	// Runtime present but nil DB: falls back to static DB.
+	c.SourceRuntime = NewSourceRuntime(nil, nil)
+	assert.Same(t, staticDB, c.resolvedDB())
+}
+
+func TestNewCursorBindsToResolvedDB(t *testing.T) {
+	staticDB := &sql.DB{Marginalia: "static"}
+	runtimeDB := &sql.DB{Marginalia: "runtime"}
+	rt := NewSourceRuntime(staticDB, runtimeTestConfig("old"))
+
+	cfg := &CursorConfig{DB: staticDB, SourceRuntime: rt}
+	table := &TableSchema{}
+
+	// Cursor created now binds to the current (static) handle.
+	cursor := cfg.NewCursor(table, NewUint64Key(0), NewUint64Key(10))
+	assert.Same(t, staticDB, cursor.DB)
+
+	// After a runtime swap, a newly created cursor binds to the new handle,
+	// while the already-created cursor keeps its handle (scan consistency).
+	_, err := rt.Replace(runtimeTestConfig("new"), nil)
+	require.NoError(t, err)
+	newHandle := rt.DB()
+
+	cursor2 := cfg.NewCursor(table, NewUint64Key(0), NewUint64Key(10))
+	assert.Same(t, newHandle, cursor2.DB)
+	assert.Same(t, staticDB, cursor.DB, "existing cursor must keep its handle")
+	_ = runtimeDB
+}
+
+func TestInlineVerifierCurrentSourceDBResetsStmtCacheOnSwap(t *testing.T) {
+	oldDB := &sql.DB{Marginalia: "old"}
+	rt := NewSourceRuntime(oldDB, runtimeTestConfig("old"))
+	v := &InlineVerifier{
+		SourceDB:        oldDB,
+		SourceRuntime:   rt,
+		sourceStmtCache: NewStmtCache(),
+	}
+
+	origCache := v.sourceStmtCache
+	db, cache := v.currentSourceDB()
+	assert.Same(t, oldDB, db)
+	assert.Same(t, origCache, cache, "no swap: cache unchanged")
+
+	// Swap the runtime; the verifier must pick up the new DB and reset its
+	// statement cache (statements were bound to the old DB).
+	_, err := rt.Replace(runtimeTestConfig("new"), nil)
+	require.NoError(t, err)
+	newHandle := rt.DB()
+
+	db, cache = v.currentSourceDB()
+	assert.Same(t, newHandle, db)
+	assert.NotSame(t, origCache, cache, "swap must reset the source statement cache")
+}
+
+func TestInlineVerifierCurrentSourceDBFallsBackWithoutRuntime(t *testing.T) {
+	db := &sql.DB{Marginalia: "static"}
+	v := &InlineVerifier{SourceDB: db, sourceStmtCache: NewStmtCache()}
+	got, _ := v.currentSourceDB()
+	assert.Same(t, db, got)
+}
+
+func TestIterativeVerifierCurrentSourceDB(t *testing.T) {
+	staticDB := &sql.DB{Marginalia: "static"}
+	v := &IterativeVerifier{SourceDB: staticDB}
+	assert.Same(t, staticDB, v.currentSourceDB())
+
+	runtimeDB := &sql.DB{Marginalia: "runtime"}
+	v.SourceRuntime = NewSourceRuntime(runtimeDB, runtimeTestConfig("h"))
+	assert.Same(t, runtimeDB, v.currentSourceDB())
+}
+
+func TestChecksumVerifierCurrentSourceDB(t *testing.T) {
+	staticDB := &sql.DB{Marginalia: "static"}
+	v := &ChecksumTableVerifier{SourceDB: staticDB}
+	assert.Same(t, staticDB, v.currentSourceDB())
+
+	runtimeDB := &sql.DB{Marginalia: "runtime"}
+	v.SourceRuntime = NewSourceRuntime(runtimeDB, runtimeTestConfig("h"))
+	assert.Same(t, runtimeDB, v.currentSourceDB())
+}

verifier.go

@@ -96,7 +96,11 @@ type ChecksumTableVerifier struct {
 	DatabaseRewrites map[string]string
 	TableRewrites    map[string]string
 	SourceDB         *sql.DB
-	TargetDB         *sql.DB
+	// SourceRuntime, when set, provides the current source connection so
+	// checksum verification follows a source master failover. When nil,
+	// SourceDB is used. See currentSourceDB.
+	SourceRuntime *SourceRuntime
+	TargetDB      *sql.DB
 
 	started *AtomicBoolean
 
@@ -107,6 +111,18 @@ type ChecksumTableVerifier struct {
 	wg     *sync.WaitGroup
 }
 
+// currentSourceDB resolves the source connection from SourceRuntime when
+// configured (so verification follows a source master failover), otherwise the
+// static SourceDB.
+func (v *ChecksumTableVerifier) currentSourceDB() *sql.DB {
+	if v.SourceRuntime != nil {
+		if db := v.SourceRuntime.DB(); db != nil {
+			return db
+		}
+	}
+	return v.SourceDB
+}
+
 func (v *ChecksumTableVerifier) VerifyBeforeCutover() error {
 	// All verification occurs in cutover for this verifier.
 	return nil
@@ -150,7 +166,7 @@ func (v *ChecksumTableVerifier) VerifyDuringCutover() (VerificationResult, error
 		go func() {
 			defer wg.Done()
 			query := fmt.Sprintf("CHECKSUM TABLE %s EXTENDED", sourceTable)
-			sourceRow := v.SourceDB.QueryRow(query)
+			sourceRow := v.currentSourceDB().QueryRow(query)
 			sourceChecksum, sourceErr = v.fetchChecksumValueFromRow(sourceRow)
 		}()