Rip out haproxy and replace it with pgbouncer

Author: Dov Alperin

.gitignore

@@ -0,0 +1 @@
+.idea/

Dockerfile

@@ -20,7 +20,7 @@ LABEL fly.version=${VERSION}
 LABEL fly.pg-version=${PG_VERSION}
 
 RUN apt-get update && apt-get install --no-install-recommends -y \
-    ca-certificates iproute2 haproxy postgresql-14-repmgr curl bash dnsutils vim procps jq \
+    ca-certificates iproute2 haproxy postgresql-14-repmgr curl bash dnsutils vim procps jq pgbouncer \
     && apt autoremove -y
 
 COPY --from=0 /fly/bin/* /usr/local/bin

cmd/event_handler/main.go

@@ -3,13 +3,15 @@ package main
 import (
 	"flag"
 	"fmt"
+	"github.com/fly-apps/postgres-flex/pkg/flypg"
 
 	"github.com/fly-apps/postgres-flex/pkg/flypg/state"
 )
 
 func main() {
 	event := flag.String("event", "", "event type")
 	nodeID := flag.Int("node-id", 0, "the node id")
+	newPrimary := flag.Int("new-node-id", 0, "the new primary node id")
 	success := flag.String("success", "", "success (1) failure (0)")
 	details := flag.String("details", "", "details")
 	flag.Parse()
@@ -33,6 +35,33 @@ func main() {
 		if err := client.RegisterPrimary(string(node.Value)); err != nil {
 			fmt.Printf("failed to register primary: %s", err)
 		}
+
+		flypgNode, err := flypg.NewNode()
+		if err != nil {
+			fmt.Printf("failed to reconfigure pgbouncer primary %s\n", err)
+		}
+
+		fmt.Println("Reconfiguring pgbouncer primary")
+		if err := flypgNode.ConfigurePGBouncerPrimary(string(node.Value), true); err != nil {
+			fmt.Printf("failed to reconfigure pgbouncer primary %s\n", err)
+		}
+	case "standby_follow":
+		client, err := state.NewConsulClient()
+		if err != nil {
+			fmt.Printf("failed to initialize consul client: %s", err)
+		}
+		node, err := client.Node(int32(*newPrimary))
+		if err != nil {
+			fmt.Printf("failed to find node: %s", err)
+		}
+		flypgNode, err := flypg.NewNode()
+		if err != nil {
+			fmt.Printf("failed to reconfigure pgbouncer primary %s\n", err)
+		}
+		fmt.Println("Reconfiguring pgbouncer primary")
+		if err := flypgNode.ConfigurePGBouncerPrimary(string(node.Value), true); err != nil {
+			fmt.Printf("failed to reconfigure pgbouncer primary %s\n", err)
+		}
 	default:
 		// noop
 	}

cmd/start/main.go

@@ -33,7 +33,7 @@ func main() {
 		for range t.C {
 
 			if err := node.PostInit(); err != nil {
-				fmt.Printf("failed post-init: %s. Retrying...", err)
+				fmt.Printf("failed post-init: %s. Retrying...\n", err)
 				continue
 			}
 
@@ -43,12 +43,7 @@ func main() {
 
 	svisor := supervisor.New("flypg", 5*time.Minute)
 
-	proxyEnv := map[string]string{
-		"FLY_APP_NAME":      os.Getenv("FLY_APP_NAME"),
-		"PRIMARY_REGION":    os.Getenv("PRIMARY_REGION"),
-		"PG_LISTEN_ADDRESS": node.PrivateIP,
-	}
-	svisor.AddProcess("proxy", "/usr/sbin/haproxy -W -db -f /fly/haproxy.cfg", supervisor.WithEnv(proxyEnv), supervisor.WithRestart(0, 1*time.Second))
+	svisor.AddProcess("pgbouncer", "/usr/sbin/pgbouncer /fly/pgbouncer.ini", supervisor.WithRestart(0, 1*time.Second))
 
 	env := map[string]string{
 		"PGDATA":     os.Getenv("PGDATA"),

config/haproxy.cfg

@@ -0,0 +1,26 @@
+[pgbouncer]
+
+listen_addr = *
+listen_port = 5432
+unix_socket_dir = /tmp
+
+auth_user = postgres
+auth_file = /data/pgbouncer.auth
+
+admin_users = postgres
+
+user = postgres
+
+pool_mode = transaction
+
+max_client_conn = 100
+default_pool_size = 20
+min_pool_size = 5
+reserve_pool_size = 5
+reserve_pool_timeout = 3
+
+log_connections = 1
+log_disconnections = 1
+log_pooler_errors = 1
+
+%include /data/pgbouncer.database.ini

config/pgbouncer.ini

@@ -148,6 +148,11 @@ func (n *Node) Init() error {
 		return fmt.Errorf("failed to configure postgres %s", err)
 	}
 
+	fmt.Println("Configuring pgbouncer auth")
+	if err := n.ConfigurePGBouncerAuth(); err != nil {
+		return fmt.Errorf("failed to configure pgbouncer auth %s", err)
+	}
+
 	return nil
 }
 
@@ -258,9 +263,24 @@ func (n *Node) PostInit() error {
 		}
 	}
 
+	primaryIP, err = client.CurrentPrimary()
+	if err != nil {
+		return fmt.Errorf("failed to query current primary: %s", err)
+	}
+
+	fmt.Println("Configuring pgbouncer primary")
+	if err := n.ConfigurePGBouncerPrimary(primaryIP, false); err != nil {
+		return fmt.Errorf("failed to configure pgbouncer primary %s", err)
+	}
+
 	return nil
 }
 
+func (n *Node) NewPGBouncerConnection(ctx context.Context) (*pgx.Conn, error) {
+	host := net.JoinHostPort(n.PrivateIP, strconv.Itoa(5432))
+	return openConnection(ctx, host, "pgbouncer", n.OperatorCredentials)
+}
+
 func (n *Node) NewLocalConnection(ctx context.Context) (*pgx.Conn, error) {
 	host := net.JoinHostPort(n.PrivateIP, strconv.Itoa(n.PGPort))
 	return openConnection(ctx, host, "postgres", n.OperatorCredentials)
@@ -325,6 +345,53 @@ func (n *Node) initializePostgres() error {
 	return err
 }
 
+func (n *Node) ConfigurePGBouncerAuth() error {
+	path := fmt.Sprintf("%s/pgbouncer.auth", "/data")
+	file, err := os.OpenFile(path, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0644)
+	if err != nil {
+		return err
+	}
+	contents := fmt.Sprintf("\"%s\" \"%s\"", n.OperatorCredentials.Username, n.OperatorCredentials.Password)
+	_, err = file.Write([]byte(contents))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (n *Node) ConfigurePGBouncerPrimary(primary string, reload bool) error {
+	path := fmt.Sprintf("%s/pgbouncer.database.ini", "/data")
+	file, err := os.OpenFile(path, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0644)
+	if err != nil {
+		return err
+	}
+	contents := fmt.Sprintf("[databases]\n* = host=%s port=%d\n", primary, 5433)
+	_, err = file.Write([]byte(contents))
+	if err != nil {
+		return err
+	}
+
+	if reload {
+		err = n.ReloadPGBouncerConfig()
+		if err != nil {
+			fmt.Printf("failed to reconfigure pgbouncer primary %s\n", err)
+		}
+	}
+	return nil
+}
+
+func (n *Node) ReloadPGBouncerConfig() error {
+	conn, err := n.NewPGBouncerConnection(context.TODO())
+	if err != nil {
+		return err
+	}
+	_, err = conn.Exec(context.TODO(), "RELOAD;")
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (n *Node) configurePostgres() error {
 	cmdStr := fmt.Sprintf("sed -i \"s/#shared_preload_libraries.*/shared_preload_libraries = 'repmgr'/\" /data/postgresql/postgresql.conf")
 	return runCommand(cmdStr)

pkg/flypg/node.go

@@ -71,7 +71,7 @@ func cloneFromPrimary(node Node, ipStr string) error {
 		return err
 	}
 
-	cmdStr = fmt.Sprintf("repmgr -h %s -d %s -U %s -f %s standby clone -F",
+	cmdStr = fmt.Sprintf("repmgr -h %s -p 5433 -d %s -U %s -f %s standby clone -F",
 		ipStr,
 		node.ManagerDatabaseName,
 		node.ManagerCredentials.Username,
@@ -92,18 +92,23 @@ func writeManagerConf(node Node) error {
 	}
 
 	conf := map[string]interface{}{
-		"node_id":                    fmt.Sprint(node.ID),
-		"node_name":                  fmt.Sprintf("'%s'", node.PrivateIP),
-		"conninfo":                   fmt.Sprintf("'host=%s port=%d user=%s dbname=%s connect_timeout=10'", node.PrivateIP, node.PGPort, node.ManagerCredentials.Username, node.ManagerDatabaseName),
-		"data_directory":             fmt.Sprintf("'%s'", node.DataDir),
-		"failover":                   "'automatic'",
-		"promote_command":            fmt.Sprintf("'repmgr standby promote -f %s --log-to-file'", node.ManagerConfigPath),
-		"follow_command":             fmt.Sprintf("'repmgr standby follow -f %s --log-to-file --upstream-node-id=%%n'", node.ManagerConfigPath),
-		"event_notification_command": fmt.Sprintf("'/usr/local/bin/event_handler -node-id %%n -event %%e -success %%s -details \"%%d\"'"),
-		"event_notifications":        "'repmgrd_failover_promote,standby_promote'",
+		"node_id":         fmt.Sprint(node.ID),
+		"node_name":       fmt.Sprintf("'%s'", node.PrivateIP),
+		"conninfo":        fmt.Sprintf("'host=%s port=%d user=%s dbname=%s connect_timeout=10'", node.PrivateIP, node.PGPort, node.ManagerCredentials.Username, node.ManagerDatabaseName),
+		"data_directory":  fmt.Sprintf("'%s'", node.DataDir),
+		"failover":        "'automatic'",
+		"promote_command": fmt.Sprintf("'repmgr standby promote -f %s --log-to-file'", node.ManagerConfigPath),
+		"follow_command":  fmt.Sprintf("'repmgr standby follow -f %s --log-to-file --upstream-node-id=%%n'", node.ManagerConfigPath),
+		// FIXME: %p does not always exist (usually on the new primary node) so event_handler throws and error, this doesn't break anything but is clunky
+		"event_notification_command": fmt.Sprintf("'/usr/local/bin/event_handler -node-id %%n -event %%e -success %%s -details \"%%d\" -new-node-id %%p'"),
+		"event_notifications":        "'repmgrd_failover_promote,standby_promote,standby_follow'",
 		"location":                   node.Region,
 	}
 
+	if !node.ValidPrimary() {
+		conf["priority"] = "0"
+	}
+
 	for key, value := range conf {
 		str := fmt.Sprintf("%s=%s\n", key, value)
 		_, err := file.Write([]byte(str))