make pg_unregister work with new names

Author: Ben Iofel

cmd/pg_unregister/main.go

@@ -3,11 +3,13 @@ package main
 import (
 	"context"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"os"
 
 	"github.com/fly-apps/postgres-flex/internal/flypg"
 	"github.com/fly-apps/postgres-flex/internal/utils"
+	"github.com/jackc/pgx/v5"
 )
 
 func main() {
@@ -40,7 +42,13 @@ func processUnregistration(ctx context.Context) error {
 	defer func() { _ = conn.Close(ctx) }()
 
 	member, err := node.RepMgr.MemberByHostname(ctx, conn, string(hostnameBytes))
-	if err != nil {
+	if errors.Is(err, pgx.ErrNoRows) {
+		// for historical reasons, flyctl passes in the 6pn as the hostname
+		member, err = node.RepMgr.MemberBy6PN(ctx, conn, string(hostnameBytes))
+		if err != nil {
+			return fmt.Errorf("failed to resolve member by hostname and 6pn: %s", err)
+		}
+	} else if err != nil {
 		return fmt.Errorf("failed to resolve member: %s", err)
 	}
 

internal/flypg/repmgr.go

@@ -449,6 +449,37 @@ func (*RepMgr) MemberByHostname(ctx context.Context, pg *pgx.Conn, hostname stri
 	return &member, nil
 }
 
+// MemberBy6PN returns a member by its 6PN address.
+// This assumes the hostnames are hostnames
+func (r *RepMgr) MemberBy6PN(ctx context.Context, pg *pgx.Conn, ip string) (*Member, error) {
+	members, err := r.Members(ctx, pg)
+	if err != nil {
+		return nil, err
+	}
+
+	resolver := privnet.GetResolver()
+	var lastErr error
+	for _, member := range members {
+		ips, err := resolver.LookupIPAddr(ctx, member.Hostname)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+
+		for _, addr := range ips {
+			if addr.IP.String() == ip {
+				return &member, nil
+			}
+		}
+	}
+
+	if lastErr != nil {
+		return nil, fmt.Errorf("no matches found for %s, and error encountered: %s", ip, lastErr)
+	}
+
+	return nil, nil
+}
+
 func (r *RepMgr) ResolveMemberOverDNS(ctx context.Context) (*Member, error) {
 	machineIds, err := r.InRegionPeerMachines(ctx)
 	if err != nil {

internal/privnet/sixpn.go

@@ -14,7 +14,7 @@ func AllPeers(ctx context.Context, appName string) ([]net.IPAddr, error) {
 }
 
 func Get6PN(ctx context.Context, hostname string) ([]net.IPAddr, error) {
-	r := getResolver()
+	r := GetResolver()
 	ips, err := r.LookupIPAddr(ctx, hostname)
 
 	if err != nil {
@@ -47,7 +47,7 @@ type Machine struct {
 }
 
 func AllMachines(ctx context.Context, appName string) ([]Machine, error) {
-	r := getResolver()
+	r := GetResolver()
 	txts, err := r.LookupTXT(ctx, fmt.Sprintf("vms.%s.internal", appName))
 	if err != nil {
 		return nil, err
@@ -67,7 +67,7 @@ func AllMachines(ctx context.Context, appName string) ([]Machine, error) {
 	return machines, nil
 }
 
-func getResolver() *net.Resolver {
+func GetResolver() *net.Resolver {
 	nameserver := os.Getenv("FLY_NAMESERVER")
 	if nameserver == "" {
 		nameserver = "fdaa::3"