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 index c11b07f097..c5c59ec0a4 100644 --- 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 createProxy(Class type) { * Takes a snapshot of the state machine state. */ private Map snapshot() { - return data; + Map 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 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) snapshot.get("data"); } @Override