Update based on review/feedback

Author: Charanya Sudharsanan

helix-core/src/main/java/org/apache/helix/controller/stages/MessageGenerationPhase.java

@@ -164,14 +164,14 @@ private void generateMessage(final Resource resource, final BaseControllerDataPr
       Map<String, List<Message>> messageMap = new HashMap<>();
 
       // Calculate metadata for message prioritization
-      int totalInstanceCount = instanceStateMap.keySet().size();
+      int partitionInstanceCount = instanceStateMap.keySet().size();
 
       // Get pending upward state transition messages to second top or top state
-      List<Message> pendingUpwardStateTransitionMessages = getPendingUpwardStateTransitionMessages(
+      List<Message> pendingUpwardStateTransitionMessagesToTopOrSecondTopStates = getPendingTransitionsToTopOrSecondTopStates(
           resourceName, partition, currentStateOutput, stateModelDef);
 
       // Initialize replica counter for prioritization
-      int initialReplicaNumber = totalInstanceCount - pendingUpwardStateTransitionMessages.size();
+      int initialReplicaNumber = partitionInstanceCount - pendingUpwardStateTransitionMessagesToTopOrSecondTopStates.size();
 
       for (String instanceName : instanceStateMap.keySet()) {
 
@@ -260,16 +260,14 @@ private void generateMessage(final Resource resource, final BaseControllerDataPr
                     pendingMessage, manager, resource, partition, sessionIdMap, instanceName,
                     stateModelDef, cancellationMessage, isCancellationEnabled);
           } else {
-            // Default currentReplicaNumber is -1 (no prioritization)
+            // Default currentReplicaNumber is -1 (provides metadata for participant-side prioritization)
             int currentReplicaNumber = -1;
             // Check if this is an upward state transition from non-second top state to second top
             // or top state
-            if (isUpwardStateTransition(currentState, nextState, stateModelDef)
+            if (stateModelDef.isUpwardStateTransition(currentState, nextState)
                 && !stateModelDef.getSecondTopStates().contains(currentState)
                 && (isSecondTopState(nextState, stateModelDef)
-                    || isTopState(nextState, stateModelDef))
-                && !isInPendingMessages(resourceName, partition, instanceName, currentState,
-                    nextState, pendingUpwardStateTransitionMessages)) {
+                    || isTopState(nextState, stateModelDef))) {
 
               // Assign the replica number for prioritization
               currentReplicaNumber = initialReplicaNumber--;
@@ -335,57 +333,6 @@ private boolean isSecondTopState(String state, StateModelDefinition stateModelDe
     return stateModelDef.getSecondTopStates().contains(state);
   }
 
-  /**
-   * Check if a state transition is already in the pending messages
-   * @param resourceName The resource name
-   * @param partition The partition
-   * @param instanceName The instance name
-   * @param fromState The from state
-   * @param toState The to state
-   * @param pendingMessages The list of pending messages
-   * @return True if the state transition is already in pending messages, false otherwise
-   */
-  private boolean isInPendingMessages(String resourceName, Partition partition, String instanceName,
-      String fromState, String toState, List<Message> pendingMessages) {
-
-    for (Message message : pendingMessages) {
-      if (message.getResourceName().equals(resourceName)
-          && message.getPartitionName().equals(partition.getPartitionName())
-          && message.getTgtName().equals(instanceName) && message.getFromState().equals(fromState)
-          && message.getToState().equals(toState)) {
-        return true;
-      }
-    }
-
-    return false;
-  }
-
-  /**
-   * Check if a state transition is upward
-   * @param fromState The from state
-   * @param toState The to state
-   * @param stateModelDef The state model definition
-   * @return True if it's an upward state transition, false otherwise
-   */
-  private boolean isUpwardStateTransition(String fromState, String toState,
-      StateModelDefinition stateModelDef) {
-
-    if (fromState == null || toState == null) {
-      return false;
-    }
-
-    Map<String, Integer> statePriorityMap = stateModelDef.getStatePriorityMap();
-
-    Integer fromStateWeight = statePriorityMap.get(fromState);
-    Integer toStateWeight = statePriorityMap.get(toState);
-
-    if (fromStateWeight == null || toStateWeight == null) {
-      return false;
-    }
-
-    return toStateWeight < fromStateWeight;
-  }
-
   /**
    * Get pending upward state transition messages from non-second top state to second top or top
    * state
@@ -395,7 +342,7 @@ private boolean isUpwardStateTransition(String fromState, String toState,
    * @param stateModelDef The state model definition
    * @return List of pending messages for upward state transitions
    */
-  private List<Message> getPendingUpwardStateTransitionMessages(String resourceName,
+  private List<Message> getPendingTransitionsToTopOrSecondTopStates(String resourceName,
       Partition partition, CurrentStateOutput currentStateOutput,
       StateModelDefinition stateModelDef) {
     List<Message> pendingUpwardSTMessages = new ArrayList<>();
@@ -411,7 +358,7 @@ private List<Message> getPendingUpwardStateTransitionMessages(String resourceNam
 
         // Check if it's an upward state transition from non-second top state to second top or top
         // state
-        if (isUpwardStateTransition(fromState, toState, stateModelDef)
+        if (stateModelDef.isUpwardStateTransition(fromState, toState)
             && !isSecondTopState(fromState, stateModelDef)
             && (isSecondTopState(toState, stateModelDef) || isTopState(toState, stateModelDef))) {
           pendingUpwardSTMessages.add(message);

helix-core/src/main/java/org/apache/helix/model/Message.java

@@ -138,12 +138,8 @@ public enum STRebalanceType {
   /**
    * Compares the creation time of two Messages
    */
-  public static final Comparator<Message> CREATE_TIME_COMPARATOR = new Comparator<Message>() {
-    @Override
-    public int compare(Message m1, Message m2) {
-      return new Long(m1.getCreateTimeStamp()).compareTo(new Long(m2.getCreateTimeStamp()));
-    }
-  };
+  public static final Comparator<Message> CREATE_TIME_COMPARATOR =
+      (m1, m2) -> Long.compare(m2.getCreateTimeStamp(), m1.getCreateTimeStamp());
 
   /**
    * Instantiate a message
@@ -937,8 +933,14 @@ public void setSrcClusterName(String clusterName) {
   }
 
   /**
-   * Set the currentReplicaNumber for transition-related messages
-   * @param currentReplicaNumber the replica count
+   * Set replica number for participant-side message prioritization.
+   * This field indicates the number of top and second-top state replicas at the time a state transition message is generated.
+   * It is used to prioritize messages, with lower values indicating higher priority.
+   * Participants can use this in custom thread pools or message handlers to process
+   * critical transitions first during recovery scenarios.
+   * Default value is -1 for transitions that don't require prioritization.
+   * @param currentReplicaNumber the replica priority number (-1 for no prioritization, >=0 for
+   *          prioritized)
    */
   public void setCurrentReplicaNumber(int currentReplicaNumber) {
     _record.setIntField(Attributes.CURRENT_REPLICA_NUMBER.name(), currentReplicaNumber);

helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java

@@ -505,4 +505,28 @@ public static Map<String, Integer> getStateCounts(Map<String, String> stateMap)
     }
     return stateCounts;
   }
+
+  /**
+   * Check if a state transition is upward
+   * @param fromState source state
+   * @param toState destination state
+   * @return True if it's an upward state transition, false otherwise
+   */
+  public boolean isUpwardStateTransition(String fromState, String toState) {
+
+    if (fromState == null || toState == null) {
+      return false;
+    }
+
+    Map<String, Integer> statePriorityMap = getStatePriorityMap();
+
+    Integer fromStateWeight = statePriorityMap.get(fromState);
+    Integer toStateWeight = statePriorityMap.get(toState);
+
+    if (fromStateWeight == null || toStateWeight == null) {
+      return false;
+    }
+
+    return toStateWeight < fromStateWeight;
+  }
 }

helix-core/src/main/java/org/apache/helix/util/MessageUtil.java

@@ -48,7 +48,7 @@ public static Message createStateTransitionCancellationMessage(String srcInstanc
           toState);
 
       Message message =
-          createStateTransitionMessage(Message.MessageType.STATE_TRANSITION_CANCELLATION,
+          createBasicStateTransitionMessage(Message.MessageType.STATE_TRANSITION_CANCELLATION,
               srcInstanceName, srcSessionId, resource, partitionName, instanceName, currentState,
               nextState, sessionId, stateModelDefName);
 
@@ -60,28 +60,6 @@ public static Message createStateTransitionCancellationMessage(String srcInstanc
     return null;
   }
 
-  public static Message createStateTransitionMessage(String srcInstanceName, String srcSessionId,
-      Resource resource, String partitionName, String instanceName, String currentState,
-      String nextState, String tgtSessionId, String stateModelDefName) {
-    Message message =
-        createStateTransitionMessage(Message.MessageType.STATE_TRANSITION, srcInstanceName,
-            srcSessionId, resource, partitionName, instanceName, currentState, nextState, tgtSessionId,
-            stateModelDefName);
-
-    // Set the retry count for state transition messages.
-    // TODO: make the retry count configurable in ClusterConfig or IdealState
-    message.setRetryCount(DEFAULT_STATE_TRANSITION_MESSAGE_RETRY_COUNT);
-
-    if (resource.getResourceGroupName() != null) {
-      message.setResourceGroupName(resource.getResourceGroupName());
-    }
-    if (resource.getResourceTag() != null) {
-      message.setResourceTag(resource.getResourceTag());
-    }
-
-    return message;
-  }
-
   /**
    * Creates a message to change participant status
    * {@link org.apache.helix.model.LiveInstance.LiveInstanceStatus}
@@ -121,7 +99,7 @@ private static Message createBasicMessage(Message.MessageType messageType, Strin
   }
 
   /* Creates state transition or state transition cancellation message */
-  private static Message createStateTransitionMessage(Message.MessageType messageType,
+  private static Message createBasicStateTransitionMessage(Message.MessageType messageType,
       String srcInstanceName, String srcSessionId, Resource resource, String partitionName,
       String instanceName, String currentState, String nextState, String tgtSessionId,
       String stateModelDefName) {
@@ -138,30 +116,62 @@ private static Message createStateTransitionMessage(Message.MessageType messageT
   }
 
   /**
-   * Create a state transition message with currentReplicaNumber for prioritization
-   * @param msgSender message sender
-   * @param sessionId session id
+   * Create a state transition message with replica prioritization metadata
+   * @param srcInstanceName source instance name
+   * @param srcSessionId source session id
    * @param resource resource
    * @param partitionName partition name
-   * @param instanceName instance name
-   * @param fromState from state
-   * @param toState to state
-   * @param sessionIdForInstance session id for instance
-   * @param stateModelDefName state model def name
-   * @param currentReplicaNumber the current replica number (for prioritization)
-   * @return message
+   * @param instanceName target instance name
+   * @param currentState current state
+   * @param nextState next state
+   * @param tgtSessionId target session id
+   * @param stateModelDefName state model definition name
+   * @param currentReplicaNumber replica priority number (-1 for no prioritization, >=0 for
+   *          prioritized)
+   * @return state transition message
    */
-  public static Message createStateTransitionMessage(String msgSender, String sessionId,
-      Resource resource, String partitionName, String instanceName, String fromState,
-      String toState, String sessionIdForInstance, String stateModelDefName,
-      int currentReplicaNumber) {
+  public static Message createStateTransitionMessage(String srcInstanceName, String srcSessionId,
+      Resource resource, String partitionName, String instanceName, String currentState,
+      String nextState, String tgtSessionId, String stateModelDefName, int currentReplicaNumber) {
+    Message message = createBasicStateTransitionMessage(Message.MessageType.STATE_TRANSITION,
+        srcInstanceName, srcSessionId, resource, partitionName, instanceName, currentState,
+        nextState, tgtSessionId, stateModelDefName);
+
+    // Set the retry count for state transition messages.
+    // TODO: make the retry count configurable in ClusterConfig or IdealState
+    message.setRetryCount(DEFAULT_STATE_TRANSITION_MESSAGE_RETRY_COUNT);
 
-    Message message = createStateTransitionMessage(msgSender, sessionId, resource, partitionName,
-        instanceName, fromState, toState, sessionIdForInstance, stateModelDefName);
+    if (resource.getResourceGroupName() != null) {
+      message.setResourceGroupName(resource.getResourceGroupName());
+    }
+    if (resource.getResourceTag() != null) {
+      message.setResourceTag(resource.getResourceTag());
+    }
 
-    // Set the current replica number for prioritization
+    // Set replica number for participant-side prioritization
     message.setCurrentReplicaNumber(currentReplicaNumber);
 
     return message;
   }
+
+  /**
+   * Create a state transition message (backward compatibility)
+   * @param srcInstanceName source instance name
+   * @param srcSessionId source session id
+   * @param resource resource
+   * @param partitionName partition name
+   * @param instanceName target instance name
+   * @param currentState current state
+   * @param nextState next state
+   * @param tgtSessionId target session id
+   * @param stateModelDefName state model definition name
+   * @return state transition message
+   */
+  public static Message createStateTransitionMessage(String srcInstanceName, String srcSessionId,
+      Resource resource, String partitionName, String instanceName, String currentState,
+      String nextState, String tgtSessionId, String stateModelDefName) {
+    // currentReplicaNumber is set to -1 for ST messages needing no prioritization metadata (backward compatibility)
+    return createStateTransitionMessage(srcInstanceName, srcSessionId, resource, partitionName,
+        instanceName, currentState, nextState, tgtSessionId, stateModelDefName, -1);
+  }
 }