diff --git a/internal/flypg/node.go b/internal/flypg/node.go index fe252327..d5744423 100644 --- a/internal/flypg/node.go +++ b/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, diff --git a/internal/flypg/repmgr.go b/internal/flypg/repmgr.go index 4714d275..74a99bb3 100644 --- a/internal/flypg/repmgr.go +++ b/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) diff --git a/internal/flypg/state/store.go b/internal/flypg/state/store.go index 2efe9458..e56cd16d 100644 --- a/internal/flypg/state/store.go +++ b/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 {