package org.apache.helix.controller.stages;

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.stream.Collectors;

import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixDefinedState;
import org.apache.helix.HelixException;
import org.apache.helix.HelixManager;
import org.apache.helix.SystemPropertyKeys;
import org.apache.helix.api.config.StateTransitionTimeoutConfig;
import org.apache.helix.controller.LogUtil;
import org.apache.helix.controller.common.ResourcesStateMap;
import org.apache.helix.controller.dataproviders.BaseControllerDataProvider;
import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
import org.apache.helix.controller.pipeline.AbstractBaseStage;
import org.apache.helix.controller.pipeline.StageException;
import org.apache.helix.manager.zk.DefaultSchedulerMessageHandlerFactory;
import org.apache.helix.model.ClusterConfig;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.LiveInstance;
import org.apache.helix.model.Message;
import org.apache.helix.model.Partition;
import org.apache.helix.model.Resource;
import org.apache.helix.model.ResourceConfig;
import org.apache.helix.model.StateModelDefinition;
import org.apache.helix.util.HelixUtil;
import org.apache.helix.util.MessageUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Compares the currentState, pendingState with IdealState and generate messages
 */
public class MessageGenerationPhase extends AbstractBaseStage {
  private static final String NO_DESIRED_STATE = "NoDesiredState";

  // If we see there is any invalid pending message leaving on host, i.e. message
  // tells participant to change from SLAVE to MASTER, and the participant is already
  // at MASTER state, we wait for timeout and if the message is still not cleaned up by
  // participant, controller will cleanup them proactively to unblock further state
  // transition
  public static final long DEFAULT_OBSELETE_MSG_PURGE_DELAY = HelixUtil
      .getSystemPropertyAsLong(SystemPropertyKeys.CONTROLLER_MESSAGE_PURGE_DELAY, 60 * 1000);
  private static final String PENDING_MESSAGE = "pending message";
  private static final String STALE_MESSAGE = "stale message";
  private static final String OFFLINE = "OFFLINE";

  private static Logger logger = LoggerFactory.getLogger(MessageGenerationPhase.class);

  @Override
  public void process(ClusterEvent event) throws Exception {
    BestPossibleStateOutput bestPossibleStateOutput =
        event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());
    _eventId = event.getEventId();
    HelixManager manager = event.getAttribute(AttributeName.helixmanager.name());
    BaseControllerDataProvider cache =
        event.getAttribute(AttributeName.ControllerDataProvider.name());
    Map<String, Resource> resourceMap =
        event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
    CurrentStateOutput currentStateOutput = event.getAttribute(AttributeName.CURRENT_STATE.name());

    Map<String, Map<String, Message>> messagesToCleanUp = new HashMap<>();
    if (manager == null || cache == null || resourceMap == null || currentStateOutput == null
        || bestPossibleStateOutput == null) {
      throw new StageException("Missing attributes in event:" + event
          + ". Requires HelixManager|DataCache|RESOURCES|CURRENT_STATE|BESTPOSSIBLE_STATE");
    }

    Map<String, LiveInstance> liveInstances = cache.getLiveInstances();
    Map<String, String> sessionIdMap = new HashMap<>();

    for (LiveInstance liveInstance : liveInstances.values()) {
      sessionIdMap.put(liveInstance.getInstanceName(), liveInstance.getEphemeralOwner());
    }
    MessageOutput output = new MessageOutput();

    for (Resource resource : resourceMap.values()) {
      try {
        generateMessage(resource, cache, bestPossibleStateOutput, currentStateOutput, manager,
            sessionIdMap, event.getEventType(), output, messagesToCleanUp);
      } catch (HelixException ex) {
        LogUtil.logError(logger, _eventId,
            "Failed to generate message for resource " + resource.getResourceName(), ex);
      }
    }

    // Asynchronously GC pending messages if necessary
    if (!messagesToCleanUp.isEmpty()) {
      schedulePendingMessageCleanUp(messagesToCleanUp, cache.getAsyncTasksThreadPool(),
          manager.getHelixDataAccessor());
    }
    event.addAttribute(AttributeName.MESSAGES_ALL.name(), output);
  }

  private void generateMessage(final Resource resource, final BaseControllerDataProvider cache,
      final ResourcesStateMap resourcesStateMap, final CurrentStateOutput currentStateOutput,
      final HelixManager manager, final Map<String, String> sessionIdMap,
      final ClusterEventType eventType, MessageOutput output,
      Map<String, Map<String, Message>> messagesToCleanUp) {
    String resourceName = resource.getResourceName();

    StateModelDefinition stateModelDef = cache.getStateModelDef(resource.getStateModelDefRef());
    if (stateModelDef == null) {
      LogUtil.logError(logger, _eventId,
          "State Model Definition null, skip generating messages for resource: " + resourceName);
      return;
    }

    for (Partition partition : resource.getPartitions()) {
      Map<String, String> instanceStateMap =
          new HashMap<>(resourcesStateMap.getInstanceStateMap(resourceName, partition));
      Map<String, String> pendingStateMap =
          currentStateOutput.getPendingStateMap(resourceName, partition);
      Map<String, String> currentStateMap =
          currentStateOutput.getCurrentStateMap(resourceName, partition);

      // The operation is combing pending state with best possible state. Since some replicas have
      // been moved from one instance to another, the instance will exist in pending state but not
      // best possible. Thus Helix need to check the pending state and cancel it.

      for (String instance : pendingStateMap.keySet()) {
        if (!instanceStateMap.containsKey(instance)) {
          instanceStateMap.put(instance, NO_DESIRED_STATE);
        }
      }

      // Look through the current state map and add DROPPED message if the instance is not in the
      // resourceStateMap. This instance may not have had been dropped by the rebalance strategy.
      // This check is required to ensure that the instances removed from the ideal state stateMap
      // are properly dropped.
      // This should only solve for instance operation case where the instance is removed from the statemap but there
      // are still valid assignments in the mapping. We should not consider case where there is no mapping at all for
      // the resource, which can occur on a rebalance failure. If the resource has been removed, the partition has
      // been removed, or the replication factor has been reduced the BP will contain the DROPPED states.See method
      // AbstractRebalancer.computeBestPossibleMap - drops replica that is in current state but not in preference list.
      for (String instance : currentStateMap.keySet()) {
        if (!instanceStateMap.isEmpty() && !instanceStateMap.containsKey(instance)) {
          instanceStateMap.put(instance, HelixDefinedState.DROPPED.name());
        }
      }

      // we should generate message based on the desired-state priority
      // so keep generated messages in a temp map keyed by state
      // desired-state->list of generated-messages
      Map<String, List<Message>> messageMap = new HashMap<>();

      /*
       * Calculate the current active replica count based on state model type.
       * This represents the number of replicas currently serving traffic for this partition
       * Active replicas include: top states, secondary top states(excluding OFFLINE) and ERROR
       * states.
       * Active replicas exclude: OFFLINE and DROPPED states.
       * All qualifying state transitions for this partition will receive this same value,
       * allowing clients to understand the current availability level and prioritize accordingly.
       */
      int currentActiveReplicaCount =
          calculateCurrentActiveReplicaCount(currentStateMap, stateModelDef);

      for (String instanceName : instanceStateMap.keySet()) {

        Set<Message> staleMessages = cache.getStaleMessagesByInstance(instanceName);

        String desiredState = instanceStateMap.get(instanceName);

        String currentState =
            currentStateOutput.getCurrentState(resourceName, partition, instanceName);
        Message pendingMessage =
            currentStateOutput.getPendingMessage(resourceName, partition, instanceName);
        boolean isCancellationEnabled = cache.getClusterConfig().isStateTransitionCancelEnabled();
        Message cancellationMessage =
            currentStateOutput.getCancellationMessage(resourceName, partition, instanceName);
        String nextState = stateModelDef.getNextStateForTransition(currentState, desiredState);

        Message message = null;

        if (currentState == null) {
          currentState = stateModelDef.getInitialState();
          nextState = stateModelDef.getNextStateForTransition(currentState, desiredState);

          if (desiredState.equals(HelixDefinedState.DROPPED.name())) {
            LogUtil.logDebug(logger, _eventId, String
                .format("No current state for partition %s in resource %s, skip the drop message",
                    partition.getPartitionName(), resourceName));

            message =
                generateCancellationMessageForPendingMessage(desiredState, currentState, nextState,
                    pendingMessage, manager, resource, partition, sessionIdMap, instanceName,
                    stateModelDef, cancellationMessage, isCancellationEnabled);
            addGeneratedMessageToMap(message, messageMap, eventType, cache, desiredState,
                resourceName, partition, currentState, nextState);

            // TODO: separate logic of resource/task message generation
            if (cache instanceof ResourceControllerDataProvider) {
              ((ResourceControllerDataProvider) cache)
                  .invalidateCachedIdealStateMapping(resourceName);
            }
            continue;
          }
        }

        if (shouldCleanUpPendingMessage(pendingMessage, sessionIdMap, instanceName, currentState,
            currentStateOutput.getEndTime(resourceName, partition, instanceName))) {
          logAndAddToCleanUp(messagesToCleanUp, pendingMessage, instanceName, resourceName,
              partition, currentState, PENDING_MESSAGE);
        }

        for (Message staleMessage : staleMessages) {
          // staleMessage can be simple or batch mode
          if ((System.currentTimeMillis() - currentStateOutput
              .getEndTime(resourceName, partition, instanceName) > DEFAULT_OBSELETE_MSG_PURGE_DELAY)
              && staleMessage.getResourceName().equals(resourceName) && sessionIdMap
              .containsKey(instanceName) && (
              staleMessage.getPartitionName().equals(partition.getPartitionName()) || (
                  staleMessage.getBatchMessageMode() && staleMessage.getPartitionNames()
                      .contains(partition.getPartitionName())))) {
            logAndAddToCleanUp(messagesToCleanUp, staleMessage, instanceName, resourceName,
                partition, currentState, STALE_MESSAGE);
          }
        }

        if (desiredState.equals(NO_DESIRED_STATE) || desiredState.equalsIgnoreCase(currentState)) {
          if (shouldCreateSTCancellation(pendingMessage, desiredState,
              stateModelDef.getInitialState())) {
            message = MessageUtil
                .createStateTransitionCancellationMessage(manager.getInstanceName(),
                    manager.getSessionId(), resource, partition.getPartitionName(), instanceName,
                    sessionIdMap.get(instanceName), stateModelDef.getId(),
                    pendingMessage.getFromState(), pendingMessage.getToState(), null,
                    cancellationMessage, isCancellationEnabled, currentState);
          }
        } else {
          if (nextState == null) {
            LogUtil.logError(logger, _eventId,
                "Unable to find a next state for resource: " + resource.getResourceName()
                    + " partition: " + partition.getPartitionName() + " from stateModelDefinition"
                    + stateModelDef.getClass() + " from:" + currentState + " to:" + desiredState);
            continue;
          }

          if (pendingMessage != null) {
            message =
                generateCancellationMessageForPendingMessage(desiredState, currentState, nextState,
                    pendingMessage, manager, resource, partition, sessionIdMap, instanceName,
                    stateModelDef, cancellationMessage, isCancellationEnabled);
          } else {
            // Set currentActiveReplicaNumber to provide metadata for potential message
            // prioritization by participant.
            // Assign the current active replica count to all qualifying upward transitions for this
            // partition.
            // This ensures consistent prioritization metadata across concurrent state transitions.
            // -1 indicates no prioritization metadata, for eg:Downward ST messages get a -1.
            int currentActiveReplicaNumber = -1;

            /*
             * Assign currentActiveReplicaNumber for qualifying upward state transitions.
             * Criteria for assignment:
             * - Must be an upward state transition according to state model
             * - Target state must be considered active (according to state model type)
             */
            if (stateModelDef.isUpwardStateTransition(currentState, nextState)
                && isStateActive(nextState, stateModelDef)) {

              // All qualifying transitions for this partition get the same
              // currentActiveReplicaNumber
              currentActiveReplicaNumber = currentActiveReplicaCount;
            }

            // Create new state transition message
            message = MessageUtil.createStateTransitionMessage(manager.getInstanceName(),
                manager.getSessionId(), resource, partition.getPartitionName(), instanceName,
                currentState, nextState, sessionIdMap.get(instanceName), stateModelDef.getId(),
                currentActiveReplicaNumber);

            if (logger.isDebugEnabled()) {
              LogUtil.logDebug(logger, _eventId, String.format(
                  "Resource %s partition %s for instance %s with currentState %s, nextState %s and currentActiveReplicaNumber %d",
                  resource.getResourceName(), partition.getPartitionName(), instanceName,
                  currentState, nextState, currentActiveReplicaNumber));
            }
          }
        }
        addGeneratedMessageToMap(message, messageMap, eventType, cache, desiredState, resourceName,
            partition, currentState, nextState);
      }

      // add generated messages to output according to state priority
      List<String> statesPriorityList = stateModelDef.getStatesPriorityList();
      for (String state : statesPriorityList) {
        if (messageMap.containsKey(state)) {
          for (Message message : messageMap.get(state)) {
            // This is for a bug where a message's target session id is null
            if (!message.isValid()) {
              LogUtil.logError(logger, _eventId, String.format(
                  "An invalid message was generated! Discarding this message. sessionIdMap: %s, CurrentStateMap: %s, InstanceStateMap: %s, AllInstances: %s, LiveInstances: %s, Message: %s",
                  sessionIdMap, currentStateOutput.getCurrentStateMap(resourceName, partition),
                  instanceStateMap, cache.getAllInstances(),
                  cache.getLiveInstances().keySet(),
                  message));
              continue; // Do not add this message
            }
            output.addMessage(resourceName, partition, message);
          }
        }
      }
    } // end of for-each-partition
  }

  /**
   * Calculate the current active replica count based on state model type.
   * The count includes replicas in top states, secondary top states (excluding OFFLINE),
   * and ERROR states since helix considers them active.Count excludes OFFLINE and DROPPED states.
   * @param currentStateMap
   * @param stateModelDef
   * @return The number of replicas currently in active states, used to determine the
   *         currentActiveReplicaNumber for the partition.
   */
  private int calculateCurrentActiveReplicaCount(Map<String, String> currentStateMap,
      StateModelDefinition stateModelDef) {
    return (int) currentStateMap.values().stream()
        .filter(state -> stateModelDef.getTopState().contains(state) // Top states (MASTER, ONLINE,
                                                                     // LEADER)
            || getActiveSecondaryTopStates(stateModelDef).contains(state) // Active secondary states
                                                                          // (SLAVE, STANDBY,
                                                                          // BOOTSTRAP)
            || HelixDefinedState.ERROR.name().equals(state) // ERROR states (still considered
                                                            // active)
        // DROPPED and OFFLINE are automatically excluded by getActiveSecondaryTopStates()
        ).count();
  }

  /**
   * Get active secondary top states - states that are not non-serving states like OFFLINE and
   * DROPPED.
   * Reasons for elimination:
   * - getSecondTopStates() can include OFFLINE as a secondary top state in some state models.
   * Example - OnlineOffline:
   * getSecondTopStates() = ["OFFLINE"] as it transitions to ONLINE.
   * After filtering: activeSecondaryTopStates=[] (removes "OFFLINE" as it's not a serving state).
   * @param stateModelDef
   */
  private List<String> getActiveSecondaryTopStates(StateModelDefinition stateModelDef) {
    return stateModelDef.getSecondTopStates().stream()
        // Remove non-serving states
        .filter(state -> !OFFLINE.equals(state) && !HelixDefinedState.DROPPED.name().equals(state))
        .collect(Collectors.toList());
  }

  /**
   * Determines if the given state is considered active based on the state model type.
   * Active states include: top states, active secondary top states (excluding OFFLINE),
   * and ERROR states. Active states exclude OFFLINE and DROPPED states.
   * ERROR state replicas are always considered active in HELIX as they do not
   * affect availability.
   * @param state
   * @param stateModelDef
   * @return true if the state is considered active, false otherwise
   */
  private boolean isStateActive(String state, StateModelDefinition stateModelDef) {
    // ERROR state is always considered active regardless of state model type
    if (HelixDefinedState.ERROR.name().equals(state)) {
      return true;
    }
    return stateModelDef.getTopState().contains(state)
        || getActiveSecondaryTopStates(stateModelDef).contains(state);
  }

    private boolean shouldCreateSTCancellation(Message pendingMessage, String desiredState,
      String initialState) {
    if (pendingMessage == null) {
      return false;
    }
    if (NO_DESIRED_STATE.equals(desiredState)) {
      return true;
    }

    // Cancel the ST except below scenarios:
    // 1. pending message toState is desired state
    // 2. pending message is an ERROR reset: ERROR -> initState (eg. OFFLINE)
    return !desiredState.equalsIgnoreCase(pendingMessage.getToState()) && !(
        HelixDefinedState.ERROR.name().equals(pendingMessage.getFromState()) && initialState
            .equals(pendingMessage.getToState()));
  }

  private void logAndAddToCleanUp(Map<String, Map<String, Message>> messagesToCleanUp,
      Message message, String instanceName, String resourceName, Partition partition,
      String currentState, String cleanUpMessageType) {
    String logMsg = String.format(
        "Adding %s %s on instance %s to clean up. Msg: %s->%s, current state"
            + " of resource %s:%s is %s", cleanUpMessageType, message.getMsgId(), instanceName,
        message.getFromState(), message.getToState(), resourceName, partition, currentState);
    LogUtil.logInfo(logger, _eventId, logMsg);
    if (!messagesToCleanUp.containsKey(instanceName)) {
      messagesToCleanUp.put(instanceName, new HashMap<String, Message>());
    }
    messagesToCleanUp.get(instanceName).put(message.getMsgId(), message);
  }

  private Message generateCancellationMessageForPendingMessage(final String desiredState,
      final String currentState, final String nextState, final Message pendingMessage,
      final HelixManager manager, final Resource resource, final Partition partition,
      final Map<String, String> sessionIdMap, final String instanceName,
      final StateModelDefinition stateModelDef, final Message cancellationMessage,
      final boolean isCancellationEnabled) {

    Message message = null;

    if (pendingMessage != null) {
      String pendingState = pendingMessage.getToState();
      if (nextState.equalsIgnoreCase(pendingState)) {
        LogUtil.logInfo(logger, _eventId,
            "Message already exists for " + instanceName + " to transit " + resource
                .getResourceName() + "." + partition.getPartitionName() + " from " + currentState
                + " to " + nextState + ", isRelay: " + pendingMessage.isRelayMessage());
      } else if (currentState.equalsIgnoreCase(pendingState)) {
        LogUtil.logDebug(logger, _eventId,
            "Message hasn't been removed for " + instanceName + " to transit " + resource
                .getResourceName() + "." + partition.getPartitionName() + " to " + pendingState
                + ", desiredState: " + desiredState + ", isRelay: " + pendingMessage
                .isRelayMessage());
      } else {
        LogUtil.logDebug(logger, _eventId,
            "IdealState changed before state transition completes for " + resource.getResourceName()
                + "." + partition.getPartitionName() + " on " + instanceName + ", pendingState: "
                + pendingState + ", currentState: " + currentState + ", nextState: " + nextState
                + ", isRelay: " + pendingMessage.isRelayMessage());

        message = MessageUtil.createStateTransitionCancellationMessage(manager.getInstanceName(),
            manager.getSessionId(), resource, partition.getPartitionName(), instanceName,
            sessionIdMap.get(instanceName), stateModelDef.getId(), pendingMessage.getFromState(),
            pendingState, nextState, cancellationMessage, isCancellationEnabled, currentState);
      }
    }
    return message;
  }

  private void addGeneratedMessageToMap(final Message message,
      Map<String, List<Message>> messageMap, final ClusterEventType eventType,
      final BaseControllerDataProvider cache, final String desiredState, final String resourceName,
      final Partition partition, final String currentState, final String nextState) {
    if (message != null) {
      IdealState idealState = cache.getIdealState(resourceName);
      if (idealState != null && idealState.getStateModelDefRef()
          .equalsIgnoreCase(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE)) {
        if (idealState.getRecord().getMapField(partition.getPartitionName()) != null) {
          message.getRecord().setMapField(Message.Attributes.INNER_MESSAGE.toString(),
              idealState.getRecord().getMapField(partition.getPartitionName()));
        }
      }

      int timeout =
          getTimeOut(cache.getClusterConfig(), cache.getResourceConfig(resourceName), currentState,
              nextState, idealState, partition);
      if (timeout > 0) {
        message.setExecutionTimeout(timeout);
      }

      message.setAttribute(Message.Attributes.ClusterEventName, eventType.name());
      // output.addMessage(resourceName, partition, message);
      if (!messageMap.containsKey(desiredState)) {
        messageMap.put(desiredState, new ArrayList<Message>());
      }
      messageMap.get(desiredState).add(message);
    }
  }

  /**
   * Start a job in worker pool that asynchronously clean up pending message. Since it is possible
   * that participant failed to clean up message after processing, it is important for controller to
   * try to clean them up as well to unblock further rebalance
   *
   * @param pendingMessagesToPurge key: instance name, value: list of pending message to cleanup
   * @param workerPool             ExecutorService that job can be submitted to
   * @param accessor               Data accessor used to clean up message
   */
  private void schedulePendingMessageCleanUp(
      final Map<String, Map<String, Message>> pendingMessagesToPurge, ExecutorService workerPool,
      final HelixDataAccessor accessor) {
    workerPool.submit(new Callable<Object>() {
      @Override
      public Object call() {
        for (Map.Entry<String, Map<String, Message>> entry : pendingMessagesToPurge.entrySet()) {
          String instanceName = entry.getKey();
          for (Message msg : entry.getValue().values()) {
            if (accessor.removeProperty(msg.getKey(accessor.keyBuilder(), instanceName))) {
              LogUtil.logInfo(logger, _eventId, String
                  .format("Deleted message %s from instance %s", msg.getMsgId(), instanceName));
            }
          }
        }
        return null;
      }
    });
  }

  private boolean shouldCleanUpPendingMessage(Message pendingMsg, Map<String, String> sessionIdMap,
      String instanceName, String currentState, Long currentStateTransitionEndTime) {
    if (pendingMsg == null || !sessionIdMap.containsKey(instanceName)) {
      return false;
    }
    if (currentState.equalsIgnoreCase(pendingMsg.getToState())) {
      // If pending message's toState is same as current state, state transition is finished
      // successfully. In this case, we will wait for a timeout for participant to cleanup
      // processed message. If participant fail to do so, controller is going to proactively delete
      // the message as participant does not retry message deletion upon failure.
      return System.currentTimeMillis() - currentStateTransitionEndTime
          > DEFAULT_OBSELETE_MSG_PURGE_DELAY;
    } else {
      // Partition's current state should be either pending message's fromState or toState or
      // the message is invalid and can be safely deleted immediately.
      return !currentState.equalsIgnoreCase(pendingMsg.getFromState());
    }
  }

  private int getTimeOut(ClusterConfig clusterConfig, ResourceConfig resourceConfig,
      String currentState, String nextState, IdealState idealState, Partition partition) {
    StateTransitionTimeoutConfig stateTransitionTimeoutConfig =
        clusterConfig.getStateTransitionTimeoutConfig();
    int timeout = stateTransitionTimeoutConfig != null ? stateTransitionTimeoutConfig
        .getStateTransitionTimeout(currentState, nextState) : -1;

    String timeOutStr = null;
    // Check IdealState whether has timeout set
    if (idealState != null) {
      String stateTransition = currentState + "-" + nextState + "_" + Message.Attributes.TIMEOUT;
      timeOutStr = idealState.getRecord().getSimpleField(stateTransition);
      if (timeOutStr == null && idealState.getStateModelDefRef()
          .equalsIgnoreCase(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE)) {
        // scheduled task queue
        if (idealState.getRecord().getMapField(partition.getPartitionName()) != null) {
          timeOutStr = idealState.getRecord().getMapField(partition.getPartitionName())
              .get(Message.Attributes.TIMEOUT.toString());
        }
      }
    }
    if (timeOutStr != null) {
      try {
        timeout = Integer.parseInt(timeOutStr);
      } catch (Exception e) {
        LogUtil.logError(logger, _eventId, "", e);
      }
    }

    if (resourceConfig != null) {
      // If resource config has timeout, replace the cluster timeout.
      stateTransitionTimeoutConfig = resourceConfig.getStateTransitionTimeoutConfig();
      timeout = stateTransitionTimeoutConfig != null ? stateTransitionTimeoutConfig
          .getStateTransitionTimeout(currentState, nextState) : -1;
    }

    return timeout;
  }
}
