From e768a01099fa7ecae301c84262dcc0147a8a1cc3 Mon Sep 17 00:00:00 2001
From: Mike Drob <mdrob@apache.org>
Date: Wed, 8 Jul 2015 11:04:30 -0500
Subject: [PATCH] CURATOR-218 Reorder ConnectionState process event

Address a race condition in ConnectionState.process where it will
trigger watchers first before updating its own state. This can lead to
inconsistencies when blocking until connected.
---
 .../org/apache/curator/ConnectionState.java   | 26 +++++++++----------
 .../imps/TestBlockUntilConnected.java         | 24 +++++++++++++++++
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/curator-client/src/main/java/org/apache/curator/ConnectionState.java b/curator-client/src/main/java/org/apache/curator/ConnectionState.java
index d3900a1ca4..46ae9fd30a 100644
--- a/curator-client/src/main/java/org/apache/curator/ConnectionState.java
+++ b/curator-client/src/main/java/org/apache/curator/ConnectionState.java
@@ -41,7 +41,7 @@ class ConnectionState implements Watcher, Closeable
 {
     private static final int MAX_BACKGROUND_EXCEPTIONS = 10;
     private static final boolean LOG_EVENTS = Boolean.getBoolean(DebugUtils.PROPERTY_LOG_EVENTS);
-    private final Logger log = LoggerFactory.getLogger(getClass());
+    private static final Logger log = LoggerFactory.getLogger(ConnectionState.class);
     private final HandleHolder zooKeeper;
     private final AtomicBoolean isConnected = new AtomicBoolean(false);
     private final EnsembleProvider ensembleProvider;
@@ -145,24 +145,22 @@ public void process(WatchedEvent event)
             log.debug("ConnectState watcher: " + event);
         }
 
-        for ( Watcher parentWatcher : parentWatchers )
-        {
-            TimeTrace timeTrace = new TimeTrace("connection-state-parent-process", tracer.get());
-            parentWatcher.process(event);
-            timeTrace.commit();
-        }
-
-        boolean wasConnected = isConnected.get();
-        boolean newIsConnected = wasConnected;
         if ( event.getType() == Watcher.Event.EventType.None )
         {
-            newIsConnected = checkState(event.getState(), wasConnected);
+            boolean wasConnected = isConnected.get();
+            boolean newIsConnected = checkState(event.getState(), wasConnected);
+            if ( newIsConnected != wasConnected )
+            {
+                isConnected.set(newIsConnected);
+                connectionStartMs = System.currentTimeMillis();
+            }
         }
 
-        if ( newIsConnected != wasConnected )
+        for ( Watcher parentWatcher : parentWatchers )
         {
-            isConnected.set(newIsConnected);
-            connectionStartMs = System.currentTimeMillis();
+            TimeTrace timeTrace = new TimeTrace("connection-state-parent-process", tracer.get());
+            parentWatcher.process(event);
+            timeTrace.commit();
         }
     }
 
diff --git a/curator-framework/src/test/java/org/apache/curator/framework/imps/TestBlockUntilConnected.java b/curator-framework/src/test/java/org/apache/curator/framework/imps/TestBlockUntilConnected.java
index f649afb0cf..a6dc7abf7d 100644
--- a/curator-framework/src/test/java/org/apache/curator/framework/imps/TestBlockUntilConnected.java
+++ b/curator-framework/src/test/java/org/apache/curator/framework/imps/TestBlockUntilConnected.java
@@ -232,4 +232,28 @@ public void run()
             CloseableUtils.closeQuietly(client);
         }
     }
+
+    /**
+     * Test that we are actually connected every time that we block until connection is established in a tight loop.
+     */
+    @Test
+    public void testBlockUntilConnectedTightLoop() throws InterruptedException
+    {
+        CuratorFramework client;
+        for(int i = 0 ; i < 50 ; i++)
+        {
+            client = CuratorFrameworkFactory.newClient(server.getConnectString(), new RetryOneTime(100));
+            try
+            {
+                client.start();
+                client.blockUntilConnected();
+
+                Assert.assertTrue(client.getZookeeperClient().isConnected(), "Not connected after blocking for connection #" + i);
+            }
+            finally
+            {
+                client.close();
+            }
+        }
+    }
 }