Better zk connection error handling

- Modifying concepts around DistributedStateMachine and
  ZookeeperStateMachineEnsemble to get better handling when
  zk connection is lost.
- New jepsen test showing that after a brain split machine
  will get back to a consistent state.
- Polish javadocs
- Hopefully fixes #91

jepsen/spring-statemachine-jepsen/src/spring_statemachine_jepsen/core.clj

@@ -272,6 +272,48 @@
     )
   )
 
+(defn event-gen-4
+  "Generates event and checks states while splitting network"
+  []
+
+  (gen/phases
+    (gen/clients
+      (gen/each
+        (gen/once {:type :invoke
+                   :f    :status})))
+
+    (gen/nemesis
+      (gen/once {:type :info :f :start}))
+
+    (gen/sleep 30)
+
+    (gen/clients
+      (gen/each
+        (gen/once {:type :invoke
+                   :f    :status})))
+
+    (gen/nemesis
+      (gen/once {:type :info :f :stop}))
+
+    (gen/clients
+      (gen/each
+        (gen/once {:type :invoke
+                   :f    :status})))
+
+    (gen/clients
+      (gen/once {:type :invoke
+                 :f    :event
+                 :e    "C"}))
+
+    (gen/clients
+      (gen/each
+        (gen/once {:type :invoke
+                   :f    :states
+                   :s    ["S0","S2","S21","S211"]})))
+
+    )
+  )
+
 (defn statemachine-test
   "Defaults for testing state machine."
   [name opts]
@@ -309,3 +351,10 @@
   (event-test "send-isolated-event-with-variable"
                {:nemesis   nemesis/noop
                 :generator (event-gen-3)}))
+
+(defn partition-half-test
+  "Does a half brain split and checks that machines are healing."
+  []
+  (event-test "partition-half"
+               {:nemesis   (nemesis/partition-random-halves)
+                :generator (event-gen-4)}))

jepsen/spring-statemachine-jepsen/test/spring_statemachine_jepsen/core_test.clj

@@ -22,3 +22,6 @@
 
 (deftest send-isolated-event-with-variable
   (run-statemachine-test! (send-isolated-event-with-variable-test)))
+
+(deftest partition-half
+  (run-statemachine-test! (partition-half-test)))

spring-statemachine-core/src/main/java/org/springframework/statemachine/ensemble/DistributedStateMachine.java

@@ -247,7 +247,12 @@ public class DistributedStateMachine<S, E> extends LifecycleObjectSupport implem
 
 		@Override
 		public void stateMachineJoined(final StateMachine<S, E> stateMachine, final StateMachineContext<S, E> context) {
+			if (log.isDebugEnabled()) {
+				log.debug("Event stateMachineJoined stateMachine=[" + stateMachine + "] context=[" + context + "]");
+			}
 			if (stateMachine != null && stateMachine == DistributedStateMachine.this) {
+				delegate.stop();
+				setStateMachineError(null);
 				if (context != null) {
 					// I'm now successfully joined, so set delegating
 					// sm to current known state by a context.

spring-statemachine-core/src/main/java/org/springframework/statemachine/ensemble/EnsembleListeger.java

@@ -29,7 +29,13 @@ import org.springframework.statemachine.StateMachineContext;
 public interface EnsembleListeger<S, E> {
 
 	/**
-	 * Called when state machine joined an ensemble.
+	 * Called when state machine joined an ensemble. This callback
+	 * is guaranteed to be called for a {@link StateMachine} who
+	 * requested a join. User of this listener should check that a
+	 * {@link StateMachine} is the one interested of. Implementation
+	 * may choose to notify other {@link StateMachine} joins if it is
+	 * able to do so. This may be called multiple time in case ensemble
+	 * has made a choice to leave machine due to ensemble errors.
 	 *
 	 * @param stateMachine the state machine
 	 * @param context the state machine context
@@ -37,7 +43,12 @@ public interface EnsembleListeger<S, E> {
 	void stateMachineJoined(StateMachine<S, E> stateMachine, StateMachineContext<S, E> context);
 
 	/**
-	 * Called when state machine left an ensemble.
+	 * Called when state machine left an ensemble. This callback
+	 * is guaranteed to be called for a {@link StateMachine} who
+	 * requested a leave. User of this listener should check that a
+	 * {@link StateMachine} is the one interested of. Implementation
+	 * may choose to notify other {@link StateMachine} leaves if it is
+	 * able to do so.
 	 *
 	 * @param stateMachine the state machine
 	 * @param context the state machine context

spring-statemachine-core/src/main/java/org/springframework/statemachine/ensemble/StateMachineEnsemble.java

@@ -30,14 +30,23 @@ import org.springframework.statemachine.StateMachineContext;
 public interface StateMachineEnsemble<S, E> {
 
 	/**
-	 * Request a join to a state machine ensemble.
+	 * Request a join to a state machine ensemble. This method
+	 * is a request to join an ensemble and doesn't guarantee
+	 * a requester will eventually successfully join. Join operation
+	 * needs to be used together with {@link EnsembleListeger} and
+	 * {@link EnsembleListeger#stateMachineJoined(StateMachine, StateMachineContext)}
+	 * is called with a {@link StateMachine} instance for successful join.
 	 *
 	 * @param stateMachine the state machine
 	 */
 	void join(StateMachine<S, E> stateMachine);
 
 	/**
-	 * Request a leave from an ensemble.
+	 * Request a leave from an ensemble. This method is a request to
+	 * leave an ensemble. After this method is called no further processing
+	 * is done for a instance of {@link StateMachine}. Additionally
+	 * {@link EnsembleListeger#stateMachineLeft(StateMachine, StateMachineContext)}
+	 * is called when leave request is fully processed.
 	 *
 	 * @param stateMachine the state machine
 	 */
@@ -58,16 +67,16 @@ public interface StateMachineEnsemble<S, E> {
 	void removeEnsembleListener(EnsembleListeger<S, E> listener);
 
 	/**
-	 * Sets the state.
+	 * Sets the state as a {@link StateMachineContext}.
 	 *
-	 * @param context the context
+	 * @param context the state machine context
 	 */
 	void setState(StateMachineContext<S, E> context);
 
 	/**
-	 * Gets the state.
+	 * Gets the state as a {@link StateMachineContext}.
 	 *
-	 * @return the state
+	 * @return the state machine context
 	 */
 	StateMachineContext<S, E> getState();
 

spring-statemachine-zookeeper/src/main/java/org/springframework/statemachine/zookeeper/ZookeeperStateMachineEnsemble.java

@@ -16,6 +16,8 @@
 package org.springframework.statemachine.zookeeper;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Queue;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -32,6 +34,8 @@ import org.apache.curator.framework.imps.CuratorFrameworkState;
 import org.apache.curator.framework.recipes.locks.InterProcessSemaphoreMutex;
 import org.apache.curator.framework.recipes.nodes.PersistentEphemeralNode;
 import org.apache.curator.framework.recipes.nodes.PersistentEphemeralNode.Mode;
+import org.apache.curator.framework.state.ConnectionState;
+import org.apache.curator.framework.state.ConnectionStateListener;
 import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.data.Stat;
 import org.springframework.statemachine.StateMachine;
@@ -73,6 +77,9 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 	private final CuratorWatcher watcher = new StateWatcher();
 	private PersistentEphemeralNode node;
 	private final Queue<StateMachine<S, E>> joinQueue = new ConcurrentLinkedQueue<StateMachine<S, E>>();
+	private final List<StateMachine<S, E>> joined = new ArrayList<StateMachine<S,E>>();
+	private final Object joinLock = new Object();
+	private final ConnectionStateListener connectionListener = new LocalConnectionStateListener();
 
 	/**
 	 * Instantiates a new zookeeper state machine ensemble.
@@ -126,7 +133,7 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 				log.error("Error reading current state during start", e);
 			}
 		}
-		joinQueued();
+		curatorClient.getConnectionStateListenable().addListener(connectionListener);
 	}
 
 	@Override
@@ -139,6 +146,7 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 				node = null;
 			}
 		}
+		curatorClient.getConnectionStateListenable().removeListener(connectionListener);
 	}
 
 	@Override
@@ -147,28 +155,57 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 			joinQueue.add(stateMachine);
 		} else {
 			StateWrapper stateWrapper = stateRef.get();
+			synchronized (joinLock) {
+				joined.add(stateMachine);
+			}
 			notifyJoined(stateMachine, stateWrapper != null ? stateWrapper.context : null);
 		}
 	}
 
 	private void joinQueued() {
-		StateWrapper stateWrapper = stateRef.get();
 		StateMachine<S, E> stateMachine = null;
-		while ((stateMachine = joinQueue.poll()) != null) {
-			notifyJoined(stateMachine, stateWrapper != null ? stateWrapper.context : null);
+		synchronized (joinLock) {
+			while ((stateMachine = joinQueue.poll()) != null) {
+				joined.add(stateMachine);
+			}
+		}
+	}
+
+	private void notifyJoined() {
+		StateWrapper stateWrapper = stateRef.get();
+		synchronized (joinLock) {
+			for (StateMachine<S, E> stateMachine : joined) {
+				notifyJoined(stateMachine, stateWrapper != null ? stateWrapper.context : null);
+			}
+		}
+	}
+
+	private void notifyLeft() {
+		StateWrapper stateWrapper = stateRef.get();
+		synchronized (joinLock) {
+			for (StateMachine<S, E> stateMachine : joined) {
+				notifyLeft(stateMachine, stateWrapper != null ? stateWrapper.context : null);
+			}
 		}
 	}
 
 	@Override
 	public void leave(StateMachine<S, E> stateMachine) {
+		// TODO: think when to close
 		if (node != null) {
 			try {
 				node.close();
 			} catch (IOException e) {
 			}
 		}
-		StateWrapper stateWrapper = stateRef.get();
-		notifyLeft(stateMachine, stateWrapper != null ? stateWrapper.context : null);
+		boolean removed = false;
+		synchronized (joinLock) {
+			removed = joined.remove(stateMachine);
+		}
+		if (removed) {
+			StateWrapper stateWrapper = stateRef.get();
+			notifyLeft(stateMachine, stateWrapper != null ? stateWrapper.context : null);
+		}
 	}
 
 	@Override
@@ -197,6 +234,19 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 		return readCurrentContext().context;
 	}
 
+	private void handleZkConnect() {
+		log.info("Handling Zookeeper connect");
+		joinQueued();
+		notifyJoined();
+		registerWatcherForStatePath();
+	}
+
+	private void handleZkDisconnect() {
+		log.info("Handling Zookeeper disconnect");
+		notifyError(new StateMachineEnsembleException("Lost connection to zookeeper"));
+		notifyLeft();
+	}
+
 	private StateWrapper readCurrentContext() {
 		try {
 			Stat stat = new Stat();
@@ -271,7 +321,7 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 				curatorClient.checkExists().usingWatcher(watcher).forPath(statePath);
 			}
 		} catch (Exception e) {
-			log.warn("Registering wacher for path " + statePath + " threw error", e);
+			log.warn("Registering watcher for path " + statePath + " threw error", e);
 		}
 	}
 
@@ -398,6 +448,30 @@ public class ZookeeperStateMachineEnsemble<S, E> extends StateMachineEnsembleObj
 		return false;
 	}
 
+	private class LocalConnectionStateListener implements ConnectionStateListener {
+
+		@Override
+		public void stateChanged(CuratorFramework client, ConnectionState newState) {
+			if (curatorClient == client) {
+				switch (newState) {
+				case CONNECTED:
+				case RECONNECTED:
+					handleZkConnect();
+					break;
+				case READ_ONLY:
+					break;
+				case LOST:
+				case SUSPENDED:
+					handleZkDisconnect();
+					break;
+				default:
+					break;
+				}
+			}
+		}
+
+	};
+
 	/**
 	 * Wrapper object for a {@link StateMachineContext} and its
 	 * current version.

spring-statemachine-zookeeper/src/test/java/org/springframework/statemachine/zookeeper/ZookeeperStateMachineTests.java

@@ -431,6 +431,149 @@ public class ZookeeperStateMachineTests extends AbstractZookeeperTests {
 		assertThat(machine2.getState().getIds(), containsInAnyOrder("S1"));
 	}
 
+	@Test
+	public void testConnectionLoss1() throws Exception {
+		context.register(ZkServerConfig.class, BaseConfig.class);
+		context.refresh();
+
+		CuratorFramework curatorClient =
+				context.getBean("curatorClient", CuratorFramework.class);
+
+		StateMachine<String, String> machine1 =
+				buildTestStateMachine(curatorClient);
+
+		StateMachineTestPlan<String, String> plan1 =
+				StateMachineTestPlanBuilder.<String, String>builder()
+					.defaultAwaitTime(2)
+					.stateMachine(machine1)
+					.step()
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.step()
+						.sendEvent("C", machine1)
+						.expectStateChanged(3)
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.build();
+		plan1.test();
+
+		Object ensemble = TestUtils.readField("ensemble", machine1);
+		try {
+			TestUtils.callMethod("handleZkDisconnect", ensemble);
+		} catch (Exception e) {
+		}
+		TestUtils.callMethod("handleZkConnect", ensemble);
+
+		StateMachineTestPlan<String, String> plan2 =
+				StateMachineTestPlanBuilder.<String, String>builder()
+					.defaultAwaitTime(2)
+					.stateMachine(machine1)
+					.step()
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.step()
+						.sendEvent("C", machine1)
+						.expectStateChanged(2)
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.build();
+		plan2.test();
+	}
+
+	@Test
+	public void testConnectionLoss2() throws Exception {
+		context.register(ZkServerConfig.class, BaseConfig.class);
+		context.refresh();
+
+		CuratorFramework curatorClient =
+				context.getBean("curatorClient", CuratorFramework.class);
+
+		StateMachine<String, String> machine1 =
+				buildTestStateMachine(curatorClient);
+		StateMachine<String, String> machine2 =
+				buildTestStateMachine(curatorClient);
+		StateMachine<String, String> machine3 =
+				buildTestStateMachine(curatorClient);
+		StateMachine<String, String> machine4 =
+				buildTestStateMachine(curatorClient);
+		StateMachine<String, String> machine5 =
+				buildTestStateMachine(curatorClient);
+
+		StateMachineTestPlan<String, String> plan1 =
+				StateMachineTestPlanBuilder.<String, String>builder()
+					.defaultAwaitTime(2)
+					.stateMachine(machine1)
+					.stateMachine(machine2)
+					.stateMachine(machine3)
+					.stateMachine(machine4)
+					.stateMachine(machine5)
+					.step()
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.step()
+						.sendEvent("C", machine1)
+						.expectStateChanged(3)
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.build();
+		plan1.test();
+
+		Object ensemble2 = TestUtils.readField("ensemble", machine2);
+		Object ensemble3 = TestUtils.readField("ensemble", machine3);
+		Object ensemble4 = TestUtils.readField("ensemble", machine4);
+		Object ensemble5 = TestUtils.readField("ensemble", machine5);
+		try {
+			TestUtils.callMethod("handleZkDisconnect", ensemble2);
+			TestUtils.callMethod("handleZkDisconnect", ensemble3);
+			TestUtils.callMethod("handleZkDisconnect", ensemble4);
+			TestUtils.callMethod("handleZkDisconnect", ensemble5);
+		} catch (Exception e) {
+		}
+		TestUtils.callMethod("handleZkConnect", ensemble2);
+		TestUtils.callMethod("handleZkConnect", ensemble3);
+		TestUtils.callMethod("handleZkConnect", ensemble4);
+		TestUtils.callMethod("handleZkConnect", ensemble5);
+
+		StateMachineTestPlan<String, String> plan2 =
+				StateMachineTestPlanBuilder.<String, String>builder()
+					.defaultAwaitTime(2)
+					.stateMachine(machine1)
+					.stateMachine(machine2)
+					.stateMachine(machine3)
+					.stateMachine(machine4)
+					.stateMachine(machine5)
+					.step()
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.step()
+						.sendEvent("C", machine1)
+						.expectStateChanged(2)
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.step()
+						.sendEvent("C", machine2)
+						.expectStateChanged(3)
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.step()
+						.sendEvent("C", machine3)
+						.expectStateChanged(2)
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.step()
+						.sendEvent("C", machine4)
+						.expectStateChanged(3)
+						.expectStates("S0", "S2", "S21", "S211")
+						.and()
+					.step()
+						.sendEvent("C", machine5)
+						.expectStateChanged(2)
+						.expectStates("S0", "S1", "S11")
+						.and()
+					.build();
+		plan2.test();
+	}
+
 	@Configuration
 	@EnableStateMachine(name = "sm1")
 	static class Config1 extends SharedConfig1 {