Persist state machine states in snapshots.

atomix · Feb 3, 2015 · e58fe12 · e58fe12
1 parent 0207a6e
commit e58fe12
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/state-machine/src/main/java/net/kuujo/copycat/state/internal/DefaultStateMachine.java b/state-machine/src/main/java/net/kuujo/copycat/state/internal/DefaultStateMachine.java
@@ -106,14 +106,29 @@ public <U> U createProxy(Class<U> type) {
    * Takes a snapshot of the state machine state.
    */
   private Map<String, Object> snapshot() {
-    return data;
+    Map<String, Object> snapshot = new HashMap<>(2);
+    snapshot.put("state", state.getClass().getName());
+    snapshot.put("data", data);
+    return snapshot;
   }
 
   /**
    * Installs a snapshot of the state machine state.
    */
+  @SuppressWarnings("unchecked")
   private void install(Map<String, Object> snapshot) {
-    this.data = snapshot;
+    Object stateClassName = snapshot.get("state");
+    if (stateClassName == null) {
+      throw new IllegalStateException("Invalid snapshot");
+    }
+    try {
+      Class<?> stateClass = Class.forName(stateClassName.toString());
+      this.state = (T) stateClass.newInstance();
+      initialize();
+    } catch (ClassNotFoundException | InstantiationException | IllegalAccessException e) {
+      throw new IllegalStateException("Invalid snapshot state");
+    }
+    this.data = (Map<String, Object>) snapshot.get("data");
   }
 
   @Override