add keepalive

registry/registry.go

@@ -31,14 +31,15 @@ type zookeeperRegistry struct {
 	conn           *zk.Conn
 	authOpen       bool
 	user, password string
+	sessionTimeout time.Duration
 }
 
 func NewZookeeperRegistry(servers []string, sessionTimeout time.Duration) (registry.Registry, error) {
 	conn, _, err := zk.Connect(servers, sessionTimeout)
 	if err != nil {
 		return nil, err
 	}
-	return &zookeeperRegistry{conn: conn}, nil
+	return &zookeeperRegistry{conn: conn, sessionTimeout: sessionTimeout}, nil
 }
 
 func NewZookeeperRegistryWithAuth(servers []string, sessionTimeout time.Duration, user, password string) (registry.Registry, error) {
@@ -54,7 +55,7 @@ func NewZookeeperRegistryWithAuth(servers []string, sessionTimeout time.Duration
 	if err != nil {
 		return nil, err
 	}
-	return &zookeeperRegistry{conn: conn, authOpen: true, user: user, password: password}, nil
+	return &zookeeperRegistry{conn: conn, authOpen: true, user: user, password: password, sessionTimeout: sessionTimeout}, nil
 }
 
 func (z *zookeeperRegistry) Register(info *registry.Info) error {
@@ -66,11 +67,13 @@ func (z *zookeeperRegistry) Register(info *registry.Info) error {
 	if err != nil {
 		return err
 	}
-	return z.createNode(path, content, true)
+	err = z.createNode(path, content, true)
+	go z.keepalive(path, content)
+	return err
 }
 
-//  path format as follows:
-//  /{serviceName}/{ip}:{port}
+// path format as follows:
+// /{serviceName}/{ip}:{port}
 func buildPath(info *registry.Info) (string, error) {
 	var path string
 	if info == nil {
@@ -150,3 +153,42 @@ func (z *zookeeperRegistry) deleteNode(path string) error {
 	}
 	return nil
 }
+
+func (z *zookeeperRegistry) keepalive(path string, content []byte) {
+	sessionID := z.conn.SessionID()
+	ticker := time.NewTicker(z.sessionTimeout / 2)
+	defer ticker.Stop()
+	for range ticker.C {
+		cur := z.conn.SessionID()
+		if cur > 0 && sessionID != cur {
+			if err := z.ensureName(path, content, zk.FlagEphemeral); err == nil {
+				sessionID = cur
+			}
+		}
+	}
+}
+
+func (z *zookeeperRegistry) ensureName(path string, data []byte, flags int32) error {
+	exists, stat, err := z.conn.Exists(path)
+	if err != nil {
+		return err
+	}
+	if exists && flags&zk.FlagEphemeral == zk.FlagEphemeral {
+		err = z.conn.Delete(path, stat.Version)
+		if err != nil && err != zk.ErrNoNode {
+			return err
+		}
+		exists = false
+	}
+	if !exists {
+		if z.authOpen {
+			_, err = z.conn.Create(path, data, flags, zk.DigestACL(zk.PermAll, z.user, z.password))
+		} else {
+			_, err = z.conn.Create(path, data, flags, zk.WorldACL(zk.PermAll))
+		}
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}