Update to use linkedin/go-zk (#2)

The primary driver behind this is to be able to make use of the new default host provider added in linkedin/go-zk#6 to resolve connection issues when zookeeper host IPs change.
bugsnag · Jul 8, 2024 · 3d1ceef · 3d1ceef
1 parent b0398fc
commit 3d1ceef
Show file tree

Hide file tree

Showing 9 changed files with 19 additions and 28 deletions.
diff --git a/NOTICE b/NOTICE
@@ -37,7 +37,7 @@ This product includes/uses eapache/queue (https://github.com/eapache/queue/)
 Copyright (c) 2014 Evan Huus
 License: MIT
 
-This product includes/uses go-zookeeper (https://github.com/samuel/go-zookeeper/)
+This product includes/uses go-zk (https://github.com/linkedin/go-zk/)
 Copyright (c) 2013, Samuel Stauffer <[email protected]>
 License: BSD
 

diff --git a/core/internal/consumer/kafka_zk_client.go b/core/internal/consumer/kafka_zk_client.go
@@ -15,7 +15,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 

diff --git a/core/internal/consumer/kafka_zk_test.go b/core/internal/consumer/kafka_zk_test.go
@@ -17,7 +17,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 

diff --git a/core/internal/helpers/zookeeper.go b/core/internal/helpers/zookeeper.go
@@ -12,14 +12,13 @@ package helpers
 import (
 	"crypto/tls"
 	"crypto/x509"
-	"net"
 	"os"
 	"time"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/viper"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap"
 
@@ -36,12 +35,7 @@ type BurrowZookeeperClient struct {
 // timeout it's possible to reestablish a connection to a different server and keep the same session. This is means any
 // ephemeral nodes and watches are maintained.
 func ZookeeperConnect(servers []string, sessionTimeout time.Duration, logger *zap.Logger) (protocol.ZookeeperClient, <-chan zk.Event, error) {
-	// We need a function to set the logger for the ZK connection
-	zkSetLogger := func(c *zk.Conn) {
-		c.SetLogger(zap.NewStdLog(logger))
-	}
-
-	zkconn, connEventChan, err := zk.Connect(servers, sessionTimeout, zkSetLogger)
+	zkconn, connEventChan, err := zk.Connect(servers, sessionTimeout)
 	return &BurrowZookeeperClient{client: zkconn}, connEventChan, err
 }
 
@@ -57,23 +51,18 @@ func ZookeeperConnectTLS(servers []string, sessionTimeout time.Duration, logger
 
 	logger.Info("starting zookeeper (TLS)", zap.String("caFile", caFile), zap.String("certFile", certFile), zap.String("keyFile", keyFile))
 
-	dialer, err := newTLSDialer(servers[0], caFile, certFile, keyFile)
+	dialer, err := newTLSDialer(caFile, certFile, keyFile)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	// We need a function to set the logger for the ZK connection
-	zkSetLogger := func(c *zk.Conn) {
-		c.SetLogger(zap.NewStdLog(logger))
-	}
-
-	zkconn, connEventChan, err := zk.Connect(servers, sessionTimeout, zk.WithDialer(dialer), zkSetLogger)
+	zkconn, connEventChan, err := zk.Connect(servers, sessionTimeout, zk.WithDialer(dialer))
 	return &BurrowZookeeperClient{client: zkconn}, connEventChan, err
 }
 
 // newTLSDialer creates a dialer with TLS configured. It will install caFile as root CA and if both certFile and keyFile are
 // set, it will add those as a certificate.
-func newTLSDialer(addr, caFile, certFile, keyFile string) (zk.Dialer, error) {
+func newTLSDialer(caFile, certFile, keyFile string) (zk.Dialer, error) {
 	caCert, err := os.ReadFile(caFile)
 	if err != nil {
 		return nil, errors.New("could not read caFile: " + err.Error())
@@ -96,9 +85,11 @@ func newTLSDialer(addr, caFile, certFile, keyFile string) (zk.Dialer, error) {
 		tlsConfig.Certificates = []tls.Certificate{cert}
 	}
 
-	return func(string, string, time.Duration) (net.Conn, error) {
-		return tls.Dial("tcp", addr, tlsConfig)
-	}, nil
+	tlsDialer := &tls.Dialer{
+		Config: tlsConfig,
+	}
+
+	return tlsDialer, nil
 }
 
 // Close shuts down the connection to the Zookeeper ensemble.

diff --git a/core/internal/zookeeper/coordinator.go b/core/internal/zookeeper/coordinator.go
@@ -18,7 +18,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 

diff --git a/core/internal/zookeeper/coordinator_test.go b/core/internal/zookeeper/coordinator_test.go
@@ -20,7 +20,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 

diff --git a/core/protocol/protocol.go b/core/protocol/protocol.go
@@ -16,7 +16,7 @@ package protocol
 import (
 	"sync"
 
-	"github.com/samuel/go-zookeeper/zk"
+	"github.com/linkedin/go-zk"
 	"go.uber.org/zap"
 )
 

diff --git a/go.mod b/go.mod
@@ -7,10 +7,10 @@ require (
 	github.com/OneOfOne/xxhash v1.2.8
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/karrick/goswarm v1.10.0
+	github.com/linkedin/go-zk v0.1.4
 	github.com/pborman/uuid v1.2.1
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.19.0
-	github.com/samuel/go-zookeeper v0.0.0-20201211165307-7117e9ea2414
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.9.0
 	github.com/xdg/scram v1.0.5

diff --git a/go.sum b/go.sum
@@ -63,6 +63,8 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/linkedin/go-zk v0.1.4 h1:ZB/u/DaNbHUiuymbtD6C0Bf6s+4O3J36Wqd1Txkztig=
+github.com/linkedin/go-zk v0.1.4/go.mod h1:X1Id+YYjM0pt6UHVD0eIUrLkhFL/0rAUn+cvc2EDwhg=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
@@ -94,8 +96,6 @@ github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6ke
 github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4=
 github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE=
 github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ=
-github.com/samuel/go-zookeeper v0.0.0-20201211165307-7117e9ea2414 h1:AJNDS0kP60X8wwWFvbLPwDuojxubj9pbfK7pjHw0vKg=
-github.com/samuel/go-zookeeper v0.0.0-20201211165307-7117e9ea2414/go.mod h1:gi+0XIa01GRL2eRQVjQkKGqKF3SF9vZR/HnPullcV2E=
 github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
 github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=