/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package backtype.storm.topology;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import backtype.storm.Config;
import backtype.storm.generated.GlobalStreamId;
import backtype.storm.state.KeyValueState;
import backtype.storm.state.State;
import backtype.storm.state.StateFactory;
import backtype.storm.task.OutputCollector;
import backtype.storm.task.TopologyContext;
import backtype.storm.tuple.Tuple;
import backtype.storm.windowing.WindowLifecycleListener;
/**
* Wraps a {@link IStatefulWindowedBolt} and handles the execution. Saves the last expired
* and evaluated states of the window during checkpoint and restores the state during recovery.
*/
public class StatefulWindowedBoltExecutor<T extends State> extends WindowedBoltExecutor implements IStatefulBolt<T> {
private static final Logger LOG = LoggerFactory.getLogger(StatefulWindowedBoltExecutor.class);
private final IStatefulWindowedBolt<T> statefulWindowedBolt;
private transient String msgIdFieldName;
private transient TopologyContext topologyContext;
private transient OutputCollector outputCollector;
// last evaluated and last expired message ids per task stream (source taskid + stream-id)
private transient KeyValueState<TaskStream, WindowState> streamState;
private transient List<Tuple> pendingTuples;
// the states to be recovered
private transient Map<TaskStream, WindowState> recoveryStates;
private transient boolean stateInitialized;
private transient boolean prePrepared;
public StatefulWindowedBoltExecutor(IStatefulWindowedBolt<T> bolt) {
super(bolt);
this.statefulWindowedBolt = bolt;
}
@Override
public void prepare(Map stormConf, TopologyContext context, OutputCollector collector) {
prepare(stormConf, context, collector, getWindowState(stormConf, context));
}
// package access for unit tests
void prepare(Map stormConf, TopologyContext context, OutputCollector collector,
KeyValueState<TaskStream, WindowState> windowState) {
init(stormConf, context, collector, windowState);
super.prepare(stormConf, context, collector);
}
private void init(Map stormConf, TopologyContext context, OutputCollector collector,
KeyValueState<TaskStream, WindowState> windowState) {
if (stormConf.containsKey(Config.TOPOLOGY_BOLTS_MESSAGE_ID_FIELD_NAME)) {
msgIdFieldName = (String) stormConf.get(Config.TOPOLOGY_BOLTS_MESSAGE_ID_FIELD_NAME);
} else {
throw new IllegalArgumentException(Config.TOPOLOGY_BOLTS_MESSAGE_ID_FIELD_NAME + " is not set");
}
topologyContext = context;
outputCollector = collector;
streamState = windowState;
pendingTuples = new ArrayList<>();
recoveryStates = new HashMap<>();
stateInitialized = false;
prePrepared = false;
}
@Override
public void execute(Tuple input) {
if (!isStateInitialized()) {
throw new IllegalStateException("execute invoked before initState with input tuple " + input);
} else if (isRecovering()) {
handleRecovery(input);
} else {
super.execute(input);
}
}
@Override
protected void start() {
if (!isStateInitialized() || isRecovering()) {
LOG.debug("Will invoke start after recovery is complete.");
} else {
super.start();
}
}
private void handleRecovery(Tuple input) {
long msgId = getMsgId(input);
TaskStream taskStream = TaskStream.fromTuple(input);
WindowState state = recoveryStates.get(taskStream);
LOG.debug("handleRecovery, recoveryStates {}", recoveryStates);
if (state != null) {
LOG.debug("Tuple msgid {}, saved state {}", msgId, state);
if (msgId <= state.lastExpired) {
LOG.debug("Ignoring tuple since msg id {} <= lastExpired id {}", msgId, state.lastExpired);
outputCollector.ack(input);
} else if (msgId <= state.lastEvaluated) {
super.execute(input);
} else {
LOG.debug("Tuple msg id {} > lastEvaluated id {}, adding to pendingTuples and clearing recovery state " +
"for taskStream {}", msgId, state.lastEvaluated, taskStream);
pendingTuples.add(input);
clearRecoveryState(taskStream);
}
} else {
pendingTuples.add(input);
}
}
private void clearRecoveryState(TaskStream stream) {
recoveryStates.remove(stream);
if (!isRecovering()) {
super.start();
LOG.debug("Recovery complete, processing {} pending tuples", pendingTuples.size());
for (Tuple tuple : pendingTuples) {
super.execute(tuple);
}
}
}
private boolean isRecovering() {
return !recoveryStates.isEmpty();
}
private boolean isStateInitialized() {
return stateInitialized;
}
@Override
public void initState(T state) {
if (stateInitialized) {
LOG.warn("State is already initialized. Ignoring initState");
return;
}
statefulWindowedBolt.initState((T) state);
// query the streamState for each input task stream and compute recoveryStates
for (GlobalStreamId streamId : topologyContext.getThisSources().keySet()) {
for (int taskId : topologyContext.getComponentTasks(streamId.get_componentId())) {
WindowState windowState = streamState.get(new TaskStream(taskId, streamId));
if (windowState != null) {
recoveryStates.put(new TaskStream(taskId, streamId), windowState);
}
}
}
LOG.debug("recoveryStates {}", recoveryStates);
stateInitialized = true;
start();
}
@Override
public void preCommit(long txid) {
if (!isStateInitialized() || (!isRecovering() && prePrepared)) {
LOG.debug("Commit streamState, txid {}", txid);
streamState.commit(txid);
} else {
LOG.debug("Still recovering, ignoring preCommit and not committing streamState.");
}
}
@Override
public void prePrepare(long txid) {
if (!isStateInitialized()) {
LOG.warn("Cannot prepare before initState");
} else if (!isRecovering()) {
LOG.debug("Prepare streamState, txid {}", txid);
streamState.prepareCommit(txid);
prePrepared = true;
} else {
LOG.debug("Still recovering, ignoring prePrepare and not preparing streamState.");
}
}
@Override
public void preRollback() {
LOG.debug("Rollback streamState, stateInitialized {}", stateInitialized);
streamState.rollback();
}
@Override
protected WindowLifecycleListener<Tuple> newWindowLifecycleListener() {
final WindowLifecycleListener<Tuple> parentListener = super.newWindowLifecycleListener();
return new WindowLifecycleListener<Tuple>() {
@Override
public void onExpiry(List<Tuple> events) {
parentListener.onExpiry(events);
}
@Override
public void onActivation(List<Tuple> events, List<Tuple> newEvents, List<Tuple> expired) {
if (isRecovering()) {
String msg = String.format("Unexpected activation with events %s, newEvents %s, expired %s in recovering state. " +
"recoveryStates %s ", events, newEvents, expired, recoveryStates);
LOG.error(msg);
throw new IllegalStateException(msg);
} else {
parentListener.onActivation(events, newEvents, expired);
updateWindowState(expired, newEvents);
}
}
};
}
private void updateWindowState(List<Tuple> expired, List<Tuple> newEvents) {
LOG.debug("Update window state, {} expired, {} new events", expired.size(), newEvents.size());
Map<TaskStream, WindowState> state = new HashMap<>();
updateState(state, expired, false);
updateState(state, newEvents, true);
updateStreamState(state);
}
private void updateState(Map<TaskStream, WindowState> state, List<Tuple> tuples, boolean newEvents) {
for (Tuple tuple : tuples) {
TaskStream taskStream = TaskStream.fromTuple(tuple);
WindowState curState = state.get(taskStream);
WindowState newState;
if ((newState = getUpdatedState(curState, getMsgId(tuple), newEvents)) != null) {
state.put(taskStream, newState);
}
}
}
// update streamState based on stateUpdates
private void updateStreamState(Map<TaskStream, WindowState> state) {
for (Map.Entry<TaskStream, WindowState> entry : state.entrySet()) {
TaskStream taskStream = entry.getKey();
WindowState newState = entry.getValue();
WindowState curState = streamState.get(taskStream);
if (curState == null) {
streamState.put(taskStream, newState);
} else {
WindowState updatedState = new WindowState(Math.max(newState.lastExpired, curState.lastExpired),
Math.max(newState.lastEvaluated, curState.lastEvaluated));
LOG.debug("Update window state, taskStream {}, curState {}, newState {}", taskStream, curState, updatedState);
streamState.put(taskStream, updatedState);
}
}
}
private WindowState getUpdatedState(WindowState state, long msgId, boolean newEvents) {
WindowState result = null;
if (newEvents) {
if (state == null) {
result = new WindowState(Long.MIN_VALUE, msgId);
} else if (msgId > state.lastEvaluated) {
result = new WindowState(state.lastExpired, msgId);
}
} else {
if (state == null) {
result = new WindowState(msgId, Long.MIN_VALUE);
} else if (msgId > state.lastExpired) {
result = new WindowState(msgId, state.lastEvaluated);
}
}
return result;
}
private long getMsgId(Tuple input) {
return input.getLongByField(msgIdFieldName);
}
private KeyValueState<TaskStream, WindowState> getWindowState(Map stormConf, TopologyContext context) {
String namespace = context.getThisComponentId() + "-" + context.getThisTaskId() + "-window";
return (KeyValueState<TaskStream, WindowState>) StateFactory.getState(namespace, stormConf, context);
}
static class WindowState {
private long lastExpired;
private long lastEvaluated;
// for kryo
WindowState() {
}
WindowState(long lastExpired, long lastEvaluated) {
this.lastExpired = lastExpired;
this.lastEvaluated = lastEvaluated;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
WindowState that = (WindowState) o;
if (lastExpired != that.lastExpired) return false;
return lastEvaluated == that.lastEvaluated;
}
@Override
public int hashCode() {
int result = (int) (lastExpired ^ (lastExpired >>> 32));
result = 31 * result + (int) (lastEvaluated ^ (lastEvaluated >>> 32));
return result;
}
@Override
public String toString() {
return "WindowState{" +
"lastExpired=" + lastExpired +
", lastEvaluated=" + lastEvaluated +
'}';
}
}
static class TaskStream {
private int sourceTask;
private GlobalStreamId streamId;
// for kryo
TaskStream() {
}
TaskStream(int sourceTask, GlobalStreamId streamId) {
this.sourceTask = sourceTask;
this.streamId = streamId;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
TaskStream that = (TaskStream) o;
if (sourceTask != that.sourceTask) return false;
return streamId != null ? streamId.equals(that.streamId) : that.streamId == null;
}
@Override
public int hashCode() {
int result = streamId != null ? streamId.hashCode() : 0;
result = 31 * result + sourceTask;
return result;
}
@Override
public String toString() {
return "TaskStream{" +
"sourceTask=" + sourceTask +
", streamId=" + streamId +
'}';
}
static TaskStream fromTuple(Tuple input) {
return new TaskStream(input.getSourceTask(), input.getSourceGlobalStreamid());
}
}
}