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

/**
 * Create and return an instance of {@link StateModel} based on given definition
 * @param resourceName the resource name for which this state model is being created.
 * @param partitionName the partition name for which this state model is being created.
 * @return an instance of {@link StateModel}.
 */
@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  StateModel stateModelToReturn;
  switch (clustermapConfig.clustermapStateModelDefinition) {
    case AmbryStateModelDefinition.AMBRY_LEADER_STANDBY_MODEL:
      stateModelToReturn =
          new AmbryPartitionStateModel(resourceName, partitionName, partitionStateChangeListener, clustermapConfig);
      break;
    case LeaderStandbySMD.name:
      stateModelToReturn = new DefaultLeaderStandbyStateModel();
      break;
    default:
      // Code won't get here because state model def is already validated in ClusterMapConfig. We keep exception here
      // in case the validation logic is changed in ClusterMapConfig.
      throw new IllegalArgumentException(
          "Unsupported state model definition: " + clustermapConfig.clustermapStateModelDefinition);
  }
  return stateModelToReturn;
}
 

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

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

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

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

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

@Override public StateModel createNewStateModel(String resourceName, String partitionName) {
  LeaderStandbyStateModel stateModel = new LeaderStandbyStateModel();
  stateModel.setDelay(_delay);
  stateModel.setInstanceName(_instanceName);
  stateModel.setPartitionName(partitionName);
  return stateModel;
}
 

@Override
public StateModel createNewStateModel(String resourceName, String stateUnitKey) {
  OnlineOfflineStateModel stateModel = new OnlineOfflineStateModel();
  stateModel.setDelay(_delay);
  stateModel.setInstanceName(_instanceName);
  return stateModel;
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  MasterSlaveStateModel stateModel = new MasterSlaveStateModel();
  stateModel.setInstanceName(_instanceName);
  stateModel.setDelay(_delay);
  stateModel.setPartitionName(partitionName);
  return stateModel;
}
 

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

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

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

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

public HelixStateTransitionCancellationHandler(StateModel stateModel, Message message,
    NotificationContext context) {
  super(message, context);
  _stateModel = stateModel;
}
 

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

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return STATE_MODEL_INSTANCE;
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return new FakeSegmentStateModel();
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return STATE_MODEL_INSTANCE;
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return new BrokerResourceOnlineOfflineStateModel();
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return new SegmentOnlineOfflineStateModel();
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partitionName) {
  return new LeadControllerResourceMasterSlaveStateModel(_leadControllerManager, partitionName);
}
 

@Override
public StateModel createNewStateModel(String resourceName, String partition) {
  LOG.info("Received request to process partition = " + partition + ", for " +
          "resource = " + resourceName + ", at " + instanceName);
  return new OnlineOfflineStateModel();
}
 

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

@Override
public StateModel createNewStateModel(String resourceName, String stateUnitKey) {
  return new BootstrapStateModel(stateUnitKey);
}
 

@Override
public StateModel createNewStateModel(String partitionName) {
  final SegmentOnlineOfflineStateModel SegmentOnlineOfflineStateModel = new SegmentOnlineOfflineStateModel();
  return SegmentOnlineOfflineStateModel;
}
 

public StateModel createNewStateModel(String resourceName) {
  return new BrokerResourceOnlineOfflineStateModelFactory.BrokerResourceOnlineOfflineStateModel();
}
 

@Override
public StateModel createNewStateModel(String resource, String stateModelUnit) {
  OnlineOfflineStateModel stateModel = new OnlineOfflineStateModel();
  return stateModel;
}