Fix race in primary election.

internal/flypg/node.go

@@ -139,8 +139,10 @@ func (n *Node) Init(ctx context.Context) error {
 		return fmt.Errorf("failed initialize cluster state store: %s", err)
 	}
 
-	// Check to see if cluster has already been initialized.
-	clusterInitialized, err := store.IsInitializationFlagSet()
+	// Atomically claim primary status for this cluster. Exactly one node
+	// across the whole app ever wins this; everyone else must join as a
+	// standby (or witness) against whoever did.
+	isPrimary, err := n.tryClaimPrimary(store)
 	if err != nil {
 		return fmt.Errorf("failed to verify cluster state %s", err)
 	}
@@ -153,7 +155,7 @@ func (n *Node) Init(ctx context.Context) error {
 	}
 
 	// Remote restores are only eligible on uninitialized clusters.
-	if !clusterInitialized {
+	if isPrimary {
 		// Determine if we are performing a remote restore.
 		if err := n.handleRemoteRestore(ctx, store); err != nil {
 			return fmt.Errorf("failed to handle remote restore: %s", err)
@@ -176,7 +178,7 @@ func (n *Node) Init(ctx context.Context) error {
 	}
 
 	if !n.PGConfig.isInitialized() {
-		if clusterInitialized {
+		if !isPrimary {
 			if n.RepMgr.Witness {
 				log.Println("Provisioning witness")
 				if err := n.PGConfig.writePasswordFile(n.OperatorCredentials.Password); err != nil {
@@ -188,7 +190,7 @@ func (n *Node) Init(ctx context.Context) error {
 				}
 			} else {
 				log.Println("Provisioning standby")
-				cloneTarget, err := n.RepMgr.ResolvePrimaryOverDNS(ctx)
+				cloneTarget, err := n.RepMgr.WaitForPrimaryOverDNS(ctx)
 				if err != nil {
 					return fmt.Errorf("failed to resolve member over dns: %s", err)
 				}
@@ -346,30 +348,24 @@ func (n *Node) PostInit(ctx context.Context) error {
 	} else {
 		// New member
 
-		// Check with consul to see if the cluster has already been initialized
+		// Check with consul to see if this node is the cluster's primary. This
+		// is the same durable claim made in Init - calling it again here is
+		// idempotent (same machineID, same stored owner) and correctly
+		// resumes as primary across a crash/restart between Init and here.
 		store, err := state.NewStore()
 		if err != nil {
 			return fmt.Errorf("failed initialize cluster state store. %v", err)
 		}
 
-		// The initialization flag is set after the primary is registered.
-		clusterInitialized, err := store.IsInitializationFlagSet()
+		isPrimary, err := n.tryClaimPrimary(store)
 		if err != nil {
 			return fmt.Errorf("failed to verify cluster state: %s", err)
 		}
 
-		if !clusterInitialized {
+		if isPrimary {
 			// Configure as primary
 			log.Println("Registering primary")
 
-			// Verify we reside within the clusters primary region
-			if !n.RepMgr.eligiblePrimary() {
-				return fmt.Errorf("unable to configure as the primary. expected region: %q, got: %q",
-					n.PrimaryRegion,
-					n.RepMgr.Region,
-				)
-			}
-
 			// Create required users
 			if err := n.setupCredentials(ctx, conn); err != nil {
 				return fmt.Errorf("failed to create required users: %s", err)
@@ -385,12 +381,6 @@ func (n *Node) PostInit(ctx context.Context) error {
 				return fmt.Errorf("failed to register repmgr primary: %s", err)
 			}
 
-			// Set initialization flag within consul so future members know they are joining
-			// an existing cluster.
-			if err := store.SetInitializationFlag(); err != nil {
-				return fmt.Errorf("failed to register cluster with consul")
-			}
-
 			// Let the boot process know that we've already been configured.
 			if err := issueRegistrationCert(); err != nil {
 				return fmt.Errorf("failed to issue registration certificate: %s", err)
@@ -439,6 +429,16 @@ func (n *Node) PostInit(ctx context.Context) error {
 	return nil
 }
 
+// tryClaimPrimary ensures only non-witness nodes in the primary region can
+// participate in the initial primary election.
+func (n *Node) tryClaimPrimary(store *state.Store) (bool, error) {
+	if n.RepMgr.Witness || !n.RepMgr.eligiblePrimary() {
+		return false, nil
+	}
+
+	return store.TryClaimPrimary(n.MachineID)
+}
+
 func (n *Node) setupCredentials(ctx context.Context, conn *pgx.Conn) error {
 	requiredCredentials := []admin.Credential{
 		n.OperatorCredentials,

internal/flypg/repmgr.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/fly-apps/postgres-flex/internal/privnet"
 	"github.com/fly-apps/postgres-flex/internal/utils"
@@ -26,6 +27,9 @@ const (
 	UnknownRoleName = ""
 
 	repmgrConsulKey = "repmgr"
+
+	primaryResolutionRetryInterval = time.Second
+	primaryResolutionTimeout       = time.Minute
 )
 
 type RepMgr struct {
@@ -512,6 +516,31 @@ func (r *RepMgr) ResolvePrimaryOverDNS(ctx context.Context) (*Member, error) {
 	return target, nil
 }
 
+// WaitForPrimaryOverDNS waits for the elected primary to finish initializing
+// and become available for cloning.
+func (r *RepMgr) WaitForPrimaryOverDNS(ctx context.Context) (*Member, error) {
+	waitCtx, cancel := context.WithTimeout(ctx, primaryResolutionTimeout)
+	defer cancel()
+
+	ticker := time.NewTicker(primaryResolutionRetryInterval)
+	defer ticker.Stop()
+
+	var lastErr error
+	for {
+		primary, err := r.ResolvePrimaryOverDNS(waitCtx)
+		if err == nil {
+			return primary, nil
+		}
+		lastErr = err
+
+		select {
+		case <-waitCtx.Done():
+			return nil, fmt.Errorf("waiting for cloneable primary: %w (last error: %v)", waitCtx.Err(), lastErr)
+		case <-ticker.C:
+		}
+	}
+}
+
 func (r *RepMgr) InRegionPeerIPs(ctx context.Context) ([]net.IPAddr, error) {
 	targets := fmt.Sprintf("%s.%s", r.PrimaryRegion, r.AppName)
 	return privnet.AllPeers(ctx, targets)

internal/flypg/state/store.go

@@ -36,23 +36,42 @@ func NewStore() (*Store, error) {
 	}, nil
 }
 
-func (c *Store) SetInitializationFlag() error {
-	kv := &api.KVPair{Key: c.targetKey("INITIALIZED"), Value: []byte("true")}
-	_, err := c.Client.KV().Put(kv, nil)
-	return err
-}
-
-func (c *Store) IsInitializationFlagSet() (bool, error) {
-	result, _, err := c.Client.KV().Get(c.targetKey("INITIALIZED"), nil)
+// TryClaimPrimary attempts to atomically claim this cluster for the given
+// machine. It replaces the old check-then-act pattern (read the flag, then
+// separately decide to become primary), which raced whenever more than one
+// node's Init/PostInit ran the check before the eventual primary had gotten
+// around to setting the flag - both would see it unset and both would try
+// to become primary.
+//
+// Consul's CAS with ModifyIndex: 0 only succeeds if the key does not already
+// exist, so when multiple nodes call this concurrently, exactly one write
+// wins across the whole cluster - Consul's Raft log serializes it, so there
+// is no window between "check" and "act" for another node to slip through.
+//
+// Returns true if this machine is the primary: either it just won the
+// claim, or it already won an earlier attempt (e.g. retrying after its own
+// crash), recognized by the stored value matching its own machineID rather
+// than someone else's.
+func (c *Store) TryClaimPrimary(machineID string) (bool, error) {
+	key := c.targetKey("INITIALIZED")
+
+	ok, _, err := c.Client.KV().CAS(&api.KVPair{
+		Key:         key,
+		Value:       []byte(machineID),
+		ModifyIndex: 0,
+	}, nil)
 	if err != nil {
-		return false, err
+		return false, fmt.Errorf("failed to claim primary: %s", err)
 	}
-
-	if result == nil {
-		return false, nil
+	if ok {
+		return true, nil
+	}
+	pair, _, err := c.Client.KV().Get(key, nil)
+	if err != nil {
+		return false, fmt.Errorf("failed to verify primary claim: %s", err)
 	}
 
-	return true, nil
+	return pair != nil && string(pair.Value) == machineID, nil
 }
 
 func (c *Store) PushUserConfig(key string, config []byte) error {