Merge pull request #9728 from dolthub/elian/9725

#9725 - Fix AUTO_INCREMENT reuse after HA failover by refreshing trackers

Author: elianddb

go/libraries/doltcore/sqle/cluster/controller.go

@@ -53,6 +53,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/clusterdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate"
 	"github.com/dolthub/dolt/go/libraries/utils/config"
 	"github.com/dolthub/dolt/go/libraries/utils/filesys"
 	"github.com/dolthub/dolt/go/libraries/utils/jwtauth"
@@ -930,6 +931,12 @@ func (c *Controller) immediateTransitionToStandby() error {
 func (c *Controller) transitionToPrimary(saveConnID int) error {
 	c.setProviderIsStandby(false)
 	c.killRunningQueries(saveConnID)
+	// Promotion: ensure next AUTO_INCREMENT values do not reuse existing IDs.
+	// Refresh trackers for databases already loaded in this session so first
+	// writes on the new primary allocate IDs at or above current table maxima.
+	if err := c.refreshAutoIncrementTrackersForSessionDatabases(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -947,6 +954,69 @@ func (c *Controller) killRunningQueries(saveConnID int) {
 	}
 }
 
+// refreshAutoIncrementTrackersForSessionDatabases re-initializes AUTO_INCREMENT
+// trackers for databases already loaded in the current session by using their
+// working sets. This immediately aligns in-memory sequences with replicated
+// table metadata after a promotion so the first post-promotion inserts do not
+// reuse primary keys. This method only refreshes session-known databases; any
+// database not yet loaded is skipped here and its tracker initializes lazily on
+// first use. Only working sets are read; branches and remotes are not scanned
+// as part of this promotion path.
+func (c *Controller) refreshAutoIncrementTrackersForSessionDatabases() error {
+	if c == nil || c.sqlCtxFactory == nil {
+		return fmt.Errorf("cluster/controller: auto-inc refresh: missing sql context factory")
+	}
+
+	// Create a SQL context to access session and provider
+	sqlCtx, err := c.sqlCtxFactory(context.Background())
+	if err != nil {
+		return fmt.Errorf("cluster/controller: auto-inc refresh: create sql context: %w", err)
+	}
+	// Ensure proper session lifecycle for any storage-layer ops
+	defer sql.SessionEnd(sqlCtx.Session)
+	sql.SessionCommandBegin(sqlCtx.Session)
+	defer sql.SessionCommandEnd(sqlCtx.Session)
+
+	sess := dsess.DSessFromSess(sqlCtx.Session)
+	provider := sess.Provider()
+
+	for _, sdb := range provider.DoltDatabases() {
+		name := sdb.Name()
+		if name == clusterdb.DoltClusterDbName {
+			continue
+		}
+
+		// Load DB and its AI tracker
+		db, err := provider.Database(sqlCtx, name)
+		if err != nil {
+			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: provider db: %w", name, err)
+		}
+		gsp, ok := db.(globalstate.GlobalStateProvider)
+		if !ok {
+			// Non-versioned DBs don't participate in AUTO_INCREMENT global state
+			continue
+		}
+		ai, err := gsp.GetGlobalState().AutoIncrementTracker(sqlCtx)
+		if err != nil {
+			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: tracker: %w", name, err)
+		}
+
+		// Get working set roots only
+		state, ok, err := sess.LookupDbState(sqlCtx, name)
+		if err != nil {
+			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: lookup db state: %w", name, err)
+		}
+		if !ok || state.WorkingSet() == nil {
+			// Not loaded in session; defer to lazy initialization on first use
+			continue
+		}
+		if err := ai.InitWithRoots(sqlCtx, state.WorkingSet()); err != nil {
+			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: init: %w", name, err)
+		}
+	}
+	return nil
+}
+
 // called with c.mu held
 func (c *Controller) setProviderIsStandby(standby bool) {
 	if c.standbyCallback != nil {

integration-tests/go-sql-server-driver/tests/sql-server-cluster.yaml

@@ -2028,3 +2028,74 @@ tests:
     - on: server1
       queries:
         - exec: "set foreign_key_checks=0"
+
+- name: auto_increment maintained across failover
+  multi_repos:
+  - name: server1
+    with_files:
+    - name: server.yaml
+      contents: |
+        log_level: trace
+        listener:
+          host: 0.0.0.0
+          port: {{get_port "server1"}}
+        cluster:
+          standby_remotes:
+          - name: standby
+            remote_url_template: http://localhost:{{get_port "server2_cluster"}}/{database}
+          bootstrap_role: primary
+          bootstrap_epoch: 1
+          remotesapi:
+            port: {{get_port "server1_cluster"}}
+    server:
+      args: ["--config", "server.yaml"]
+      dynamic_port: server1
+  - name: server2
+    with_files:
+    - name: server.yaml
+      contents: |
+        log_level: trace
+        listener:
+          host: 0.0.0.0
+          port: {{get_port "server2"}}
+        cluster:
+          standby_remotes:
+          - name: standby
+            remote_url_template: http://localhost:{{get_port "server1_cluster"}}/{database}
+          bootstrap_role: standby
+          bootstrap_epoch: 1
+          remotesapi:
+            port: {{get_port "server2_cluster"}}
+    server:
+      args: ["--config", "server.yaml"]
+      dynamic_port: server2
+  connections:
+  - on: server1
+    queries:
+    - exec: 'CREATE DATABASE repo1'
+    - exec: 'USE repo1'
+    - exec: 'SET @@GLOBAL.dolt_cluster_ack_writes_timeout_secs = 10'
+    - exec: 'CREATE TABLE t (id INT AUTO_INCREMENT PRIMARY KEY, v INT)'
+    - exec: 'INSERT INTO t(v) VALUES (0),(0),(0),(0),(0)'
+  - on: server2
+    queries:
+    - exec: 'USE repo1'
+    - query: 'SELECT * FROM t ORDER BY id'
+      result:
+        columns: ['id', 'v']
+        rows: [['1','0'], ['2','0'], ['3','0'], ['4','0'], ['5','0']]
+      retry_attempts: 100
+  - on: server1
+    queries:
+    - exec: "CALL dolt_assume_cluster_role('standby', 2)"
+  - on: server2
+    queries:
+    - exec: "CALL dolt_assume_cluster_role('primary', 2)"
+  - on: server2
+    queries:
+    - exec: 'USE repo1'
+    - exec: 'INSERT INTO t(v) VALUES (1)'
+    - query: 'SELECT * FROM t ORDER BY id'
+      result:
+        columns: ['id', 'v']
+        rows: [['1','0'], ['2','0'], ['3','0'], ['4','0'], ['5','0'], ['6','1']]