Allow reconnection to a session from a saved session id and passwd

conn.go

@@ -99,6 +99,7 @@ type Conn struct {
 	watchers     map[watchPathType][]chan Event
 	watchersLock sync.Mutex
 	closeChan    chan struct{} // channel to tell send loop stop
+	sessionLock  sync.RWMutex
 
 	// Debug (used by unit tests)
 	reconnectLatch   chan struct{}
@@ -310,6 +311,15 @@ func WithMaxConnBufferSize(maxBufferSize int) connOption {
 	}
 }
 
+// WithSessionIdAndPasswd sets the session id and password to be used
+// for the first connection attempt
+func WithSessionIdAndPasswd(sessionID int64, passwd []byte) connOption {
+	return func(c *Conn) {
+		c.sessionID = sessionID
+		c.passwd = passwd
+	}
+}
+
 // Close will submit a close request with ZK and signal the connection to stop
 // sending and receiving packets.
 func (c *Conn) Close() {
@@ -330,7 +340,16 @@ func (c *Conn) State() State {
 
 // SessionID returns the current session id of the connection.
 func (c *Conn) SessionID() int64 {
-	return atomic.LoadInt64(&c.sessionID)
+	c.sessionLock.RLock()
+	defer c.sessionLock.RUnlock()
+	return c.sessionID
+}
+
+// SessionPassword returns the current session password of the connection.
+func (c *Conn) SessionPasswd() []byte {
+	c.sessionLock.RLock()
+	defer c.sessionLock.RUnlock()
+	return c.passwd
 }
 
 // SetLogger sets the logger to be used for printing errors.
@@ -637,7 +656,7 @@ func (c *Conn) authenticate() error {
 		LastZxidSeen:    c.lastZxid,
 		TimeOut:         c.sessionTimeoutMs,
 		SessionID:       c.SessionID(),
-		Passwd:          c.passwd,
+		Passwd:          c.SessionPasswd(),
 	})
 	if err != nil {
 		return err
@@ -676,16 +695,20 @@ func (c *Conn) authenticate() error {
 		return err
 	}
 	if r.SessionID == 0 {
-		atomic.StoreInt64(&c.sessionID, int64(0))
+		c.sessionLock.Lock()
+		c.sessionID = 0
 		c.passwd = emptyPassword
+		c.sessionLock.Unlock()
 		c.lastZxid = 0
 		c.setState(StateExpired)
 		return ErrSessionExpired
 	}
 
-	atomic.StoreInt64(&c.sessionID, r.SessionID)
-	c.setTimeouts(r.TimeOut)
+	c.sessionLock.Lock()
+	c.sessionID = r.SessionID
 	c.passwd = r.Passwd
+	c.sessionLock.Unlock()
+	c.setTimeouts(r.TimeOut)
 	c.setState(StateHasSession)
 
 	return nil

server_java_test.go

@@ -153,19 +153,19 @@ func (sc ServerConfig) Marshall(w io.Writer) error {
 
 // this is a helper to wait for the zk connection to at least get to the HasSession state
 func waitForSession(ctx context.Context, eventChan <-chan Event) error {
-	select {
-	case event, ok := <-eventChan:
-		// The eventChan is used solely to determine when the ZK conn has
-		// stopped.
-		if !ok {
-			return fmt.Errorf("connection closed before state reached")
+	for {
+		select {
+		case event, ok := <-eventChan:
+			// The eventChan is used solely to determine when the ZK conn has
+			// stopped.
+			if !ok {
+				return fmt.Errorf("connection closed before state reached")
+			}
+			if event.State == StateHasSession {
+				return nil
+			}
+		case <-ctx.Done():
+			return ctx.Err()
 		}
-		if event.State == StateHasSession {
-			return nil
-		}
-	case <-ctx.Done():
-		return ctx.Err()
 	}
-
-	return nil
 }

zk_test.go

@@ -1112,6 +1112,60 @@ func TestMaxBufferSize(t *testing.T) {
 	}
 }
 
+func TestSessionReconnectionFromSavedID(t *testing.T) {
+	ts, err := StartTestCluster(t, 1, nil, logWriter{t: t, p: "[ZKERR] "})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ts.Stop()
+
+	zk, eventChan, err := ts.ConnectAll()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer zk.Close()
+	// Prevent reconnection from this client
+	zk.reconnectLatch = make(chan struct{})
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+	err = waitForSession(ctx, eventChan)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sessionID := zk.SessionID()
+	passwd := zk.SessionPasswd()
+
+	// Simulate network error by brutally closing the network connection.
+	_ = zk.conn.Close()
+
+	// Re-establish the session from a different client with the session id and passwd
+	zk, eventChan, err = ts.ConnectWithOptions(15*time.Second, WithSessionIdAndPasswd(sessionID, passwd))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer zk.Close()
+
+	ctx, cancel = context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+	err = waitForSession(ctx, eventChan)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sessionID2 := zk.SessionID()
+	passwd2 := zk.SessionPasswd()
+
+	if sessionID != sessionID2 {
+		t.Fatalf("session id mismatch: %d %d", sessionID, sessionID2)
+	}
+
+	if !reflect.DeepEqual(passwd, passwd2) {
+		t.Fatalf("session password mismatch: %v %v", passwd, passwd2)
+	}
+}
+
 func startSlowProxy(t *testing.T, up, down Rate, upstream string, adj func(ln *Listener)) (string, chan bool, error) {
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {