Refactor session state management

1 files changed, 19 insertions, 17 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
index 273d5d89ac2..b9293abdee3 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
@@ -377,6 +377,12 @@ public class MessageBusVisitorSession implements VisitorSession {
         }
     }
 
+    private void updateStateUnlessAlreadyFailed(StateDescription newState) {
+        if (!state.failed()) {
+            state = newState;
+        } // else: don't override existing failure state
+    }
+
     /**
      * Attempt to transition to a new state. Depending on the current state,
      * some transitions may be disallowed, such as transitioning from ABORTED
@@ -391,25 +397,21 @@ public class MessageBusVisitorSession implements VisitorSession {
      */
     private StateDescription transitionTo(StateDescription newState) {
         log.log(LogLevel.DEBUG, sessionName + ": attempting transition to state " + newState);
-        if (newState.getState() == State.WORKING) {
-            assert(state.getState() == State.NOT_STARTED);
-            state = newState;
-        } else if (newState.getState() == State.COMPLETED) {
-            if (!state.failed()) {
+        switch (newState.getState()) {
+            case WORKING:
+                assert(state.getState() == State.NOT_STARTED);
                 state = newState;
-            } // else: don't override existing failure state
-        } else if (newState.getState() == State.ABORTED) {
-            state = newState;
-        } else if (newState.getState() == State.FAILED) {
-            if (state.getState() != State.FAILED) {
-                state = newState;
-            } // else: don't override existing failure state
-        } else if (newState.getState() == State.TIMED_OUT) {
-            if (!state.failed()) {
+                break;
+            case ABORTED:
                 state = newState;
-            } // else: don't override existing failure state
-        } else {
-            assert(false);
+                break;
+            case COMPLETED:
+            case FAILED:
+            case TIMED_OUT:
+                updateStateUnlessAlreadyFailed(newState);
+                break;
+            default:
+                assert(false);
         }
         log.log(LogLevel.DEBUG, "Session '" + sessionName + "' is now in state " +  state);
         return state;