org.apache.helix.participant.statemachine.StateModelFactory Java Examples

@Override
public boolean registerStateModelFactory(String stateModelName,
    StateModelFactory<? extends StateModel> factory, String factoryName) {
  if (stateModelName == null || factory == null || factoryName == null) {
    throw new HelixException("stateModelDef|stateModelFactory|factoryName cannot be null");
  }

  logger.info("Register state model factory for state model " + stateModelName
      + " using factory name " + factoryName + " with " + factory);

  if (!_stateModelFactoryMap.containsKey(stateModelName)) {
    _stateModelFactoryMap.put(stateModelName,
        new ConcurrentHashMap<String, StateModelFactory<? extends StateModel>>());
  }

  if (_stateModelFactoryMap.get(stateModelName).containsKey(factoryName)) {
    logger.warn("stateModelFactory for " + stateModelName + " using factoryName " + factoryName
        + " has already been registered.");
    return false;
  }

  _stateModelFactoryMap.get(stateModelName).put(factoryName, factory);
  sendNopMessage();
  return true;
}
 

public HelixStateMachineEngine(HelixManager manager) {
  _stateModelParser = new StateModelParser();
  _manager = manager;

  _stateModelFactoryMap =
      new ConcurrentHashMap<String, Map<String, StateModelFactory<? extends StateModel>>>();
  _stateModelDefs = new ConcurrentHashMap<String, StateModelDefinition>();
}
 

public void start(String stateModelName, StateModelFactory<T> stateModelFactory) {
    try {
        manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
                InstanceType.PARTICIPANT, zkAddr);

        StateMachineEngine stateMachine = manager.getStateMachineEngine();
        stateMachine.registerStateModelFactory(stateModelName, stateModelFactory);

        manager.connect();
    } catch (Exception e) {
        logger.error("failed to connect manager", e);
        throw new RuntimeException("failed to start HelixPartitionManager");
    }
}
 

@Override
public StateModelFactory<? extends StateModel> getStateModelFactory(String stateModelName,
    String factoryName) {
  if (!_stateModelFactoryMap.containsKey(stateModelName)) {
    return null;
  }
  return _stateModelFactoryMap.get(stateModelName).get(factoryName);
}
 

@Override
public void reset() {
  logger.info("Resetting HelixStateMachineEngine");
  for (Map<String, StateModelFactory<? extends StateModel>> ftyMap : _stateModelFactoryMap
      .values()) {
    for (StateModelFactory<? extends StateModel> stateModelFactory : ftyMap.values()) {
      for (String resourceName : stateModelFactory.getResourceSet()) {
        for (String partitionKey : stateModelFactory.getPartitionSet(resourceName)) {
          logger.info("Resetting {}::{}", resourceName, partitionKey);
          StateModel stateModel = stateModelFactory.getStateModel(resourceName, partitionKey);
          if (stateModel != null) {
            stateModel.reset();
            String initialState = _stateModelParser.getInitialState(stateModel.getClass());
            stateModel.updateState(initialState);
            // TODO probably should update the state on ZK. Shi confirm what needs
            // to be done here.
          } else {
            // TODO: If stateModel is null, we might need to do something here
            // This reset() is not synchronized. We observed that during a shutdown (where
            // resources
            // are all dropped), an NPE could be possible due to stateModel being null
            // Two cases are possible: 1) removing a partition/resource 2) adding a
            // partition/resource
            // We may need to add more processing here to make sure things are being set to
            // initialState. Otherwise, there might be inconsistencies that might cause partitions
            // to be stuck in some state (because reset() would be a NOP here)
            logger.warn(
                "Failed to reset due to StateModel being null! Resource: {}, Partition: {}",
                resourceName, partitionKey);
          }
        }
      }
    }
  }
  logger.info("Successfully reset HelixStateMachineEngine");
}
 

/**
 * check state-model factory contains state-models same as in expect-state-model map
 * @param fty
 * @param expectStateModelMap
 */
static void checkStateModelMap(StateModelFactory<? extends StateModel> fty,
    Map<String, String> expectStateModelMap) {
  Assert.assertEquals(fty.getPartitionSet("TestDB0").size(), expectStateModelMap.size());
  for (String partition : fty.getPartitionSet("TestDB0")) {
    StateModel stateModel = fty.getStateModel("TestDB0", partition);
    String actualState = stateModel.getCurrentState();
    String expectState = expectStateModelMap.get(partition);
    LOG.debug(partition + " actual state: " + actualState + ", expect state: " + expectState);
    Assert.assertEquals(actualState, expectState, "partition: " + partition
        + " should be in state: " + expectState + " but was " + actualState);
  }
}
 

public HelixStateTransitionHandler(StateModelFactory<? extends StateModel> stateModelFactory,
    StateModel stateModel, Message message, NotificationContext context,
    CurrentState currentStateDelta) {
  super(message, context);
  _stateModel = stateModel;
  _statusUpdateUtil = new StatusUpdateUtil();
  _transitionMethodFinder = new StateModelParser();
  _currentStateDelta = currentStateDelta;
  _manager = _notificationContext.getManager();
  _stateModelFactory = stateModelFactory;
}
 

StateModelFactory getStateModelFactory() {
  return stateModelFactory;
}
 

@Override
public void start()
    throws Exception {
  LOGGER.info("Starting Pinot server");
  long startTimeMs = System.currentTimeMillis();

  LOGGER.info("Initializing Helix manager with zkAddress: {}, clusterName: {}, instanceId: {}", _zkAddress,
      _helixClusterName, _instanceId);
  setupHelixSystemProperties();
  _helixManager =
      HelixManagerFactory.getZKHelixManager(_helixClusterName, _instanceId, InstanceType.PARTICIPANT, _zkAddress);

  LOGGER.info("Initializing server instance and registering state model factory");
  Utils.logVersions();
  ServerSegmentCompletionProtocolHandler
      .init(_serverConf.subset(SegmentCompletionProtocol.PREFIX_OF_CONFIG_OF_SEGMENT_UPLOADER));
  ServerConf serverInstanceConfig = DefaultHelixStarterServerConfig.getDefaultHelixServerConfig(_serverConf);
  _serverInstance = new ServerInstance(serverInstanceConfig, _helixManager);
  InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
  SegmentFetcherAndLoader fetcherAndLoader = new SegmentFetcherAndLoader(_serverConf, instanceDataManager);
  StateModelFactory<?> stateModelFactory =
      new SegmentOnlineOfflineStateModelFactory(_instanceId, instanceDataManager, fetcherAndLoader);
  _helixManager.getStateMachineEngine()
      .registerStateModelFactory(SegmentOnlineOfflineStateModelFactory.getStateModelName(), stateModelFactory);
  // Start the server instance as a pre-connect callback so that it starts after connecting to the ZK in order to
  // access the property store, but before receiving state transitions
  _helixManager.addPreConnectCallback(_serverInstance::start);

  LOGGER.info("Connecting Helix manager");
  _helixManager.connect();
  _helixAdmin = _helixManager.getClusterManagmentTool();
  updateInstanceConfigIfNeeded(_host, _port);

  // Start restlet server for admin API endpoint
  String accessControlFactoryClass =
      _serverConf.getString(ACCESS_CONTROL_FACTORY_CLASS, DEFAULT_ACCESS_CONTROL_FACTORY_CLASS);
  LOGGER.info("Using class: {} as the AccessControlFactory", accessControlFactoryClass);
  final AccessControlFactory accessControlFactory;
  try {
    accessControlFactory = PluginManager.get().createInstance(accessControlFactoryClass);
  } catch (Exception e) {
    throw new RuntimeException(
        "Caught exception while creating new AccessControlFactory instance using class '" + accessControlFactoryClass
            + "'", e);
  }

  int adminApiPort = _serverConf.getInt(CONFIG_OF_ADMIN_API_PORT, DEFAULT_ADMIN_API_PORT);
  _adminApiApplication = new AdminApiApplication(_serverInstance, accessControlFactory);
  _adminApiApplication.start(adminApiPort);
  setAdminApiPort(adminApiPort);

  ServerMetrics serverMetrics = _serverInstance.getServerMetrics();
  // Register message handler factory
  SegmentMessageHandlerFactory messageHandlerFactory =
      new SegmentMessageHandlerFactory(fetcherAndLoader, instanceDataManager, serverMetrics);
  _helixManager.getMessagingService()
      .registerMessageHandlerFactory(Message.MessageType.USER_DEFINE_MSG.toString(), messageHandlerFactory);

  serverMetrics.addCallbackGauge(INSTANCE_CONNECTED_METRIC_NAME, () -> _helixManager.isConnected() ? 1L : 0L);
  _helixManager
      .addPreConnectCallback(() -> serverMetrics.addMeteredGlobalValue(ServerMeter.HELIX_ZOOKEEPER_RECONNECTS, 1L));

  // Register the service status handler
  registerServiceStatusHandler();

  ControllerLeaderLocator.create(_helixManager);

  if (_serverConf
      .getBoolean(CONFIG_OF_STARTUP_ENABLE_SERVICE_STATUS_CHECK, DEFAULT_STARTUP_ENABLE_SERVICE_STATUS_CHECK)) {
    long endTimeMs = startTimeMs + _serverConf.getLong(CONFIG_OF_STARTUP_TIMEOUT_MS, DEFAULT_STARTUP_TIMEOUT_MS);
    startupServiceStatusCheck(endTimeMs);
  }
  setShuttingDownStatus(false);
  LOGGER.info("Pinot server ready");

  // Create metrics for mmap stuff
  serverMetrics.addCallbackGauge("memory.directBufferCount", PinotDataBuffer::getDirectBufferCount);
  serverMetrics.addCallbackGauge("memory.directBufferUsage", PinotDataBuffer::getDirectBufferUsage);
  serverMetrics.addCallbackGauge("memory.mmapBufferCount", PinotDataBuffer::getMmapBufferCount);
  serverMetrics.addCallbackGauge("memory.mmapBufferUsage", PinotDataBuffer::getMmapBufferUsage);
  serverMetrics.addCallbackGauge("memory.allocationFailureCount", PinotDataBuffer::getAllocationFailureCount);

  _realtimeLuceneIndexRefreshState = RealtimeLuceneIndexRefreshState.getInstance();
  _realtimeLuceneIndexRefreshState.start();
}
 

@Test()
public void testInvocationAnnotated() throws Exception {
  System.out.println("START TestCMTaskHandler.testInvocationAnnotated() at "
      + new Date(System.currentTimeMillis()));
  HelixTaskExecutor executor = new HelixTaskExecutor();
  Message message = new Message(MessageType.STATE_TRANSITION, "Some unique id");
  message.setSrcName("cm-instance-0");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setMsgId("Some unique message id");
  message.setResourceName("TestDB");
  message.setTgtName("localhost");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  MockStateModelAnnotated stateModel = new MockStateModelAnnotated();
  NotificationContext context;

  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();

  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  context = new NotificationContext(manager);

  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  StateModelFactory<MockStateModelAnnotated> stateModelFactory =
      new StateModelFactory<MockStateModelAnnotated>() {

        @Override
        public MockStateModelAnnotated createNewStateModel(String resource, String partitionName) {
          // TODO Auto-generated method stub
          return new MockStateModelAnnotated();
        }

      };

  HelixStateTransitionHandler stHandler =
      new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);

  HelixTask handler = new HelixTask(message, context, stHandler, executor);
  handler.call();
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskHandler.testInvocationAnnotated() at "
      + new Date(System.currentTimeMillis()));
}
 

@Test()
public void testCMTaskExecutor() throws Exception {
  System.out.println("START TestCMTaskExecutor");
  String msgId = "TestMessageId";
  Message message = new Message(MessageType.TASK_REPLY, msgId);

  message.setMsgId(msgId);
  message.setSrcName("cm-instance-0");
  message.setTgtName("cm-instance-1");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setResourceName("TestDB");
  message.setStateModelDef("MasterSlave");

  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  MockHelixTaskExecutor executor = new MockHelixTaskExecutor();
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  executor.registerMessageHandlerFactory(MessageType.TASK_REPLY.name(),
      new AsyncCallbackService());

  NotificationContext context = new NotificationContext(manager);
  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  StateModelFactory<MockMasterSlaveStateModel> stateModelFactory = new StateModelFactory<MockMasterSlaveStateModel>() {

    @Override
    public MockMasterSlaveStateModel createNewStateModel(String resource, String partitionName) {
      // TODO Auto-generated method stub
      return new MockMasterSlaveStateModel();
    }

  };
  HelixStateTransitionHandler handler =
      new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);

  HelixTask task = new HelixTask(message, context, handler, executor);
  executor.scheduleTask(task);
  for (int i = 0; i < 10; i++) {
    if (!executor.isDone(task.getTaskId())) {
      Thread.sleep(500);
    }
  }
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskExecutor");
}
 

/**
 * test drop resource in error state should remove all state-models
 * @throws Exception
 */
@Test
public void testDropErrorPartition() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  int n = 2;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      4, // partitions per resource
      n, // number of nodes
      2, // replicas
      "MasterSlave", true); // do rebalance

  // start controller
  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller");
  controller.syncStart();

  MockParticipantManager[] participants = new MockParticipantManager[n];
  for (int i = 0; i < n; i++) {
    final String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    if (i == 0) {
      Map<String, Set<String>> errTransitionMap = new HashMap<String, Set<String>>();
      Set<String> partitions = new HashSet<String>();
      partitions.add("TestDB0_0");
      errTransitionMap.put("OFFLINE-SLAVE", partitions);
      participants[0].setTransition(new ErrTransition(errTransitionMap));
    }

    participants[i].syncStart();
  }

  Map<String, Map<String, String>> errStates = new HashMap<String, Map<String, String>>();
  errStates.put("TestDB0", new HashMap<String, String>());
  errStates.get("TestDB0").put("TestDB0_0", "localhost_12918");
  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName, errStates));
  Assert.assertTrue(result);

  // check state-models in state-machine
  HelixStateMachineEngine stateMachine =
      (HelixStateMachineEngine) participants[0].getStateMachineEngine();
  StateModelFactory<? extends StateModel> fty = stateMachine.getStateModelFactory("MasterSlave");
  Map<String, String> expectStateModelMap = new TreeMap<String, String>();
  expectStateModelMap.put("TestDB0_0", "ERROR");
  expectStateModelMap.put("TestDB0_1", "MASTER");
  expectStateModelMap.put("TestDB0_2", "SLAVE");
  expectStateModelMap.put("TestDB0_3", "MASTER");
  checkStateModelMap(fty, expectStateModelMap);

  // drop resource
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.dropResource(clusterName, "TestDB0");

  result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);

  // check state models have been dropped also
  Assert.assertTrue(fty.getPartitionSet("TestDB0").isEmpty(),
      "All state-models should be dropped, but was " + fty.getPartitionSet("TestDB0"));

  // cleanup
  controller.syncStop();
  for (int i = 0; i < n; i++) {
    participants[i].syncStop();
  }
  TestHelper.dropCluster(clusterName, _gZkClient);

  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * test drop resource should remove all state models
 * @throws Exception
 */
@Test
public void testDrop() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  int n = 2;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      4, // partitions per resource
      n, // number of nodes
      2, // replicas
      "MasterSlave", true); // do rebalance

  // start controller
  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller");
  controller.syncStart();

  MockParticipantManager[] participants = new MockParticipantManager[n];
  for (int i = 0; i < n; i++) {
    final String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);

  // check state-models in state-machine
  HelixStateMachineEngine stateMachine =
      (HelixStateMachineEngine) participants[0].getStateMachineEngine();
  StateModelFactory<? extends StateModel> fty = stateMachine.getStateModelFactory("MasterSlave");
  Map<String, String> expectStateModelMap = new TreeMap<String, String>();
  expectStateModelMap.put("TestDB0_0", "SLAVE");
  expectStateModelMap.put("TestDB0_1", "MASTER");
  expectStateModelMap.put("TestDB0_2", "SLAVE");
  expectStateModelMap.put("TestDB0_3", "MASTER");
  checkStateModelMap(fty, expectStateModelMap);

  // drop resource
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.dropResource(clusterName, "TestDB0");

  result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);

  // check state models have been dropped also
  Assert.assertTrue(fty.getPartitionSet("TestDB0").isEmpty(),
      "All state-models should be dropped, but was " + fty.getPartitionSet("TestDB0"));

  // cleanup
  controller.syncStop();
  for (int i = 0; i < n; i++) {
    participants[i].syncStop();
  }
  TestHelper.dropCluster(clusterName, _gZkClient);

  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/** Dedicated Thread pool can be provided in configuration or by client.
 *  This method is to check it and update the thread pool if necessary.
 */
private void updateStateTransitionMessageThreadPool(Message message, HelixManager manager) {
  if (!message.getMsgType().equals(MessageType.STATE_TRANSITION.name())) {
    return;
  }

  String resourceName = message.getResourceName();
  String factoryName = message.getStateModelFactoryName();
  String stateModelName = message.getStateModelDef();

  if (factoryName == null) {
    factoryName = HelixConstants.DEFAULT_STATE_MODEL_FACTORY;
  }
  StateModelFactory<? extends StateModel> stateModelFactory =
      manager.getStateMachineEngine().getStateModelFactory(stateModelName, factoryName);

  String perStateTransitionTypeKey =
      getStateTransitionType(getPerResourceStateTransitionPoolName(resourceName),
          message.getFromState(), message.getToState());
  if (perStateTransitionTypeKey != null && stateModelFactory != null
      && !_transitionTypeThreadpoolChecked.contains(perStateTransitionTypeKey)) {
    ExecutorService perStateTransitionTypeExecutor = stateModelFactory
        .getExecutorService(resourceName, message.getFromState(), message.getToState());
    _transitionTypeThreadpoolChecked.add(perStateTransitionTypeKey);

    if (perStateTransitionTypeExecutor != null) {
      _executorMap.put(perStateTransitionTypeKey, perStateTransitionTypeExecutor);
      LOG.info(String
          .format("Added client specified dedicate threadpool for resource %s from %s to %s",
              getPerResourceStateTransitionPoolName(resourceName), message.getFromState(),
              message.getToState()));
      return;
    }
  }

  if (!_resourcesThreadpoolChecked.contains(resourceName)) {
    int threadpoolSize = -1;
    ConfigAccessor configAccessor = manager.getConfigAccessor();
    // Changes to this configuration on thread pool size will only take effect after the participant get restarted.
    if (configAccessor != null) {
      HelixConfigScope scope =
          new HelixConfigScopeBuilder(ConfigScopeProperty.RESOURCE)
              .forCluster(manager.getClusterName()).forResource(resourceName).build();

      String threadpoolSizeStr = configAccessor.get(scope, MAX_THREADS);
      try {
        if (threadpoolSizeStr != null) {
          threadpoolSize = Integer.parseInt(threadpoolSizeStr);
        }
      } catch (Exception e) {
        LOG.error(
            "Failed to parse ThreadPoolSize from resourceConfig for resource" + resourceName, e);
      }
    }
    final String key = getPerResourceStateTransitionPoolName(resourceName);
    if (threadpoolSize > 0) {
      _executorMap.put(key, Executors.newFixedThreadPool(threadpoolSize, new ThreadFactory() {
        @Override public Thread newThread(Runnable r) {
          return new Thread(r, "GerenricHelixController-message_handle_" + key);
        }
      }));
      LOG.info("Added dedicate threadpool for resource: " + resourceName + " with size: "
          + threadpoolSize);
    } else {
      // if threadpool is not configured
      // check whether client specifies customized threadpool.
      if (stateModelFactory != null) {
        ExecutorService executor = stateModelFactory.getExecutorService(resourceName);
        if (executor != null) {
          _executorMap.put(key, executor);
          LOG.info("Added client specified dedicate threadpool for resource: " + key);
        }
      } else {
        LOG.error(String.format(
            "Fail to get dedicate threadpool defined in stateModelFactory %s: using factoryName: %s for resource %s. No stateModelFactory was found!",
            stateModelName, factoryName, resourceName));
      }
    }
    _resourcesThreadpoolChecked.add(resourceName);
  }
}
 

@Override
public boolean removeStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory, String factoryName) {
  throw new UnsupportedOperationException("Remove not yet supported");
}
 

@Override
public boolean removeStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory) {
  throw new UnsupportedOperationException("Remove not yet supported");
}
 

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

  if (!type.equals(MessageType.STATE_TRANSITION.name()) && !type
      .equals(MessageType.STATE_TRANSITION_CANCELLATION.name())) {
    throw new HelixException("Expect state-transition message type, but was "
        + message.getMsgType() + ", msgId: " + message.getMsgId());
  }

  String partitionKey = message.getPartitionName();
  String stateModelName = message.getStateModelDef();
  String resourceName = message.getResourceName();
  String sessionId = message.getTgtSessionId();
  int bucketSize = message.getBucketSize();

  if (stateModelName == null) {
    logger
        .error("Fail to create msg-handler because message does not contain stateModelDef. msgId: "
            + message.getId());
    return null;
  }

  String factoryName = message.getStateModelFactoryName();
  if (factoryName == null) {
    factoryName = HelixConstants.DEFAULT_STATE_MODEL_FACTORY;
  }

  StateModelFactory<? extends StateModel> stateModelFactory =
      getStateModelFactory(stateModelName, factoryName);
  if (stateModelFactory == null) {
    logger.warn("Fail to create msg-handler because cannot find stateModelFactory for model: "
        + stateModelName + " using factoryName: " + factoryName + " for resource: "
        + resourceName);
    return null;
  }

  // check if the state model definition exists and cache it
  if (!_stateModelDefs.containsKey(stateModelName)) {
    HelixDataAccessor accessor = _manager.getHelixDataAccessor();
    Builder keyBuilder = accessor.keyBuilder();
    StateModelDefinition stateModelDef =
        accessor.getProperty(keyBuilder.stateModelDef(stateModelName));
    if (stateModelDef == null) {
      throw new HelixException("fail to create msg-handler because stateModelDef for "
          + stateModelName + " does NOT exist");
    }
    _stateModelDefs.put(stateModelName, stateModelDef);
  }

  if (!message.getBatchMessageMode()) {
    String initState = _stateModelDefs.get(message.getStateModelDef()).getInitialState();
    StateModel stateModel = stateModelFactory.getStateModel(resourceName, partitionKey);
    if (stateModel == null) {
      stateModel = stateModelFactory.createAndAddStateModel(resourceName, partitionKey);
      if (stateModelName.equals(TaskConstants.STATE_MODEL_NAME)
          && message.getToState().equals(TaskPartitionState.DROPPED.name())) {
        // If stateModel is null, that means there was a reboot of the Participant. Then the
        // purpose of this first message must be to drop the task. We manually set the current
        // state to be the same state of fromState (which Controller inferred from JobContext) to
        // allow the Participant to successfully process this dropping transition
        stateModel.updateState(message.getFromState());
      } else {
        stateModel.updateState(initState);
      }
    }
    if (message.getMsgType().equals(MessageType.STATE_TRANSITION_CANCELLATION.name())) {
      return new HelixStateTransitionCancellationHandler(stateModel, message, context);
    } else {
      // create currentStateDelta for this partition
      // TODO: move currentStateDelta to StateTransitionMsgHandler
      CurrentState currentStateDelta = new CurrentState(resourceName);
      currentStateDelta.setSessionId(sessionId);
      currentStateDelta.setStateModelDefRef(stateModelName);
      currentStateDelta.setStateModelFactoryName(factoryName);
      currentStateDelta.setBucketSize(bucketSize);

      currentStateDelta.setState(partitionKey,
          (stateModel.getCurrentState() == null) ? initState : stateModel.getCurrentState());

      return new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);
    }
  } else {
    BatchMessageWrapper wrapper = stateModelFactory.getBatchMessageWrapper(resourceName);
    if (wrapper == null) {
      wrapper = stateModelFactory.createAndAddBatchMessageWrapper(resourceName);
    }

    // get executor-service for the message
    TaskExecutor executor = (TaskExecutor) context.get(MapKey.TASK_EXECUTOR.toString());
    if (executor == null) {
      logger.error(
          "fail to get executor-service for batch message: " + message.getId() + ". msgType: "
              + message.getMsgType() + ", resource: " + message.getResourceName());
      return null;
    }
    return new BatchMessageHandler(message, context, this, wrapper, executor);
  }
}
 

@Override
public boolean registerStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory) {
  return registerStateModelFactory(stateModelDef, factory,
      HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
}
 

@Override
public StateModelFactory<? extends StateModel> getStateModelFactory(String stateModelName) {
  return getStateModelFactory(stateModelName, HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
}
 

/**
 * Get a default state model factory for a state model definition.
 * @param stateModelName
 * @return
 */
public StateModelFactory<? extends StateModel> getStateModelFactory(String stateModelName);
 

/**
 * Get a default state model factory for a state model definition.
 * @param stateModelName
 * @param factoryName
 * @return
 */
public StateModelFactory<? extends StateModel> getStateModelFactory(String stateModelName,
    String factoryName);
 

/**
 * Remove the state model factory with a name for a state model definition
 * @param stateModelDef
 * @param factory
 * @param factoryName
 * @return
 */
public boolean removeStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory, String factoryName);
 

/**
 * Remove the default state model factory for a state model definition
 * @param stateModelDef
 * @param factory
 * @return
 */
public boolean removeStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory);
 

/**
 * Register a state model factory with a name for a state model definition
 * @param stateModelDef
 * @param factory
 * @param factoryName
 * @return
 */
public boolean registerStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory, String factoryName);
 

/**
 * Register a default state model factory for a state model definition
 * A state model definition could be, for example:
 * "MasterSlave", "OnlineOffline", "LeaderStandby", etc.
 * @param stateModelDef
 * @param factory
 * @return
 */
public boolean registerStateModelFactory(String stateModelDef,
    StateModelFactory<? extends StateModel> factory);