From b1073c34284efc1b2dd86e6dd988f27c522187d8 Mon Sep 17 00:00:00 2001
From: Kezhu Wang <kezhuw@gmail.com>
Date: Wed, 22 Jun 2022 12:45:00 +0800
Subject: [PATCH] ZOOKEEPER-3996: Fix flaky
 ReadOnlyModeTest.testConnectionEvents

The same watcher was resued across different clients. It is hard to
know which event will last during verification, disconnected from old
client or connected from new client.

A brand new watcher solves this.
---
 .../org/apache/zookeeper/test/ClientBase.java | 24 ++++++++++++-------
 .../zookeeper/test/ReadOnlyModeTest.java      |  1 +
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ClientBase.java b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ClientBase.java
index e2257c1563c..cf35fdc54bd 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ClientBase.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ClientBase.java
@@ -124,6 +124,12 @@ public synchronized void process(WatchedEvent event) {
         public synchronized boolean isConnected() {
             return connected;
         }
+
+        protected synchronized String connectionDescription() {
+            return String.format("connected(%s), syncConnected(%s), readOnlyConnected(%s)",
+                    connected, syncConnected, readOnlyConnected);
+        }
+
         public synchronized void waitForConnected(long timeout) throws InterruptedException, TimeoutException {
             long expire = Time.currentElapsedTime() + timeout;
             long left = timeout;
@@ -132,8 +138,7 @@ public synchronized void waitForConnected(long timeout) throws InterruptedExcept
                 left = expire - Time.currentElapsedTime();
             }
             if (!connected) {
-                throw new TimeoutException("Failed to connect to ZooKeeper server.");
-
+                throw new TimeoutException("Failed to connect to ZooKeeper server: " + connectionDescription());
             }
         }
         public synchronized void waitForSyncConnected(long timeout) throws InterruptedException, TimeoutException {
@@ -144,18 +149,22 @@ public synchronized void waitForSyncConnected(long timeout) throws InterruptedEx
                 left = expire - Time.currentElapsedTime();
             }
             if (!syncConnected) {
-                throw new TimeoutException("Failed to connect to read-write ZooKeeper server.");
+                throw new TimeoutException(
+                    "Failed to connect to read-write ZooKeeper server: "
+                    + connectionDescription());
             }
         }
         public synchronized void waitForReadOnlyConnected(long timeout) throws InterruptedException, TimeoutException {
-            long expire = System.currentTimeMillis() + timeout;
+            long expire = Time.currentElapsedTime() + timeout;
             long left = timeout;
             while (!readOnlyConnected && left > 0) {
                 wait(left);
-                left = expire - System.currentTimeMillis();
+                left = expire - Time.currentElapsedTime();
             }
             if (!readOnlyConnected) {
-                throw new TimeoutException("Failed to connect in read-only mode to ZooKeeper server.");
+                throw new TimeoutException(
+                    "Failed to connect in read-only mode to ZooKeeper server: "
+                    + connectionDescription());
             }
         }
         public synchronized void waitForDisconnected(long timeout) throws InterruptedException, TimeoutException {
@@ -166,8 +175,7 @@ public synchronized void waitForDisconnected(long timeout) throws InterruptedExc
                 left = expire - Time.currentElapsedTime();
             }
             if (connected) {
-                throw new TimeoutException("Did not disconnect");
-
+                throw new TimeoutException("Did not disconnect: " + connectionDescription());
             }
         }
 
diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeTest.java
index 49f871c1ce4..521a1247d16 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeTest.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeTest.java
@@ -191,6 +191,7 @@ public void testConnectionEvents() throws Exception {
 
         // Re-connect the client (in case we were connected to the shut down
         // server and the local session was not persisted).
+        watcher = new CountdownWatcher();
         zk = new ZooKeeper(qu.getConnString(), CONNECTION_TIMEOUT, watcher, true);
         long start = Time.currentElapsedTime();
         while (!(zk.getState() == States.CONNECTEDREADONLY)) {