@Override
public void onLiveInstanceChange(List<LiveInstance> liveInstances,
    NotificationContext changeContext) {
  if (changeContext.getType() != NotificationContext.Type.FINALIZE) {
    for (LiveInstance liveInstance : liveInstances) {
      if (_expectedLiveInstances.contains(liveInstance.getInstanceName())) {
        try {
          _manager.addCurrentStateChangeListener(
              (CurrentStateChangeListener) (instanceName, statesInfo, currentStateChangeContext) -> {
                // empty callback
              }, liveInstance.getInstanceName(), liveInstance.getEphemeralOwner());
        } catch (Exception e) {
          throw new HelixException("Unexpected exception in the test method.", e);
        }
      }
    }
  }
}
 

private void subscribeDataChange(String path, NotificationContext.Type callbackType) {
  if (callbackType == NotificationContext.Type.INIT
      || callbackType == NotificationContext.Type.CALLBACK) {
    if (logger.isDebugEnabled()) {
      logger.debug("{} subscribe data-change. path: {}, listener: {}", _manager.getInstanceName(),
          path, _listener);
    }
    boolean subStatus = _zkClient.subscribeDataChanges(path, this, callbackType != Type.INIT);
    logger.debug("CallbackHandler {} subscribe data path {} result {}", this, path, subStatus);
    if (!subStatus) {
      logger.info("CallbackHandler {} subscribe data path {} failed!", this, path);
    }
  } else if (callbackType == NotificationContext.Type.FINALIZE) {
    logger.info("{} unsubscribe data-change. path: {}, listener: {}",
        _manager.getInstanceName(), path, _listener);

    _zkClient.unsubscribeDataChanges(path, this);
  }
}
 

@Transition(to = "STOPPED", from = "RUNNING")
public String onBecomeStoppedFromRunning(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 %s completed with result %s.", msg.getPartitionName(), r));

  timeout_task.cancel(false);

  return r.getInfo();
}
 

/**
 * 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();
}
 

@Transition(to = "INACTIVE", from = "STANDBY")
public void onBecomeInactiveFromStandby(Message message, NotificationContext context) {
  String partitionName = message.getPartitionName();
  logger.info("Partition {} in resource {} is becoming INACTIVE from STANDBY", partitionName,
      message.getResourceName());
  if (clusterMapConfig.clustermapEnableStateModelListener) {
    partitionStateChangeListener.onPartitionBecomeInactiveFromStandby(partitionName);
  }
}
 

public HelixBatchMessageTask(Message batchMsg, List<Message> subMsgs,
    List<MessageHandler> handlers, NotificationContext context) {
  _batchMsg = batchMsg;
  _context = context;
  _subMsgs = subMsgs;
  _handlers = handlers;
}
 

private void markReadMessage(Message message, NotificationContext context,
    HelixManager manager) {
  message.setMsgState(MessageState.READ);
  message.setReadTimeStamp(new Date().getTime());
  message.setExecuteSessionId(context.getManager().getSessionId());

  _statusUpdateUtil.logInfo(message, HelixStateMachineEngine.class, "New Message", manager);
}
 

SubscribeChangeEvent(CallbackHandler handler, NotificationContext.Type callbackType,
    String path, boolean watchChild, Object listener) {
  this.handler = handler;
  this.path = path;
  this.callbackType = callbackType;
  this.listener = listener;
  this.watchChild = watchChild;
}
 

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());
}
 

@Transition(from = "ONLINE", to = "DROPPED")
public void onBecomeDroppedFromOnline(Message message, NotificationContext context) {
  _logger.info("SegmentOnlineOfflineStateModel.onBecomeDroppedFromOnline() : " + message);
  try {
    onBecomeOfflineFromOnline(message, context);
    onBecomeDroppedFromOffline(message, context);
  } catch (final Exception e) {
    _logger.error("Caught exception on ONLINE -> DROPPED state transition", e);
    Utils.rethrowException(e);
  }
}
 

@Transition(to = "*", from = "*")
public void generalTransitionHandle(Message message, NotificationContext context)
    throws InterruptedException {
  LOG.info(String
      .format("Resource %s partition %s becomes %s from %s", message.getResourceName(),
          message.getPartitionName(), message.getToState(), message.getFromState()));
  if (_transition != null) {
    _transition.doTransition(message, context);
  }
}
 

public BatchMessageHandler(Message msg, NotificationContext context, MessageHandlerFactory fty,
    BatchMessageWrapper wrapper, TaskExecutor executor) {
  super(msg, context);

  if (fty == null || executor == null) {
    throw new HelixException("MessageHandlerFactory | TaskExecutor can't be null");
  }

  _msgHandlerFty = fty;
  _batchMsgWrapper = wrapper;
  _executor = executor;

  // create sub-messages
  _subMessages = new ArrayList<Message>();
  List<String> partitionKeys = _message.getPartitionNames();
  for (String partitionKey : partitionKeys) {
    // assign a new message id, put batch-msg-id to parent-id field
    Message subMsg = new Message(_message.getRecord(), UUID.randomUUID().toString());
    subMsg.setPartitionName(partitionKey);
    subMsg.setAttribute(Attributes.PARENT_MSG_ID, _message.getId());
    subMsg.setBatchMessageMode(false);

    _subMessages.add(subMsg);
  }

  // create sub-message handlers
  _subMessageHandlers = createMsgHandlers(_subMessages, context);
}
 

@Transition(from = "ONLINE", to = "OFFLINE")
public void onBecomeOfflineFromOnline(Message message, NotificationContext context) {
  _logger.info("SegmentOnlineOfflineStateModel.onBecomeOfflineFromOnline() : " + message);
  String tableNameWithType = message.getResourceName();
  String segmentName = message.getPartitionName();
  try {
    _instanceDataManager.removeSegment(tableNameWithType, segmentName);
  } catch (Exception e) {
    _logger.error("Caught exception in state transition from ONLINE -> OFFLINE for resource: {}, partition: {}",
        tableNameWithType, segmentName, e);
    Utils.rethrowException(e);
  }
}
 

@Override
public MessageHandler createHandler(Message message, NotificationContext context) {
  // TODO Auto-generated method stub
  if (message.getMsgSubType() != null && message.getMsgSubType().equals("EXCEPTION")) {
    throw new HelixException("Test Message handler exception, can ignore");
  }
  _handlersCreated++;
  return new TestMessageHandler(message, context);
}
 

private void setupWatcher() throws Exception {

    LiveInstanceChangeListener listener = new LiveInstanceChangeListener() {
      @Override
      public void onLiveInstanceChange(List<LiveInstance> liveInstances,
          NotificationContext changeContext) {
        if (changeContext.getType() != NotificationContext.Type.FINALIZE) {
          refreshCache();
        }
      }
    };
    admin.addLiveInstanceChangeListener(listener);
  }
 

@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();
  }
}
 

@Transition(from = "OFFLINE", to = "ONLINE")
public void onBecomeOnlineFromOffline(Message message, NotificationContext context) {
  String tableNameWithType = message.getPartitionName();
  LOGGER.info("Processing transition from OFFLINE to ONLINE for table: {}", tableNameWithType);
  try {
    _routingManager.buildRouting(tableNameWithType);
    TableConfig tableConfig = ZKMetadataProvider.getTableConfig(_propertyStore, tableNameWithType);
    _queryQuotaManager.initTableQueryQuota(tableConfig,
        _helixDataAccessor.getProperty(_helixDataAccessor.keyBuilder().externalView(BROKER_RESOURCE_INSTANCE)));
  } catch (Exception e) {
    LOGGER.error("Caught exception while processing transition from OFFLINE to ONLINE for table: {}",
        tableNameWithType, e);
    throw e;
  }
}
 

@Override
public void run() {
  try {
    if (_shouldRefreshCacheOption.orElse(
        _clusterEventType.equals(ClusterEventType.PeriodicalRebalance) || _clusterEventType
            .equals(ClusterEventType.OnDemandRebalance))) {
      requestDataProvidersFullRefresh();

      HelixDataAccessor accessor = _manager.getHelixDataAccessor();
      PropertyKey.Builder keyBuilder = accessor.keyBuilder();
      List<LiveInstance> liveInstances =
          accessor.getChildValues(keyBuilder.liveInstances(), true);

      if (liveInstances != null && !liveInstances.isEmpty()) {
        NotificationContext changeContext = new NotificationContext(_manager);
        changeContext.setType(NotificationContext.Type.CALLBACK);
        synchronized (_manager) {
          checkLiveInstancesObservation(liveInstances, changeContext);
        }
      }
    }
    forceRebalance(_manager, _clusterEventType);
  } catch (Throwable ex) {
    logger.error("Time task failed. Rebalance task type: " + _clusterEventType + ", cluster: "
        + _clusterName, ex);
  }
}
 

@Override
public void doTransition(Message message, NotificationContext context) {
  // System.err.println("doReset() invoked");
  super.doTransition(message, context);
  String fromState = message.getFromState();
  String toState = message.getToState();
  if (fromState.equals("ERROR") && toState.equals("OFFLINE")) {
    _errToOfflineInvoked++;
  }
}
 

@Override
@PreFetch(enabled = true)
public void onLiveInstanceChange(List<LiveInstance> liveInstances,
    NotificationContext changeContext) {
  if (_isBlocking) {
    try {
      newEventHandlingCount.acquire();
    } catch (InterruptedException e) {
      throw new HelixException("Failed to acquire handling lock for testing.");
    }
  }
  super.onLiveInstanceChange(liveInstances, changeContext);
}
 

@Override
public void onExternalViewChange(List<ExternalView> externalViewList,
    NotificationContext changeContext) {
  if (changeContext.getType() == Type.CALLBACK) {
    cbCnt++;
  }
}
 

@Test()
public void testOnBecomeLeaderFromStandby() {
  Message message = new Message(MessageType.STATE_TRANSITION, "0");
  message.setPartitionName(clusterName);
  message.setTgtName("controller_0");
  try {
    stateModel.onBecomeLeaderFromStandby(message, new NotificationContext(null));
  } catch (Exception e) {
    LOG.error("Exception becoming leader from standby", e);
  }
  stateModel.onBecomeStandbyFromLeader(message, new NotificationContext(null));
}
 

@Override
public boolean cancelTask(MessageTask task) {
  Message message = task.getMessage();
  NotificationContext notificationContext = task.getNotificationContext();
  String taskId = task.getTaskId();

  synchronized (_lock) {
    if (_taskMap.containsKey(taskId)) {
      MessageTaskInfo taskInfo = _taskMap.get(taskId);
      // cancel timeout task
      if (taskInfo._timerTask != null) {
        taskInfo._timerTask.cancel();
      }

      // cancel task
      Future<HelixTaskResult> future = taskInfo.getFuture();
      removeMessageFromTaskAndFutureMap(message);
      _statusUpdateUtil.logInfo(message, HelixTaskExecutor.class, "Canceling task: " + taskId,
          notificationContext.getManager());

      // If the thread is still running it will be interrupted if cancel(true)
      // is called. So state transition callbacks should implement logic to
      // return if it is interrupted.
      if (future.cancel(true)) {
        _statusUpdateUtil.logInfo(message, HelixTaskExecutor.class, "Canceled task: " + taskId,
            notificationContext.getManager());
        _taskMap.remove(taskId);
        return true;
      } else {
        _statusUpdateUtil.logInfo(message, HelixTaskExecutor.class,
            "fail to cancel task: " + taskId, notificationContext.getManager());
      }
    } else {
      _statusUpdateUtil.logWarning(message, HelixTaskExecutor.class,
          "fail to cancel task: " + taskId + ", future not found",
          notificationContext.getManager());
    }
  }
  return false;
}
 

@Transition(to = "DROPPED", from = "OFFLINE")
public void onBecomeDroppedFromOffline(Message message, NotificationContext context) {
  String partitionName = message.getPartitionName();
  logger.info("Partition {} in resource {} is becoming DROPPED from OFFLINE", partitionName,
      message.getResourceName());
  if (clusterMapConfig.clustermapEnableStateModelListener) {
    partitionStateChangeListener.onPartitionBecomeDroppedFromOffline(partitionName);
  }
}
 

@Override
@PreFetch(enabled = true)
public void onLiveInstanceChange(List<LiveInstance> liveInstances,
    NotificationContext changeContext) {
  if (_sourceDataTypeMap.containsKey(PropertyType.CURRENTSTATES)) {
    // Go though the live instance list and update CurrentState listeners
    updateCurrentStatesListeners(liveInstances, changeContext);
  }
  _routerUpdater.queueEvent(changeContext, ClusterEventType.LiveInstanceChange,
      HelixConstants.ChangeType.LIVE_INSTANCE);
}
 

@Test()
public void testRollbackOnError() {
  Message message = new Message(MessageType.STATE_TRANSITION, "0");
  message.setPartitionName(clusterName);
  message.setTgtName("controller_0");
  try {
    stateModel.onBecomeLeaderFromStandby(message, new NotificationContext(null));
  } catch (Exception e) {
    LOG.error("Exception becoming leader from standby", e);
  }
  stateModel.rollbackOnError(message, new NotificationContext(null), null);
}
 

@Test
public void testUsingNameConvention() {
  StateModelParser parser = new StateModelParser();
  StateModelUsingNameConvention testModel = new StateModelUsingNameConvention();

  Method method =
      parser.getMethodForTransition(testModel.getClass(), "error", "dropped", new Class[] {
          Message.class, NotificationContext.class
      });
  Assert.assertNotNull(method);
  Assert.assertEquals(method.getName(), "onBecomeDroppedFromError");

}
 

@Override
@PreFetch(enabled = false)
public void onCustomizedStateChange(String instanceName, List<CustomizedState> statesInfo,
    NotificationContext changeContext) {
  logger.info("START: GenericClusterController.onCustomizedStateChange()");
  notifyCaches(changeContext, ChangeType.CUSTOMIZED_STATE);
  pushToEventQueues(ClusterEventType.CustomizedStateChange, changeContext, Collections
      .<String, Object>singletonMap(AttributeName.instanceName.name(), instanceName));
  logger.info("END: GenericClusterController.onCustomizedStateChange()");
}
 

@Transition(to = "MASTER", from = "SLAVE")
public void onBecomeMasterFromSlave(Message message, NotificationContext context)
    throws InterruptedException {
  LOG.info("Become MASTER from SLAVE");
  if (_transition != null && _sleep) {
    _transition.doTransition(message, context);
  }
}
 

@Override
public void onConfigChange(List<InstanceConfig> configs, NotificationContext changeContext) {
  LOG.info("onConfigChange() invoked");
  callParticipantCode(changeContext);
}