extract method for get zookeeper connection

apache · Jun 2, 2020 · 9ae184f · 9ae184f
1 parent 1d1bbe0
commit 9ae184f
Showing 1 changed file with 20 additions and 25 deletions.
diff --git a/remoting/zookeeper/client.go b/remoting/zookeeper/client.go
@@ -416,15 +416,15 @@ func (z *ZookeeperClient) CreateWithValue(basePath string, value []byte) error {
 	)
 
 	logger.Debugf("zookeeperClient.Create(basePath{%s})", basePath)
+	conn := z.getConn()
+	err = errNilZkClientConn
+	if conn == nil {
+		return perrors.WithMessagef(err, "zk.Create(path:%s)", basePath)
+	}
+
 	for _, str := range strings.Split(basePath, "/")[1:] {
 		tmpPath = path.Join(tmpPath, "/", str)
-		err = errNilZkClientConn
-		z.RLock()
-		conn := z.Conn
-		z.RUnlock()
-		if conn != nil {
-			_, err = conn.Create(tmpPath, value, 0, zk.WorldACL(zk.PermAll))
-		}
+		_, err = conn.Create(tmpPath, value, 0, zk.WorldACL(zk.PermAll))
 
 		if err != nil {
 			if err == zk.ErrNodeExists {
@@ -446,9 +446,7 @@ func (z *ZookeeperClient) Delete(basePath string) error {
 	)
 
 	err = errNilZkClientConn
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		err = conn.Delete(basePath, -1)
 	}
@@ -468,9 +466,7 @@ func (z *ZookeeperClient) RegisterTemp(basePath string, node string) (string, er
 	err = errNilZkClientConn
 	data = []byte("")
 	zkPath = path.Join(basePath) + "/" + node
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		tmpPath, err = conn.Create(zkPath, data, zk.FlagEphemeral, zk.WorldACL(zk.PermAll))
 	}
@@ -493,9 +489,7 @@ func (z *ZookeeperClient) RegisterTempSeq(basePath string, data []byte) (string,
 	)
 
 	err = errNilZkClientConn
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		tmpPath, err = conn.Create(
 			path.Join(basePath)+"/",
@@ -526,9 +520,7 @@ func (z *ZookeeperClient) GetChildrenW(path string) ([]string, <-chan zk.Event,
 	)
 
 	err = errNilZkClientConn
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		children, stat, watcher, err = conn.ChildrenW(path)
 	}
@@ -562,9 +554,7 @@ func (z *ZookeeperClient) GetChildren(path string) ([]string, error) {
 	)
 
 	err = errNilZkClientConn
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		children, stat, err = conn.Children(path)
 	}
@@ -595,9 +585,7 @@ func (z *ZookeeperClient) ExistW(zkPath string) (<-chan zk.Event, error) {
 	)
 
 	err = errNilZkClientConn
-	z.RLock()
-	conn := z.Conn
-	z.RUnlock()
+	conn := z.getConn()
 	if conn != nil {
 		exist, _, watcher, err = conn.ExistsW(zkPath)
 	}
@@ -618,3 +606,10 @@ func (z *ZookeeperClient) ExistW(zkPath string) (<-chan zk.Event, error) {
 func (z *ZookeeperClient) GetContent(zkPath string) ([]byte, *zk.Stat, error) {
 	return z.Conn.Get(zkPath)
 }
+
+// getConn gets zookeeper connection safely
+func (z *ZookeeperClient) getConn() *zk.Conn {
+	z.RLock()
+	defer z.RUnlock()
+	return z.Conn
+}