Random cleanup (#110)

* Random cleanup

* Cleanup

* Revert arg add

* bug fix

* Ensure connections and files are being closed

* Close connection

Author: davissp14

cmd/event_handler/main.go

@@ -71,7 +71,7 @@ func main() {
 }
 
 func evaluateClusterState(ctx context.Context, conn *pgx.Conn, node *flypg.Node) error {
-	primary, err := node.EvaluateClusterState(ctx, conn)
+	primary, err := flypg.PerformScreening(ctx, conn, node)
 	if errors.Is(err, flypg.ErrZombieDiagnosisUndecided) || errors.Is(err, flypg.ErrZombieDiscovered) {
 		if err := flypg.Quarantine(ctx, conn, node, primary); err != nil {
 			return fmt.Errorf("failed to quarantine failed primary: %s", err)

cmd/monitor/monitor_cluster_state.go

@@ -37,7 +37,7 @@ func clusterStateMonitorTick(ctx context.Context, node *flypg.Node) error {
 		return nil
 	}
 
-	primary, err := node.EvaluateClusterState(ctx, conn)
+	primary, err := flypg.PerformScreening(ctx, conn, node)
 	if errors.Is(err, flypg.ErrZombieDiagnosisUndecided) || errors.Is(err, flypg.ErrZombieDiscovered) {
 		if err := flypg.Quarantine(ctx, conn, node, primary); err != nil {
 			return fmt.Errorf("failed to quarantine failed primary: %s", err)

cmd/pg_unregister/main.go

@@ -34,6 +34,7 @@ func main() {
 		os.Exit(1)
 		return
 	}
+	defer conn.Close(ctx)
 
 	member, err := node.RepMgr.MemberByHostname(ctx, conn, string(hostnameBytes))
 	if err != nil {

internal/flypg/admin/admin.go

@@ -182,7 +182,6 @@ func FindDatabase(ctx context.Context, pg *pgx.Conn, name string) (*DbInfo, erro
 	`
 
 	sql = fmt.Sprintf(sql, name)
-
 	row := pg.QueryRow(ctx, sql)
 
 	db := new(DbInfo)

internal/flypg/config.go

@@ -102,6 +102,7 @@ func WriteConfigFiles(c Config) error {
 		return err
 	}
 	defer internalFile.Close()
+
 	userFile, err := os.OpenFile(c.UserConfigFile(), os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644)
 	if err != nil {
 		return err

internal/flypg/node.go

@@ -5,12 +5,10 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"math/rand"
 	"net"
 	"os"
 	"os/exec"
-	"os/user"
 	"strconv"
 	"time"
 
@@ -119,7 +117,7 @@ func (n *Node) Init(ctx context.Context) error {
 		return err
 	}
 
-	// Initiate a restore
+	// Check to see if we were just restored
 	if os.Getenv("FLY_RESTORED_FROM") != "" {
 		// Check to see if there's an active restore.
 		active, err := isRestoreActive()
@@ -134,64 +132,10 @@ func (n *Node) Init(ctx context.Context) error {
 		}
 	}
 
+	// Verify whether we are a booting zombie.
 	if ZombieLockExists() {
-		fmt.Println("Zombie lock detected!")
-		primaryStr, err := ReadZombieLock()
-		if err != nil {
-			return fmt.Errorf("failed to read zombie lock: %s", primaryStr)
-		}
-
-		// If the zombie lock contains a hostname, it means we were able to resolve the real primary and
-		// will attempt to rejoin it.
-		if primaryStr != "" {
-			ip := net.ParseIP(primaryStr)
-			if ip == nil {
-				return fmt.Errorf("zombie.lock file contains an invalid ipv6 address")
-			}
-
-			conn, err := n.RepMgr.NewRemoteConnection(ctx, ip.String())
-			if err != nil {
-				return fmt.Errorf("failed to establish a connection to our rejoin target %s: %s", ip.String(), err)
-			}
-			defer conn.Close(ctx)
-
-			primary, err := n.RepMgr.PrimaryMember(ctx, conn)
-			if err != nil {
-				return fmt.Errorf("failed to confirm primary on recover target %s: %s", ip.String(), err)
-			}
-
-			// Confirm that our rejoin target still identifies itself as the primary.
-			if primary.Hostname != ip.String() {
-				// Clear the zombie.lock file so we can attempt to re-resolve the correct primary.
-				if err := RemoveZombieLock(); err != nil {
-					return fmt.Errorf("failed to remove zombie lock: %s", err)
-				}
-
-				return ErrZombieLockPrimaryMismatch
-			}
-
-			// If the primary does not reside within our primary region, we cannot rejoin until it is.
-			if primary.Region != n.PrimaryRegion {
-				fmt.Printf("Primary region mismatch detected. The primary lives in '%s', while PRIMARY_REGION is set to '%s'\n", primary.Region, n.PrimaryRegion)
-				return ErrZombieLockRegionMismatch
-			}
-
-			if err := n.RepMgr.rejoinCluster(primary.Hostname); err != nil {
-				return fmt.Errorf("failed to rejoin cluster: %s", err)
-			}
-
-			// TODO - Wait for target cluster to register self as a standby.
-
-			if err := RemoveZombieLock(); err != nil {
-				return fmt.Errorf("failed to remove zombie lock: %s", err)
-			}
-
-			// Ensure the single instance created with the --force-rewind process is cleaned up properly.
-			utils.RunCommand("pg_ctl -D /data/postgresql/ stop", "postgres")
-		} else {
-			// TODO - Provide link to documention on how to address this
-			fmt.Println("Zombie lock file does not contain a hostname.")
-			fmt.Println("This likely means that we were unable to determine who the real primary is.")
+		if err := handleZombieLock(ctx, n); err != nil {
+			return err
 		}
 	}
 
@@ -205,18 +149,24 @@ func (n *Node) Init(ctx context.Context) error {
 		return fmt.Errorf("failed initialize cluster state store: %s", err)
 	}
 
-	if err := n.configure(ctx, store); err != nil {
-		return fmt.Errorf("failed to configure node: %s", err)
+	if err := n.configureInternal(store); err != nil {
+		return fmt.Errorf("failed to set internal config: %s", err)
+	}
+
+	if err := n.configureRepmgr(store); err != nil {
+		return fmt.Errorf("failed to configure repmgr config: %s", err)
 	}
 
 	if !n.isPGInitialized() {
-		// Check to see if repmgr cluster has been initialized.
+		// Check to see if cluster has already been initialized.
 		clusterInitialized, err := store.IsInitializationFlagSet()
 		if err != nil {
 			return fmt.Errorf("failed to verify cluster state %s", err)
 		}
 
 		if !clusterInitialized {
+			fmt.Println("Provisioning primary")
+
 			// Initialize ourselves as the primary.
 			if err := n.initializePG(); err != nil {
 				return fmt.Errorf("failed to initialize postgres %s", err)
@@ -227,6 +177,8 @@ func (n *Node) Init(ctx context.Context) error {
 			}
 
 		} else {
+			fmt.Println("Provisioning standby")
+			// Initialize ourselves as a standby
 			cloneTarget, err := n.RepMgr.ResolveMemberOverDNS(ctx)
 			if err != nil {
 				return err
@@ -300,7 +252,8 @@ func (n *Node) PostInit(ctx context.Context) error {
 			return fmt.Errorf("failed to register repmgr primary: %s", err)
 		}
 
-		// Set flag within consul to let future new members that the cluster exists
+		// 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")
 		}
@@ -326,7 +279,7 @@ func (n *Node) PostInit(ctx context.Context) error {
 
 		switch role {
 		case PrimaryRoleName:
-			primary, err := n.EvaluateClusterState(ctx, conn)
+			primary, err := PerformScreening(ctx, conn, n)
 			if errors.Is(err, ErrZombieDiagnosisUndecided) {
 				fmt.Println("Unable to confirm that we are the true primary!")
 				if err := Quarantine(ctx, conn, n, primary); err != nil {
@@ -383,13 +336,15 @@ func (n *Node) initializePG() error {
 		return nil
 	}
 
-	if err := ioutil.WriteFile("/data/.default_password", []byte(n.OperatorCredentials.Password), 0644); err != nil {
+	if err := os.WriteFile("/data/.default_password", []byte(n.OperatorCredentials.Password), 0644); err != nil {
 		return err
 	}
 	cmd := exec.Command("gosu", "postgres", "initdb", "--pgdata", n.DataDir, "--pwfile=/data/.default_password")
-	_, err := cmd.CombinedOutput()
+	if _, err := cmd.CombinedOutput(); err != nil {
+		return err
+	}
 
-	return err
+	return nil
 }
 
 func (n *Node) isPGInitialized() bool {
@@ -400,56 +355,6 @@ func (n *Node) isPGInitialized() bool {
 	return true
 }
 
-func (n *Node) configure(ctx context.Context, store *state.Store) error {
-	if err := n.configureInternal(store); err != nil {
-		return fmt.Errorf("failed to set internal config: %s", err)
-	}
-
-	if err := n.configureRepmgr(store); err != nil {
-		return fmt.Errorf("failed to configure repmgr config: %s", err)
-	}
-
-	return nil
-}
-
-func writeSSHKey() error {
-	err := os.Mkdir("/data/.ssh", 0700)
-	if err != nil && !os.IsExist(err) {
-		return err
-	}
-
-	key := os.Getenv("SSH_KEY")
-
-	keyFile, err := os.Create("/data/.ssh/id_rsa")
-	if err != nil {
-		return err
-	}
-	defer keyFile.Close()
-	_, err = keyFile.Write([]byte(key))
-	if err != nil {
-		return err
-	}
-
-	cert := os.Getenv("SSH_CERT")
-
-	certFile, err := os.Create("/data/.ssh/id_rsa-cert.pub")
-	if err != nil {
-		return err
-	}
-	defer certFile.Close()
-	_, err = certFile.Write([]byte(cert))
-	if err != nil {
-		return err
-	}
-
-	err = setSSHOwnership()
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
 func (n *Node) configureInternal(store *state.Store) error {
 	if err := n.InternalConfig.initialize(); err != nil {
 		return fmt.Errorf("failed to initialize internal config: %s", err)
@@ -605,6 +510,7 @@ func (n *Node) setDefaultHBA() error {
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
 	for _, entry := range entries {
 		str := fmt.Sprintf("%s %s %s %s %s\n", entry.Type, entry.Database, entry.User, entry.Address, entry.Method)
@@ -634,47 +540,17 @@ func openConnection(parentCtx context.Context, host string, database string, cre
 	return pgx.ConnectConfig(ctx, conf)
 }
 
-func setSSHOwnership() error {
-	cmdStr := fmt.Sprintf("chmod 600 %s %s", "/data/.ssh/id_rsa", "/data/.ssh/id_rsa-cert.pub")
-	cmd := exec.Command("sh", "-c", cmdStr)
-	_, err := cmd.Output()
-	return err
-}
-
 func setDirOwnership() error {
-	pgUser, err := user.Lookup("postgres")
-	if err != nil {
-		return err
-	}
-	pgUID, err := strconv.Atoi(pgUser.Uid)
-	if err != nil {
-		return err
-	}
-	pgGID, err := strconv.Atoi(pgUser.Gid)
+	pgUID, pgGID, err := utils.SystemUserIDs("postgres")
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to find postgres user ids: %s", err)
 	}
 
 	cmdStr := fmt.Sprintf("chown -R %d:%d %s", pgUID, pgGID, "/data")
 	cmd := exec.Command("sh", "-c", cmdStr)
-	_, err = cmd.Output()
-	return err
-}
-
-func (n *Node) EvaluateClusterState(ctx context.Context, conn *pgx.Conn) (string, error) {
-	standbys, err := n.RepMgr.StandbyMembers(ctx, conn)
-	if err != nil {
-		if !errors.Is(err, pgx.ErrNoRows) {
-			return "", fmt.Errorf("failed to query standbys")
-		}
-	}
-
-	sample, err := TakeDNASample(ctx, n, standbys)
-	if err != nil {
-		return "", fmt.Errorf("failed to evaluate cluster data: %s", err)
+	if _, err = cmd.Output(); err != nil {
+		return err
 	}
 
-	fmt.Println(DNASampleString(sample))
-
-	return ZombieDiagnosis(sample)
+	return nil
 }

internal/flypg/readonly.go

@@ -58,6 +58,7 @@ func BroadcastReadonlyChange(ctx context.Context, n *Node, enabled bool) error {
 	if err != nil {
 		return fmt.Errorf("failed to establish connection: %s", err)
 	}
+	defer conn.Close(ctx)
 
 	members, err := n.RepMgr.Members(ctx, conn)
 	if err != nil {

internal/flypg/repmgr.go

@@ -97,16 +97,16 @@ func (r *RepMgr) NewRemoteConnection(ctx context.Context, hostname string) (*pgx
 func (r *RepMgr) initialize() error {
 	r.setDefaults()
 
-	f, err := os.OpenFile(r.ConfigPath, os.O_APPEND|os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+	file, err := os.OpenFile(r.ConfigPath, os.O_APPEND|os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
 	if err != nil {
 		return nil
 	}
-	defer f.Close()
+	defer file.Close()
 
 	entries := []string{"include 'repmgr.internal.conf'\n", "include 'repmgr.user.conf'\n"}
 
 	for _, entry := range entries {
-		if _, err := f.WriteString(entry); err != nil {
+		if _, err := file.WriteString(entry); err != nil {
 			return fmt.Errorf("failed append configuration entry: %s", err)
 		}
 	}
@@ -238,6 +238,7 @@ func (r *RepMgr) writePasswdConf() error {
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
 	entries := []string{
 		fmt.Sprintf("*:*:*:%s:%s", r.Credentials.Username, r.Credentials.Password),
@@ -268,6 +269,7 @@ func (r *RepMgr) Members(ctx context.Context, pg *pgx.Conn) ([]Member, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer rows.Close()
 
 	var members []Member
 
@@ -316,7 +318,6 @@ func (r *RepMgr) StandbyMembers(ctx context.Context, conn *pgx.Conn) ([]Member,
 	}
 
 	var standbys []Member
-
 	for _, member := range members {
 		if member.Role == StandbyRoleName {
 			standbys = append(standbys, member)
@@ -391,7 +392,6 @@ func (r *RepMgr) ResolveMemberOverDNS(ctx context.Context) (*Member, error) {
 
 func (r *RepMgr) InRegionPeerIPs(ctx context.Context) ([]net.IPAddr, error) {
 	targets := fmt.Sprintf("%s.%s", r.PrimaryRegion, r.AppName)
-
 	return privnet.AllPeers(ctx, targets)
 }