Merge remote-tracking branch 'origin/master' into node-name

fly-apps · Aug 12, 2024 · fd6192e · fd6192e
2 parents 2473d7a + 7a8ac4d
commit fd6192e
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 15 deletions.
diff --git a/cmd/pg_unregister/main.go b/cmd/pg_unregister/main.go
@@ -5,9 +5,12 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"log"
 	"os"
+	"time"
 
 	"github.com/fly-apps/postgres-flex/internal/flypg"
+	"github.com/fly-apps/postgres-flex/internal/flypg/admin"
 	"github.com/fly-apps/postgres-flex/internal/utils"
 	"github.com/jackc/pgx/v5"
 )
@@ -56,5 +59,43 @@ func processUnregistration(ctx context.Context) error {
 		return fmt.Errorf("failed to unregister member: %v", err)
 	}
 
+	slotName := fmt.Sprintf("repmgr_slot_%d", member.ID)
+	if err := removeReplicationSlot(ctx, conn, slotName); err != nil {
+		return err
+	}
+
 	return nil
 }
+
+func removeReplicationSlot(ctx context.Context, conn *pgx.Conn, slotName string) error {
+	ticker := time.NewTicker(1 * time.Second)
+	timeout := time.After(10 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-timeout:
+			return fmt.Errorf("timed out trying to drop replication slot")
+		case <-ticker.C:
+			slot, err := admin.GetReplicationSlot(ctx, conn, slotName)
+			if err != nil {
+				if err == pgx.ErrNoRows {
+					return nil
+				}
+				return fmt.Errorf("failed to get replication slot %s: %v", slotName, err)
+			}
+
+			if slot.Active {
+				log.Printf("Slot %s is still active, waiting...", slotName)
+				continue
+			}
+
+			if err := admin.DropReplicationSlot(ctx, conn, slotName); err != nil {
+				return fmt.Errorf("failed to drop replication slot %s: %v", slotName, err)
+			}
+
+			return nil
+		}
+	}
+}
diff --git a/internal/flypg/admin/admin.go b/internal/flypg/admin/admin.go
@@ -130,6 +130,18 @@ type ReplicationSlot struct {
 	RetainedWalInBytes int
 }
 
+func GetReplicationSlot(ctx context.Context, pg *pgx.Conn, slotName string) (*ReplicationSlot, error) {
+	sql := fmt.Sprintf("SELECT slot_name, active, wal_status, pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn) AS retained_wal FROM pg_replication_slots where slot_name = '%s';", slotName)
+	row := pg.QueryRow(ctx, sql)
+
+	var slot ReplicationSlot
+	if err := row.Scan(&slot.Name, &slot.Active, &slot.WalStatus, &slot.RetainedWalInBytes); err != nil {
+		return nil, err
+	}
+
+	return &slot, nil
+}
+
 func ListReplicationSlots(ctx context.Context, pg *pgx.Conn) ([]ReplicationSlot, error) {
 	sql := "SELECT slot_name, active, wal_status, pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn) AS retained_wal FROM pg_replication_slots;"
 	rows, err := pg.Query(ctx, sql)
@@ -167,7 +179,6 @@ func ListReplicationSlots(ctx context.Context, pg *pgx.Conn) ([]ReplicationSlot,
 
 func DropReplicationSlot(ctx context.Context, pg *pgx.Conn, name string) error {
 	sql := fmt.Sprintf("SELECT pg_drop_replication_slot('%s');", name)
-
 	_, err := pg.Exec(ctx, sql)
 	if err != nil {
 		return err

diff --git a/pg16/Dockerfile b/pg16/Dockerfile
@@ -21,15 +21,19 @@ COPY ./bin/* /fly/bin/
 
 FROM ubuntu:24.04
 
-ENV PGDATA=/data/postgresql
-ENV PGPASSFILE=/data/.pgpass
-ENV AWS_SHARED_CREDENTIALS_FILE=/data/.aws/credentials
 ARG VERSION
 ARG PG_MAJOR_VERSION
 ARG POSTGIS_MAJOR=3
 ARG HAPROXY_VERSION=2.8
 ARG REPMGR_VERSION=5.4.1-1build2
 
+ENV PGDATA=/data/postgresql
+ENV PGPASSFILE=/data/.pgpass
+ENV AWS_SHARED_CREDENTIALS_FILE=/data/.aws/credentials
+ENV PG_MAJOR_VERSION=${PG_MAJOR_VERSION}
+ENV PATH="/usr/lib/postgresql/${PG_MAJOR_VERSION}/bin:$PATH"
+
+
 LABEL fly.app_role=postgres_cluster
 LABEL fly.version=${VERSION}
 LABEL fly.pg-version=${PG_VERSION}
@@ -63,16 +67,15 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     haproxy=$HAPROXY_VERSION.\* \
     && apt autoremove -y && apt clean
 
-
-# Add PostgreSQL bin directory to PATH
-ENV PATH="/usr/lib/postgresql/${PG_MAJOR_VERSION}/bin:$PATH"
-
 # Copy Go binaries from the builder stage
 COPY --from=builder /fly/bin/* /usr/local/bin
 
 # Copy Postgres exporter
 COPY --from=wrouesnel/postgres_exporter:latest /postgres_exporter /usr/local/bin/
 
+# Move pg_rewind into path.
+RUN ln -s /usr/lib/postgresql/${PG_MAJOR_VERSION}/bin/pg_rewind /usr/bin/pg_rewind
+
 ADD /config/* /fly/
 RUN mkdir -p /run/haproxy/
 RUN usermod -d /data postgres

diff --git a/pg16/Dockerfile-timescaledb b/pg16/Dockerfile-timescaledb
@@ -21,15 +21,18 @@ COPY ./bin/* /fly/bin/
 
 FROM ubuntu:24.04
 
-ENV PGDATA=/data/postgresql
-ENV PGPASSFILE=/data/.pgpass
-ENV AWS_SHARED_CREDENTIALS_FILE=/data/.aws/credentials
 ARG VERSION
 ARG PG_MAJOR_VERSION
 ARG POSTGIS_MAJOR=3
 ARG HAPROXY_VERSION=2.8
 ARG REPMGR_VERSION=5.4.1-1build2
 
+ENV PGDATA=/data/postgresql
+ENV PGPASSFILE=/data/.pgpass
+ENV AWS_SHARED_CREDENTIALS_FILE=/data/.aws/credentials
+ENV PG_MAJOR_VERSION=${PG_MAJOR_VERSION}
+ENV PATH="/usr/lib/postgresql/${PG_MAJOR_VERSION}/bin:$PATH"
+
 LABEL fly.app_role=postgres_cluster
 LABEL fly.version=${VERSION}
 LABEL fly.pg-version=${PG_VERSION}
@@ -69,16 +72,15 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     haproxy=$HAPROXY_VERSION.\* \
     && apt autoremove -y && apt clean
 
-
-# Add PostgreSQL bin directory to PATH
-ENV PATH="/usr/lib/postgresql/${PG_MAJOR_VERSION}/bin:$PATH"
-
 # Copy Go binaries from the builder stage
 COPY --from=builder /fly/bin/* /usr/local/bin
 
 # Copy Postgres exporter
 COPY --from=wrouesnel/postgres_exporter:latest /postgres_exporter /usr/local/bin/
 
+# Move pg_rewind into path.
+RUN ln -s /usr/lib/postgresql/${PG_MAJOR_VERSION}/bin/pg_rewind /usr/bin/pg_rewind
+
 ADD /config/* /fly/
 RUN mkdir -p /run/haproxy/
 RUN usermod -d /data postgres