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 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,20 @@ package main
 import (
 	"flag"
 	"fmt"
+	"github.com/fly-apps/postgres-flex/pkg/flypg"
+	"strconv"
 
 	"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")
+	// This might not actually always be the new primary. %p from repmgr is variably the new or
+	// old primary. In the events that we subscribe to it's always either empty or the new primary.
+	// In the future if we subscribe to repmgrd_failover_promote, then we would have to change this
+	// name.
+	newPrimary := flag.String("new-node-id", "", "the new primary node id")
 	success := flag.String("success", "", "success (1) failure (0)")
 	details := flag.String("details", "", "details")
 	flag.Parse()
@@ -33,6 +40,37 @@ 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)
+		}
+		newNodeID, err := strconv.Atoi(*newPrimary)
+		if err != nil {
+			fmt.Printf("failed to parse new node id: %s", err)
+		}
+		node, err := client.Node(int32(newNodeID))
+		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

@@ -33,6 +33,7 @@ type Node struct {
 	DataDir   string
 	Region    string
 	PGPort    int
+	ProxyPort int
 
 	SUCredentials       Credentials
 	OperatorCredentials Credentials
@@ -46,6 +47,7 @@ func NewNode() (*Node, error) {
 	node := &Node{
 		AppName:             "local",
 		PGPort:              5433,
+		ProxyPort:           5432,
 		DataDir:             "/data/postgresql",
 		ManagerDatabaseName: "repmgr",
 		ManagerConfigPath:   "/data/repmgr.conf",
@@ -148,6 +150,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 +265,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(n.ProxyPort))
+	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 +347,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, n.PGPort)
+	_, 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,8 +71,9 @@ 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 %d -d %s -U %s -f %s standby clone -F",
 		ipStr,
+		node.PGPort,
 		node.ManagerDatabaseName,
 		node.ManagerCredentials.Username,
 		node.ManagerConfigPath)
@@ -99,11 +100,15 @@ func writeManagerConf(node Node) error {
 		"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'",
+		"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))