From 852a7177840f8c7df270d120c227e1ed76af9a5b Mon Sep 17 00:00:00 2001
From: Roman Vanicek <roman.vanicek@ivasoft.cz>
Date: Thu, 4 Apr 2024 23:32:38 +0200
Subject: [PATCH] Load balance to read-only nodes

---
 main.go | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/main.go b/main.go
index 4e615d1..32cb8dd 100644
--- a/main.go
+++ b/main.go
@@ -47,7 +47,7 @@ type pgpoolForcedConfig struct {
 	num    int
 	host   string
 	port   int
-	weight int
+	weight float32
 	dir    string
 	flag   string
 }
@@ -326,6 +326,7 @@ func configure(msg pgpoolConfigMessage, nodesHistory map[pgpoolNodeKey]pgpoolNod
 	// be modified without pgpool restart
 	// Note: However there is no need to print the nodes that are down
 	hasPrimaryChanged := false
+	processedNums := make(map[int]bool)
 	for _, i := range msg.Instances {
 		nodeKey := pgpoolNodeKey{address: i.IpAddress.String(), port: i.Port}
 		var num int
@@ -337,12 +338,13 @@ func configure(msg pgpoolConfigMessage, nodesHistory map[pgpoolNodeKey]pgpoolNod
 			hasPrimaryChanged = hasPrimaryChanged || !i.IsReadOnly
 		}
 		nodesHistory[nodeKey] = pgpoolNodeInfo{num: num, isPrimary: !i.IsReadOnly}
+		processedNums[num] = true
 
-		var weight int
-		if forced, exists := forcedNodes[pgpoolForcedKey{hostName: i.HostName, port: i.Port}]; exists {
+		var weight float32
+		if forced, exists := forcedNodes[pgpoolForcedKey{hostName: i.HostName, port: i.Port}]; exists && forced.weight != -1 {
 			weight = forced.weight
 		} else {
-			weight = 0
+			weight = i.Weight
 		}
 
 		var flag string
@@ -359,6 +361,14 @@ func configure(msg pgpoolConfigMessage, nodesHistory map[pgpoolNodeKey]pgpoolNod
 		conf["backend_flag"+suffix] = flag
 	}
 
+	// Reset node primality in instances that disappeared
+	for i, j := range nodesHistory {
+		if _, exists := processedNums[j.num]; !exists {
+			j.isPrimary = false
+			nodesHistory[i] = j
+		}
+	}
+
 	/*
 
 	   read -r -a nodes <<<"$(tr ',;' ' ' <<<"${PGPOOL_BACKEND_NODES}")"
@@ -400,6 +410,8 @@ func configure(msg pgpoolConfigMessage, nodesHistory map[pgpoolNodeKey]pgpoolNod
 			switch j := j.(type) {
 			case int:
 				jS = strconv.Itoa(j)
+			case float32:
+				jS = strconv.FormatFloat(float64(j), 'f', 5, 32)
 			case string:
 				jS = "'" + j + "'"
 			case bool:
@@ -559,8 +571,8 @@ func parseForcedNodeConfig(value string) (pgpoolForcedConfig, error) {
 	}
 
 	if fieldsLen >= 4 {
-		if val, err := strconv.ParseInt(fields[3], 10, 32); err == nil {
-			result.weight = int(val)
+		if val, err := strconv.ParseFloat(fields[3], 32); err == nil {
+			result.weight = float32(val)
 		} else {
 			result.weight = -1
 		}