@Override
@SuppressWarnings("unchecked")
public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
if (mergingWindowsByKey != null) {
TupleSerializer<Tuple2<W, W>> tupleSerializer =
new TupleSerializer<>(
(Class) Tuple2.class, new TypeSerializer[] {windowSerializer, windowSerializer});
ListStateDescriptor<Tuple2<W, W>> mergeStateDescriptor =
new ListStateDescriptor<>("merging-window-set", tupleSerializer);
for (Map.Entry<K, MergingWindowSet<W>> key : mergingWindowsByKey.entrySet()) {
setKeyContext(key.getKey());
ListState<Tuple2<W, W>> mergeState =
getStateBackend()
.getPartitionedState(null, VoidSerializer.INSTANCE, mergeStateDescriptor);
mergeState.clear();
key.getValue().persist(mergeState);
}
}
StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
AbstractStateBackend.CheckpointStateOutputView out =
getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp);
snapshotTimers(out);
taskState.setOperatorState(out.closeAndGetHandle());
return taskState;
}
@Override
public void restoreState(StreamTaskState taskState, long recoveryTimestamp) throws Exception {
super.restoreState(taskState, recoveryTimestamp);
final ClassLoader userClassloader = getUserCodeClassloader();
@SuppressWarnings("unchecked")
StateHandle<DataInputView> inputState =
(StateHandle<DataInputView>) taskState.getOperatorState();
DataInputView in = inputState.getState(userClassloader);
int numKeys = in.readInt();
this.windows = new HashMap<>(numKeys);
this.processingTimeTimers = new HashMap<>();
this.watermarkTimers = new HashMap<>();
for (int i = 0; i < numKeys; i++) {
int numWindows = in.readInt();
for (int j = 0; j < numWindows; j++) {
Context context = new Context(in, userClassloader);
Map<W, Context> keyWindows = windows.get(context.key);
if (keyWindows == null) {
keyWindows = new HashMap<>(numWindows);
windows.put(context.key, keyWindows);
}
keyWindows.put(context.window, context);
}
}
}
@Override
public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
// here, we deal with key/value state snapshots
StreamTaskState state = new StreamTaskState();
if (stateBackend != null) {
HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> partitionedSnapshots =
stateBackend.snapshotPartitionedState(checkpointId, timestamp);
if (partitionedSnapshots != null) {
state.setKvStates(partitionedSnapshots);
}
}
return state;
}
@Override
@SuppressWarnings("rawtypes,unchecked")
public void restoreState(StreamTaskState state, long recoveryTimestamp) throws Exception {
// restore the key/value state. the actual restore happens lazily, when the function requests
// the state again, because the restore method needs information provided by the user function
if (stateBackend != null) {
stateBackend.injectKeyValueStateSnapshots((HashMap) state.getKvStates(), recoveryTimestamp);
}
}
@Override
public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
// we write the panes with the key/value maps into the stream
StateBackend.CheckpointStateOutputView out =
getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp);
int numKeys = windows.size();
out.writeInt(numKeys);
for (Map.Entry<K, Map<W, Context>> keyWindows : windows.entrySet()) {
int numWindows = keyWindows.getValue().size();
out.writeInt(numWindows);
for (Context context : keyWindows.getValue().values()) {
context.writeToState(out);
}
}
taskState.setOperatorState(out.closeAndGetHandle());
return taskState;
}
@Override
public void restoreState(StreamTaskState taskState, long recoveryTimestamp) throws Exception {
super.restoreState(taskState, recoveryTimestamp);
final ClassLoader userClassloader = getUserCodeClassloader();
@SuppressWarnings("unchecked")
StateHandle<DataInputView> inputState =
(StateHandle<DataInputView>) taskState.getOperatorState();
DataInputView in = inputState.getState(userClassloader);
restoreTimers(in);
}