Support arbitrary ZooKeeper config lines (vitessio#13829)

Signed-off-by: Adam Saponara <[email protected]> Signed-off-by: Matt Lord <[email protected]> Co-authored-by: Matt Lord <[email protected]>
slackhq · Sep 19, 2023 · c341b32 · c341b32
1 parent b102fee
commit c341b32
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 6 deletions.
diff --git a/go/cmd/zkctl/zkctl.go b/go/cmd/zkctl/zkctl.go
@@ -34,17 +34,20 @@ Commands:
 `
 
 var (
-	zkCfg = "6@<hostname>:3801:3802:3803"
-	myID  uint
+	zkCfg   = "6@<hostname>:3801:3802:3803"
+	myID    uint
+	zkExtra []string
 )
 
 func registerZkctlFlags(fs *pflag.FlagSet) {
 	fs.StringVar(&zkCfg, "zk.cfg", zkCfg,
 		"zkid@server1:leaderPort1:electionPort1:clientPort1,...)")
 	fs.UintVar(&myID, "zk.myid", myID,
 		"which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname")
-
+	fs.StringArrayVar(&zkExtra, "zk.extra", zkExtra,
+		"extra config line(s) to append verbatim to config (flag can be specified more than once)")
 }
+
 func init() {
 	servenv.OnParse(registerZkctlFlags)
 }
@@ -59,6 +62,7 @@ func main() {
 	args := servenv.ParseFlagsWithArgs("zkctl")
 
 	zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID))
+	zkConfig.Extra = zkExtra
 	zkd := zkctl.NewZkd(zkConfig)
 
 	action := args[0]

diff --git a/go/cmd/zkctld/zkctld.go b/go/cmd/zkctld/zkctld.go
@@ -35,8 +35,9 @@ import (
 )
 
 var (
-	zkCfg = "6@<hostname>:3801:3802:3803"
-	myID  uint
+	zkCfg   = "6@<hostname>:3801:3802:3803"
+	myID    uint
+	zkExtra []string
 )
 
 func init() {
@@ -48,7 +49,8 @@ func registerFlags(fs *pflag.FlagSet) {
 		"zkid@server1:leaderPort1:electionPort1:clientPort1,...)")
 	fs.UintVar(&myID, "zk.myid", myID,
 		"which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname")
-
+	fs.StringArrayVar(&zkExtra, "zk.extra", zkExtra,
+		"extra config line(s) to append verbatim to config (flag can be specified more than once)")
 	acl.RegisterFlags(fs)
 }
 
@@ -59,6 +61,7 @@ func main() {
 	servenv.ParseFlags("zkctld")
 	servenv.Init()
 	zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID))
+	zkConfig.Extra = zkExtra
 	zkd := zkctl.NewZkd(zkConfig)
 
 	if zkd.Inited() {

diff --git a/go/flags/endtoend/zkctl.txt b/go/flags/endtoend/zkctl.txt
@@ -21,4 +21,5 @@ Usage of zkctl:
   -v, --version                                                     print binary version
       --vmodule moduleSpec                                          comma-separated list of pattern=N settings for file-filtered logging
       --zk.cfg string                                               zkid@server1:leaderPort1:electionPort1:clientPort1,...) (default "6@<hostname>:3801:3802:3803")
+      --zk.extra stringArray                                        extra config line(s) to append verbatim to config (flag can be specified more than once)
       --zk.myid uint                                                which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname
diff --git a/go/flags/endtoend/zkctld.txt b/go/flags/endtoend/zkctld.txt
@@ -22,4 +22,5 @@ Usage of zkctld:
   -v, --version                                                     print binary version
       --vmodule moduleSpec                                          comma-separated list of pattern=N settings for file-filtered logging
       --zk.cfg string                                               zkid@server1:leaderPort1:electionPort1:clientPort1,...) (default "6@<hostname>:3801:3802:3803")
+      --zk.extra stringArray                                        extra config line(s) to append verbatim to config (flag can be specified more than once)
       --zk.myid uint                                                which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname
diff --git a/go/vt/zkctl/zkconf.go b/go/vt/zkctl/zkconf.go
@@ -49,6 +49,7 @@ type ZkConfig struct {
 	ServerId   uint32 // nolint:revive
 	ClientPort int
 	Servers    []zkServerAddr
+	Extra      []string
 	Global     bool
 }
 
@@ -117,6 +118,11 @@ func MakeZooCfg(cnfFiles []string, cnf *ZkConfig, header string) (string, error)
 		return "", dataErr
 	}
 
+	myTemplateSource.WriteString("\n") // in case `data` did not end with a newline
+	for _, extra := range cnf.Extra {
+		myTemplateSource.WriteString(fmt.Sprintf("%s\n", extra))
+	}
+
 	myTemplate, err := template.New("foo").Parse(myTemplateSource.String())
 	if err != nil {
 		return "", err

diff --git a/go/vt/zkctl/zkctl_test.go b/go/vt/zkctl/zkctl_test.go
@@ -17,6 +17,8 @@ limitations under the License.
 package zkctl
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 )
 
@@ -33,6 +35,15 @@ func TestLifeCycle(t *testing.T) {
 	myID := 255
 
 	zkConf := MakeZkConfigFromString(config, uint32(myID))
+	zkExtraConfLine := "tcpKeepAlive=true"
+	zkConf.Extra = []string{zkExtraConfLine}
+
+	if zkObservedConf, err := MakeZooCfg([]string{zkConf.ConfigFile()}, zkConf, "header"); err != nil {
+		t.Fatalf("MakeZooCfg err: %v", err)
+	} else if !strings.Contains(string(zkObservedConf), fmt.Sprintf("\n%s\n", zkExtraConfLine)) {
+		t.Fatalf("Expected zkExtraConfLine in zkObservedConf")
+	}
+
 	zkd := NewZkd(zkConf)
 	if err := zkd.Init(); err != nil {
 		t.Fatalf("Init() err: %v", err)
@@ -49,4 +60,5 @@ func TestLifeCycle(t *testing.T) {
 	if err := zkd.Teardown(); err != nil {
 		t.Fatalf("Teardown() err: %v", err)
 	}
+
 }