@Override
public void doTransition(Message message, NotificationContext context)
throws InterruptedException {
String instance = message.getTgtName();
PartitionId partition = message.getPartitionId();
if (instance.equals("localhost_12918")
&& partition.equals("TestDB0_0")
&& _done.getAndSet(true) == false) {
_startCountdown.countDown();
// this await will be interrupted since we cancel the task during handleNewSession
_endCountdown.await();
}
}
/**
* Set an inner message
*
* @param inner message
*/
public void setInnerMessage(Message message) {
_record.setMapField(Attributes.INNER_MESSAGE.name(), message.getRecord().getSimpleFields());
}
/**
* Create a reply based on an incoming message
*
* @param srcMessage the incoming message
* @param instanceName the instance that is the source of the reply
* @param taskResultMap the result of executing the incoming message
* @return the reply Message
*/
public static Message createReplyMessage(
Message srcMessage, String instanceName, Map<String, String> taskResultMap) {
if (srcMessage.getCorrelationId() == null) {
throw new HelixException(
"Message " + srcMessage.getMessageId() + " does not contain correlation id");
}
Message replyMessage =
new Message(MessageType.TASK_REPLY, MessageId.from(UUID.randomUUID().toString()));
replyMessage.setCorrelationId(srcMessage.getCorrelationId());
replyMessage.setResultMap(taskResultMap);
replyMessage.setTgtSessionId(SessionId.from("*"));
replyMessage.setMsgState(MessageState.NEW);
replyMessage.setSrcName(instanceName);
if (srcMessage.getSrcInstanceType() == InstanceType.CONTROLLER) {
replyMessage.setTgtName("Controller");
} else {
replyMessage.setTgtName(srcMessage.getMsgSrc());
}
return replyMessage;
}
@Override
public int compare(Message m1, Message m2) {
// long t1 = m1.getCreateTimeStamp();
// long t2 = m2.getCreateTimeStamp();
return (int) (m1.getCreateTimeStamp() - m2.getCreateTimeStamp());
}
private Map<String, String> getMessageContent(final Message message) {
return message.getRecord().getMapField("content");
}
private MessageType buildMessageTypeFromReply(Message message) {
final Map<String, String> result =
message.getRecord().getMapField(Message.Attributes.MESSAGE_RESULT.toString());
return buildMessageType(result.get("type"));
}
/**
* greedy message selection algorithm: 1) calculate CS+PS state lower/upper-bounds 2) group
* messages by state transition and sorted by priority 3) from highest priority to lowest, for
* each message group with the same transition add message one by one and make sure state
* constraint is not violated update state lower/upper-bounds when a new message is selected
*
* @param currentStates
* @param pendingStates
* @param messages
* @param stateConstraints : STATE -> bound (lower:upper)
* @param stateTransitionPriorities : FROME_STATE-TO_STATE -> priority
* @return: selected messages
*/
List<Message> selectMessages(
Map<ParticipantId, Participant> liveParticipants,
Map<ParticipantId, State> currentStates,
Map<ParticipantId, State> pendingStates,
List<Message> messages,
Map<State, Bounds> stateConstraints,
final Map<String, Integer> stateTransitionPriorities,
State initialState) {
if (messages == null || messages.isEmpty()) {
return Collections.emptyList();
}
List<Message> selectedMessages = new ArrayList<Message>();
Map<State, Bounds> bounds = new HashMap<State, Bounds>();
// count currentState, if no currentState, count as in initialState
for (ParticipantId liveParticipantId : liveParticipants.keySet()) {
State state = initialState;
if (currentStates.containsKey(liveParticipantId)) {
state = currentStates.get(liveParticipantId);
}
if (!bounds.containsKey(state)) {
bounds.put(state, new Bounds(0, 0));
}
bounds.get(state).increaseLowerBound();
bounds.get(state).increaseUpperBound();
}
// count pendingStates
for (ParticipantId participantId : pendingStates.keySet()) {
State state = pendingStates.get(participantId);
if (!bounds.containsKey(state)) {
bounds.put(state, new Bounds(0, 0));
}
// TODO: add lower bound, need to refactor pendingState to include fromState also
bounds.get(state).increaseUpperBound();
}
// group messages based on state transition priority
Map<Integer, List<Message>> messagesGroupByStateTransitPriority =
new TreeMap<Integer, List<Message>>();
for (Message message : messages) {
State fromState = message.getTypedFromState();
State toState = message.getTypedToState();
String transition = fromState.toString() + "-" + toState.toString();
int priority = Integer.MAX_VALUE;
if (stateTransitionPriorities.containsKey(transition)) {
priority = stateTransitionPriorities.get(transition);
}
if (!messagesGroupByStateTransitPriority.containsKey(priority)) {
messagesGroupByStateTransitPriority.put(priority, new ArrayList<Message>());
}
messagesGroupByStateTransitPriority.get(priority).add(message);
}
// select messages
for (List<Message> messageList : messagesGroupByStateTransitPriority.values()) {
for (Message message : messageList) {
State fromState = message.getTypedFromState();
State toState = message.getTypedToState();
if (!bounds.containsKey(fromState)) {
LOG.error("Message's fromState is not in currentState. message: " + message);
continue;
}
if (!bounds.containsKey(toState)) {
bounds.put(toState, new Bounds(0, 0));
}
// check lower bound of fromState
if (stateConstraints.containsKey(fromState)) {
int newLowerBound = bounds.get(fromState).getLowerBound() - 1;
if (newLowerBound < 0) {
LOG.error(
"Number of currentState in "
+ fromState
+ " is less than number of messages transiting from "
+ fromState);
continue;
}
if (newLowerBound < stateConstraints.get(fromState).getLowerBound()) {
continue;
}
}
// check upper bound of toState
if (stateConstraints.containsKey(toState)) {
int newUpperBound = bounds.get(toState).getUpperBound() + 1;
if (newUpperBound > stateConstraints.get(toState).getUpperBound()) {
continue;
}
}
selectedMessages.add(message);
bounds.get(fromState).increaseLowerBound();
bounds.get(toState).increaseUpperBound();
}
}
return selectedMessages;
}