/**
 * Create a statusupdate that is related to a cluster manager message.
 * @param message
 *          the related cluster manager message
 * @param level
 *          the error level
 * @param classInfo
 *          class info about the class that reports the status update
 * @param additionalInfo
 *          info the additional debug information
 */
public ZNRecord createMessageStatusUpdateRecord(Message message, Level level, Class classInfo,
    String additionalInfo) {
  ZNRecord result = createEmptyStatusUpdateRecord(getStatusUpdateRecordName(message));
  Map<String, String> contentMap = new TreeMap<String, String>();

  contentMap.put("Message state",
      (message.getMsgState() == null ? "NULL" : message.getMsgState().toString()));
  contentMap.put("AdditionalInfo", additionalInfo);
  contentMap.put("Class", classInfo.toString());
  contentMap.put("MSG_ID", message.getMsgId());

  result.setMapField(generateMapFieldId(level, getRecordIdForMessage(message)), contentMap);

  return result;
}
 

@Override
public void onBecomeLeaderFromStandby(Message message, NotificationContext context)
    throws Exception {
  String clusterName = message.getPartitionName();
  String controllerName = message.getTgtName();

  logger.info(controllerName + " becoming leader from standby for " + clusterName);

  if (_controller == null) {
    _controller =
        HelixManagerFactory.getZKHelixManager(clusterName, controllerName,
            InstanceType.CONTROLLER, _zkAddr);
    _controller.setEnabledControlPipelineTypes(_enabledPipelineTypes);
    _controller.connect();
    _controller.startTimerTasks();
    logStateTransition("STANDBY", "LEADER", clusterName, controllerName);
  } else {
    logger.error("controller already exists:" + _controller.getInstanceName() + " for "
        + clusterName);
  }

}
 

private boolean isJobFinished(JobContext jobContext, String jobResource,
    CurrentStateOutput currentStateOutput) {
  for (int pId : jobContext.getPartitionSet()) {
    TaskPartitionState state = jobContext.getPartitionState(pId);
    Partition partition = new Partition(pName(jobResource, pId));
    String instance = jobContext.getAssignedParticipant(pId);
    Message pendingMessage =
        currentStateOutput.getPendingMessage(jobResource, partition, instance);
    // If state is INIT but is pending INIT->RUNNING, it's not yet safe to say the job finished
    if (state == TaskPartitionState.RUNNING
        || (state == TaskPartitionState.INIT && pendingMessage != null)) {
      return false;
    }
  }
  return true;
}
 

public synchronized final void onReply(Message message) {
  _logger.info("OnReply msg " + message.getMsgId());
  if (!isDone()) {
    _messageReplied.add(message);
    try {
      onReplyMessage(message);
    } catch (Exception e) {
      _logger.error(e.toString());
    }
  }
  if (isDone()) {
    if (_timer != null) {
      _timer.cancel();
    }
    notifyAll();
  }
}
 

private void setStateMessage(String resourceName, Partition partition, String instanceName,
    Message message, Map<String, Map<Partition, Map<String, Message>>> stateMessageMap) {
  if (!stateMessageMap.containsKey(resourceName)) {
    stateMessageMap.put(resourceName, new HashMap<Partition, Map<String, Message>>());
  }
  if (!stateMessageMap.get(resourceName).containsKey(partition)) {
    stateMessageMap.get(resourceName).put(partition, new HashMap<String, Message>());
  }
  stateMessageMap.get(resourceName).get(partition).put(instanceName, message);
}
 

@GET
@Path("{clusterId}/controller/messages/{messageId}")
public Response getClusterControllerMessages(@PathParam("clusterId") String clusterId,
    @PathParam("messageId") String messageId) {
  HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
  Message message =
      dataAccessor.getProperty(dataAccessor.keyBuilder().controllerMessage(messageId));
  return JSONRepresentation(message.getRecord());
}
 

@Test()
public void testOnBecomeStandbyFromLeader() {
  Message message = new Message(MessageType.STATE_TRANSITION, "0");
  message.setPartitionName(clusterName);
  message.setTgtName("controller_0");
  stateModel.onBecomeStandbyFromLeader(message, new NotificationContext(null));
}
 

DefaultMessageHandler(Message message, ServerMetrics metrics, NotificationContext context) {
  super(message, context);
  _segmentName = message.getPartitionName();
  _tableNameWithType = message.getResourceName();
  _metrics = metrics;
  _logger = LoggerFactory.getLogger(_tableNameWithType + "-" + this.getClass().getSimpleName());
}
 

private void forwardRelayMessages(HelixDataAccessor accessor, Message message,
    long taskCompletionTime) {
  if (message.hasRelayMessages()) {
    Map<String, Message> relayMessages = message.getRelayMessages();
    Builder keyBuilder = accessor.keyBuilder();

    // Ignore all relay messages if participant's session has changed.
    if (!_manager.getSessionId().equals(message.getTgtSessionId())) {
      logger.info(
          "Session id has been changed, ignore all relay messages attached with " + message
              .getId());
      return;
    }

    for (String instance : relayMessages.keySet()) {
      Message msg = relayMessages.get(instance);
      if (msg.getMsgSubType().equals(MessageType.RELAYED_MESSAGE.name())) {
        msg.setRelayTime(taskCompletionTime);
        if (msg.isExpired()) {
          logger.info(
              "Relay message expired, ignore " + msg.getId() + " to instance " + instance);
          continue;
        }
        PropertyKey msgKey = keyBuilder.message(instance, msg.getId());
        boolean success = accessor.getBaseDataAccessor()
            .create(msgKey.getPath(), msg.getRecord(), AccessOption.PERSISTENT);
        if (!success) {
          logger.warn("Failed to send relay message " + msg.getId() + " to " + instance);
        } else {
          logger.info("Send relay message " + msg.getId() + " to " + instance);
        }
      }
    }
  }
}
 

List<MessageHandler> createMsgHandlers(List<Message> msgs, NotificationContext context) {

    List<MessageHandler> handlers = new ArrayList<MessageHandler>();
    for (Message msg : msgs) {
      MessageHandler handler = _msgHandlerFty.createHandler(msg, context);
      handlers.add(handler);
    }
    return handlers;
  }
 

/**
 * Update message count per instance and per resource
 * @param messages a list of messages
 */
public void increaseMessageReceived(List<Message> messages) {
  Map<String, Long> messageCountPerInstance = new HashMap<>();
  Map<String, Long> messageCountPerResource = new HashMap<>();

  // Aggregate messages
  for (Message message : messages) {
    String instanceName = message.getAttribute(Message.Attributes.TGT_NAME);
    String resourceName = message.getAttribute(Message.Attributes.RESOURCE_NAME);

    if (instanceName != null) {
      if (!messageCountPerInstance.containsKey(instanceName)) {
        messageCountPerInstance.put(instanceName, 0L);
      }
      messageCountPerInstance.put(instanceName, messageCountPerInstance.get(instanceName) + 1L);
    }

    if (resourceName != null) {
      if (!messageCountPerResource.containsKey(resourceName)) {
        messageCountPerResource.put(resourceName, 0L);
      }
      messageCountPerResource.put(resourceName, messageCountPerResource.get(resourceName) + 1L);
    }
  }

  // Update message count per instance and per resource
  for (String instance : messageCountPerInstance.keySet()) {
    InstanceMonitor instanceMonitor = _instanceMonitorMap.get(instance);
    if (instanceMonitor != null) {
      instanceMonitor.increaseMessageCount(messageCountPerInstance.get(instance));
    }
  }
  for (String resource : messageCountPerResource.keySet()) {
    ResourceMonitor resourceMonitor = _resourceMonitorMap.get(resource);
    if (resourceMonitor != null) {
      resourceMonitor.increaseMessageCount(messageCountPerResource.get(resource));
    }
  }
}
 

private List<Message> throttle(Map<String, Integer> throttleMap, ClusterConstraints constraint,
    List<Message> messages, final boolean needThrottle) {

  List<Message> throttleOutputMsgs = new ArrayList<Message>();
  for (Message message : messages) {
    Map<ConstraintAttribute, String> msgAttr = ClusterConstraints.toConstraintAttributes(message);

    Set<ConstraintItem> matches = constraint.match(msgAttr);
    matches = selectConstraints(matches, msgAttr);

    boolean msgThrottled = false;
    for (ConstraintItem item : matches) {
      String key = item.filter(msgAttr).toString();
      if (!throttleMap.containsKey(key)) {
        throttleMap.put(key, valueOf(item.getConstraintValue()));
      }
      int value = throttleMap.get(key);
      throttleMap.put(key, --value);

      if (needThrottle && value < 0) {
        msgThrottled = true;

        if (LOG.isDebugEnabled()) {
          // TODO: printout constraint item that throttles the message
          LogUtil.logDebug(LOG, _eventId,
              "message: " + message + " is throttled by constraint: " + item);
        }
      }
    }
    if (!msgThrottled) {
      throttleOutputMsgs.add(message);
    }
  }

  return throttleOutputMsgs;
}
 

@Transition(
    from = "OFFLINE",
    to = "ONLINE"
)
public void onBecomeOnlineFromOffline(Message message, NotificationContext context) {
    LOGGER.info("Become Online from Offline");
}
 

/**
 * Remove the given message from ZK using the given accessor. This function will
 * not throw exception
 * @param accessor HelixDataAccessor
 * @param msg message to remove
 * @param instanceName name of the instance on which the message sits
 * @return true if success else false
 */
public static boolean removeMessageFromZK(HelixDataAccessor accessor, Message msg,
    String instanceName) {
  try {
    return accessor.removeProperty(msg.getKey(accessor.keyBuilder(), instanceName));
  } catch (Exception e) {
    LOG.error("Caught exception while removing message {}.", msg, e);
  }
  return false;
}
 

@Transition(to = "ONLINE", from = "OFFLINE")
public void onBecomeOnlineFromOffline(Message message, NotificationContext context) {
    logger.info("becomes ONLINE from OFFLINE for {}", partitionId);
    if (handler.compareAndSet(null, handlerFactory.createParitionHandler(partitionId))) {
        handler.get().startServing();
    }
}
 

@Override
public MessageHandler createHandler(Message message, NotificationContext context) {
  String type = message.getMsgType();

  if (!type.equals(getMessageType())) {
    throw new HelixException("Unexpected msg type for message " + message.getMsgId() + " type:"
        + message.getMsgType());
  }

  return new DefaultSchedulerMessageHandler(message, context, _manager);
}
 

@Override
public void doTransition(Message message, NotificationContext context)
    throws InterruptedException {
  String instance = message.getTgtName();
  String partition = message.getPartitionName();
  if (instance.equals("localhost_12918") && partition.equals("TestDB0_0")
      && !_done.getAndSet(true)) {
    _startCountdown.countDown();
    // this await will be interrupted since we cancel the task during handleNewSession
    _endCountdown.await();
  }
}
 

@Override
public HelixTaskResult handleMessage() {
  logger.warn(String.format("No handling setup for %s message of subtype: %s",
      Message.MessageType.USER_DEFINE_MSG.toString(), this._message.getMsgSubType()));

  HelixTaskResult helixTaskResult = new HelixTaskResult();
  helixTaskResult.setSuccess(true);
  return helixTaskResult;
}
 

public ControllerUserDefinedMessageHandler(Message message, NotificationContext context, String serviceName,
    boolean flowCatalogLocalCommit, GobblinServiceJobScheduler scheduler,
    GobblinServiceFlowConfigResourceHandler resourceHandler) {
  super(message, context);
  this.serviceName = serviceName;
  this.flowCatalogLocalCommit = flowCatalogLocalCommit;
  this.jobScheduler = scheduler;
  this.resourceHandler = resourceHandler;
}
 

@Transition(to = "INIT", from = "RUNNING")
public void onBecomeInitFromRunning(Message msg, NotificationContext context) {
  String taskPartition = msg.getPartitionName();
  if (_taskRunner == null) {
    throw new IllegalStateException(String
        .format("Invalid state transition. There is no running task for partition %s.", taskPartition));
  }

  _taskRunner.cancel();
  TaskResult r = _taskRunner.waitTillDone();
  LOG.info(String.format("Task partition %s returned result %s.", msg.getPartitionName(), r));
  _taskRunner = null;
}
 

@Override
public HelixTaskResult handleMessage() throws InterruptedException {
  LOGGER.warn(String
      .format("No handling setup for %s message of subtype: %s", Message.MessageType.USER_DEFINE_MSG.toString(),
          this._message.getMsgSubType()));

  final HelixTaskResult helixTaskResult = new HelixTaskResult();
  helixTaskResult.setSuccess(true);
  return helixTaskResult;
}
 

/**
 * If the node is slave, start the rsync thread if it is not started
 * @param message
 * @param context
 * @throws Exception
 */

@Transition(from = "OFFLINE", to = "SLAVE")
public void onBecomeSlaveFromOffline(Message message, NotificationContext context)
    throws Exception {
  System.out.println(_serverId + " transitioning from " + message.getFromState() + " to "
      + message.getToState() + " for " + _partition);

  replicator.start();
  System.out.println(_serverId + " transitioned from " + message.getFromState() + " to "
      + message.getToState() + " for " + _partition);
}
 

private void cacheRelayMessage(Message relayMessage, Message hostMessage) {
  String instanceName = relayMessage.getTgtName();
  if (!_relayMessageCache.containsKey(instanceName)) {
    _relayMessageCache.put(instanceName, Maps.<String, Message> newHashMap());
  }
  if (!_relayMessageCache.get(instanceName).containsKey(relayMessage.getId())) {
    // Only log if the message doesn't already exist in the cache
    LOG.info("Add relay message to relay cache " + relayMessage.getMsgId() + ", hosted message "
        + hostMessage.getMsgId());
  }
  _relayMessageCache.get(instanceName).put(relayMessage.getId(), relayMessage);
  _relayHostMessageCache.put(relayMessage.getMsgId(), hostMessage);
}
 

@GET
@Path("messages/{messageId}")
public Response getMessageOnInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName,
    @PathParam("messageId") String messageId) throws IOException {
  HelixDataAccessor accessor = getDataAccssor(clusterId);
  Message message = accessor.getProperty(accessor.keyBuilder().message(instanceName, messageId));
  if (message != null) {
    return JSONRepresentation(message.getRecord());
  }

  return notFound();
}
 

@Test
public void testP2PWithStateCancellationMessage() {
  ClusterEvent event = generateClusterEvent();
  runStage(event, new ResourceMessageGenerationPhase());
  MessageOutput messageOutput = event.getAttribute(AttributeName.MESSAGES_ALL.name());
  // No message should be sent for partition 0
  Assert.assertEquals(messageOutput.getMessages(RESOURCE_NAME, new Partition("0")).size(), 0);

  // One cancellation message should be sent out for partition 1
  List<Message> messages = messageOutput.getMessages(RESOURCE_NAME, new Partition("1"));
  Assert.assertEquals(messages.size(), 1);
  Assert.assertEquals(messages.get(0).getMsgType(),
      Message.MessageType.STATE_TRANSITION_CANCELLATION.name());
}
 

public void logError(Message message, Class classInfo, Exception e, String additionalInfo,
    HelixManager manager) {
  StringWriter sw = new StringWriter();
  PrintWriter pw = new PrintWriter(sw);
  e.printStackTrace(pw);
  logMessageStatusUpdateRecord(message, Level.HELIX_ERROR, classInfo,
      additionalInfo + sw.toString(), manager);
}
 

/**
 * Run the task. The parallelism of this is dictated by the constraints that are set.
 * @param message
 * @param context
 * @throws InterruptedException
 */
public void onBecomeOnlineFromOffline(final Message message, NotificationContext context)
    throws InterruptedException {
  // Do the work, and then finally remove the instance from the preference list for this
  // partition
  HelixManager manager = context.getManager();
  LOG.info("START onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
      + manager.getInstanceName());
  int oldSize;
  synchronized (_instanceList) {
    oldSize = _instanceList.size();
    _instanceList.add(manager.getInstanceName());
  }
  Assert.assertEquals(oldSize, 0); // ensure these transitions are fully synchronized

  Thread.sleep(TRANSITION_TIME); // a sleep simulates work

  // Need to disable in order to get the transition the next time
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  removeInstanceFromPreferences(accessor, manager.getInstanceName(), message.getResourceName(),
      message.getPartitionName());
  LOG.info("FINISH onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
      + manager.getInstanceName());

  int newSize;
  synchronized (_instanceList) {
    _instanceList.remove(_instanceList.size() - 1);
    newSize = _instanceList.size();
  }
  Assert.assertEquals(newSize, oldSize); // ensure nothing came in during this time
  _onlineLatch.countDown();
}
 

public void onBecomeOnlineFromOffline(Message message, NotificationContext context) {
  String partitionName = message.getPartitionName();
  LOG.info("Received transition to Online from Offline for partition: " + partitionName);
  TimelineMetricAggregator.AGGREGATOR_TYPE type = PARTITION_AGGREGATION_TYPES.get(partitionName);
  taskRunner.setPartitionAggregationFunction(type);
}
 

@Transition(from = "OFFLINE", to = "ONLINE")
public void onBecomeOnlineFromOffline(Message message, NotificationContext context) {
  LOGGER.info("ControllerStateModel.onBecomeOnlineFromOffline() for resource: "
    + message.getResourceName() + ", partition: " + message.getPartitionName());
  handleStateChange(message);
}
 

@Transition(to = "OFFLINE", from = "ERROR")
public void onBecomeOfflineFromError(Message message, NotificationContext context) {
  String partitionName = message.getPartitionName();
  String instanceName = message.getTgtName();
  LOGGER.info(instanceName + " becomes OFFLINE from ERROR for " + partitionName);
}