Introduce SourceRuntime to own the source connection (GTID failover 1/3)

ferry.go

@@ -54,6 +54,8 @@ type Ferry struct {
 	SourceDB *sql.DB
 	TargetDB *sql.DB
 
+	sourceRuntime *SourceRuntime
+
 	ControlServer *ControlServer
 
 	BinlogStreamer *BinlogStreamer
@@ -91,6 +93,10 @@ type Ferry struct {
 	rowCopyCompleteCh chan struct{}
 }
 
+func (f *Ferry) SourceRuntime() *SourceRuntime {
+	return f.sourceRuntime
+}
+
 func (f *Ferry) NewDataIterator() *DataIterator {
 	f.ensureInitialized()
 
@@ -367,6 +373,8 @@ func (f *Ferry) Initialize() (err error) {
 		return err
 	}
 
+	f.sourceRuntime = NewSourceRuntime(f.SourceDB, f.Config.Source)
+
 	err = f.checkConnection("source", f.SourceDB)
 	if err != nil {
 		f.logger.WithError(err).Error("source connection checking failed")

source_runtime.go

@@ -0,0 +1,72 @@
+package ghostferry
+
+import (
+	"sync"
+
+	sql "github.com/Shopify/ghostferry/sqlwrapper"
+)
+
+// SourceRuntime owns the source connection and its config; safe for concurrent use.
+type SourceRuntime struct {
+	mu     sync.RWMutex
+	db     *sql.DB
+	config *DatabaseConfig
+
+	// retired holds replaced handles, kept open because in-flight work may
+	// still reference them; released by CloseRetired.
+	retired []*sql.DB
+}
+
+// NewSourceRuntime wraps an already-open source connection and its config.
+func NewSourceRuntime(db *sql.DB, config *DatabaseConfig) *SourceRuntime {
+	return &SourceRuntime{
+		db:     db,
+		config: config,
+	}
+}
+
+func (r *SourceRuntime) DB() *sql.DB {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return r.db
+}
+
+func (r *SourceRuntime) Config() *DatabaseConfig {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return r.config
+}
+
+// Replace opens a new connection from config and installs it as current,
+// retiring (not closing) the previous handle; on failure the source is unchanged.
+func (r *SourceRuntime) Replace(config *DatabaseConfig, logger Logger) (*sql.DB, error) {
+	newDB, err := config.SqlDB(logger)
+	if err != nil {
+		return nil, err
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.db != nil {
+		r.retired = append(r.retired, r.db)
+	}
+	r.db = newDB
+	r.config = config
+	return newDB, nil
+}
+
+// CloseRetired closes the handles retired by Replace; callers must ensure no
+// source work is still in flight.
+func (r *SourceRuntime) CloseRetired() {
+	r.mu.Lock()
+	retired := r.retired
+	r.retired = nil
+	r.mu.Unlock()
+
+	for _, db := range retired {
+		if db != nil {
+			_ = db.Close()
+		}
+	}
+}

source_runtime_test.go

@@ -0,0 +1,96 @@
+package ghostferry
+
+import (
+	"context"
+	"testing"
+
+	sql "github.com/Shopify/ghostferry/sqlwrapper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// SqlDB opens lazily (no dial), so tests get real handles without a server.
+func runtimeTestConfig(host string) *DatabaseConfig {
+	return &DatabaseConfig{Host: host, Port: 3306, User: "root", Net: "tcp"}
+}
+
+// isClosed: a cancelled context makes PingContext fail fast without dialing.
+func isClosed(t *testing.T, db *sql.DB) bool {
+	t.Helper()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err := db.PingContext(ctx)
+	require.Error(t, err)
+	return err != context.Canceled
+}
+
+func TestSourceRuntimeAccessors(t *testing.T) {
+	db := &sql.DB{Marginalia: "initial"}
+	cfg := runtimeTestConfig("initial-host")
+
+	rt := NewSourceRuntime(db, cfg)
+	assert.Same(t, db, rt.DB())
+	assert.Same(t, cfg, rt.Config())
+}
+
+func TestSourceRuntimeReplaceInstallsNewSourceAndRetiresOld(t *testing.T) {
+	oldDB := &sql.DB{Marginalia: "old"}
+	oldCfg := runtimeTestConfig("old-host")
+	rt := NewSourceRuntime(oldDB, oldCfg)
+
+	newCfg := runtimeTestConfig("new-host")
+	newDB, err := rt.Replace(newCfg, nil)
+	require.NoError(t, err)
+	require.NotNil(t, newDB)
+
+	assert.Same(t, newDB, rt.DB())
+	assert.Same(t, newCfg, rt.Config())
+	assert.NotSame(t, oldDB, rt.DB())
+
+	assert.Contains(t, rt.retired, oldDB)
+}
+
+func TestSourceRuntimeReplaceRetainsAllPriorHandles(t *testing.T) {
+	rt := NewSourceRuntime(&sql.DB{Marginalia: "gen0"}, runtimeTestConfig("h0"))
+
+	_, err := rt.Replace(runtimeTestConfig("h1"), nil)
+	require.NoError(t, err)
+	_, err = rt.Replace(runtimeTestConfig("h2"), nil)
+	require.NoError(t, err)
+
+	assert.Len(t, rt.retired, 2)
+}
+
+func TestSourceRuntimeCloseRetiredClosesOnlyRetiredHandles(t *testing.T) {
+	initial, err := runtimeTestConfig("gen0").SqlDB(nil)
+	require.NoError(t, err)
+	rt := NewSourceRuntime(initial, runtimeTestConfig("gen0"))
+
+	replaced, err := rt.Replace(runtimeTestConfig("gen1"), nil)
+	require.NoError(t, err)
+	current, err := rt.Replace(runtimeTestConfig("gen2"), nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, current.Close()) })
+
+	assert.False(t, isClosed(t, initial))
+	assert.False(t, isClosed(t, replaced))
+
+	rt.CloseRetired()
+	assert.True(t, isClosed(t, initial), "initial handle must be closed")
+	assert.True(t, isClosed(t, replaced), "intermediate handle must be closed")
+	assert.False(t, isClosed(t, current), "current handle must stay open")
+	assert.Same(t, current, rt.DB())
+
+	rt.CloseRetired()
+	assert.False(t, isClosed(t, current))
+}
+
+func TestSourceRuntimeNilInitial(t *testing.T) {
+	rt := NewSourceRuntime(nil, nil)
+	assert.Nil(t, rt.DB())
+	assert.Nil(t, rt.Config())
+
+	_, err := rt.Replace(runtimeTestConfig("h1"), nil)
+	require.NoError(t, err)
+	assert.Empty(t, rt.retired)
+}