Rework the standby cleaner a bit to remove replication slots

Author: davissp14

cmd/standby_cleaner/main.go

@@ -7,16 +7,14 @@ import (
 	"time"
 
 	"github.com/fly-apps/postgres-flex/pkg/flypg"
+	"github.com/fly-apps/postgres-flex/pkg/flypg/admin"
+	"github.com/jackc/pgx/v4"
 )
 
-// We need to adjust this to look at replication slots.
-// Pull ids from repmgr.show_nodes
-// Pull replication_slot ids that are inactive.
-
-// Remove replication slot if the slot_name id is inactive and is not
-// present as a repmgr node.
-
-var Minute int64 = 60
+var (
+	monitorFrequency           = time.Minute * 5
+	deadMemberRemovalThreshold = time.Hour * 24
+)
 
 func main() {
 	ctx := context.Background()
@@ -26,44 +24,91 @@ func main() {
 		os.Exit(1)
 	}
 
+	// TODO - We should connect using the flypgadmin user so we can  differentiate between
+	// internal admin connection usage and the actual repmgr process.
 	conn, err := flypgNode.RepMgr.NewLocalConnection(ctx)
 	if err != nil {
 		fmt.Printf("failed to open local connection: %s\n", err)
 		os.Exit(1)
 	}
 
-	ticker := time.NewTicker(5 * time.Second)
-	defer ticker.Stop()
+	seenAt := map[int]time.Time{}
 
-	seenAt := map[int]int64{}
+	ticker := time.NewTicker(monitorFrequency)
+	defer ticker.Stop()
 
-	for _ = range ticker.C {
+	for range ticker.C {
 		role, err := flypgNode.RepMgr.CurrentRole(ctx, conn)
 		if err != nil {
-			fmt.Printf("Failed to check role: %s", err)
+			fmt.Printf("Failed to check role: %s\n", err)
 			continue
 		}
+
 		if role != flypg.PrimaryRoleName {
 			continue
 		}
+
 		standbys, err := flypgNode.RepMgr.Standbys(ctx, conn)
 		if err != nil {
-			fmt.Printf("Failed to get standbys: %s", err)
+			fmt.Printf("Failed to query standbys: %s\n", err)
 			continue
 		}
+
 		for _, standby := range standbys {
 			newConn, err := flypgNode.RepMgr.NewRemoteConnection(ctx, standby.Ip)
+			defer newConn.Close(ctx)
 			if err != nil {
-				if time.Now().Unix()-seenAt[standby.Id] >= 10*Minute {
+				// TODO - Verify the exception that's getting thrown.
+				if time.Now().Sub(seenAt[standby.Id]) >= deadMemberRemovalThreshold {
 					if err := flypg.UnregisterMemberByID(ctx, int32(standby.Id)); err != nil {
-						fmt.Println(err.Error())
+						fmt.Printf("failed to unregister member %d: %v\n", standby.Id, err.Error())
+						continue
 					}
 
 					delete(seenAt, standby.Id)
 				}
-			} else {
-				seenAt[standby.Id] = time.Now().Unix()
-				newConn.Close(ctx)
+
+				continue
+			}
+
+			seenAt[standby.Id] = time.Now()
+		}
+
+		removeOrphanedReplicationSlots(ctx, conn, standbys)
+	}
+}
+
+func removeOrphanedReplicationSlots(ctx context.Context, conn *pgx.Conn, standbys []flypg.Standby) {
+	var orphanedSlots []admin.ReplicationSlot
+
+	slots, err := admin.ListReplicationSlots(ctx, conn)
+	if err != nil {
+		fmt.Printf("failed to list replication slots: %s", err)
+	}
+
+	// Identify orphaned replication slots.  active replication slots with the listed standbys
+	for _, slot := range slots {
+		matchFound := false
+		for _, standby := range standbys {
+			if slot.MemberID == int32(standby.Id) {
+				matchFound = true
+			}
+		}
+
+		if !matchFound && !slot.Active {
+			orphanedSlots = append(orphanedSlots, slot)
+		}
+	}
+
+	if len(orphanedSlots) > 0 {
+		fmt.Printf("%d orphaned replication slots detected", len(orphanedSlots))
+
+		for _, slot := range orphanedSlots {
+			fmt.Printf("dropping replication slot: %s", slot.Name)
+
+			if err := admin.DropReplicationSlot(ctx, conn, slot.Name); err != nil {
+				fmt.Printf("failed to drop replication slot %s: %v", slot.Name, err)
+				continue
 			}
 		}
 	}

pkg/flypg/admin/admin.go

@@ -3,6 +3,8 @@ package admin
 import (
 	"context"
 	"fmt"
+	"strconv"
+	"strings"
 
 	"github.com/jackc/pgx/v4"
 )
@@ -75,6 +77,14 @@ func DeleteDatabase(ctx context.Context, pg *pgx.Conn, name string) error {
 	return nil
 }
 
+type ReplicationSlot struct {
+	MemberID  int32
+	Name      string
+	Type      string
+	Active    bool
+	WalStatus string
+}
+
 func ListReplicationSlots(ctx context.Context, pg *pgx.Conn) ([]ReplicationSlot, error) {
 	sql := fmt.Sprintf("SELECT slot_name, slot_type, active, wal_status from pg_replication_slots;")
 	rows, err := pg.Query(ctx, sql)
@@ -91,7 +101,21 @@ func ListReplicationSlots(ctx context.Context, pg *pgx.Conn) ([]ReplicationSlot,
 			return nil, err
 		}
 
-		slots = append(slots, slot)
+		slotArr := strings.Split(slot.Name, "_")
+		// Only look at repmgr replication slots.
+		if slotArr[0] == "repmgr" {
+			// Resolve member id from slot name.
+			idStr := slotArr[2]
+
+			num, err := strconv.ParseInt(idStr, 10, 32)
+			if err != nil {
+				fmt.Printf("failed to parse member id %s", idStr)
+				continue
+			}
+
+			slot.MemberID = int32(num)
+			slots = append(slots, slot)
+		}
 	}
 
 	return slots, nil
@@ -174,13 +198,6 @@ type DbInfo struct {
 	Users []string `json:"users"`
 }
 
-type ReplicationSlot struct {
-	Name      string
-	Type      string
-	Active    bool
-	WalStatus string
-}
-
 func ListUsers(ctx context.Context, pg *pgx.Conn) ([]UserInfo, error) {
 	sql := `
 		select u.usename,