Auto Exit MM with timeout

Author: LZD-PratyushBhatt

helix-core/src/main/java/org/apache/helix/HelixAdmin.java

@@ -422,6 +422,19 @@ void autoEnableMaintenanceMode(String clusterName, boolean enabled, String reaso
   void manuallyEnableMaintenanceMode(String clusterName, boolean enabled, String reason,
       Map<String, String> customFields);
 
+  /**
+   * Manually enable maintenance mode with timeout. To be called by the REST client.
+   * The cluster will automatically exit maintenance mode after the specified timeout.
+   * @param clusterName the cluster name
+   * @param enabled if true, enter maintenance mode; if false, exit maintenance mode
+   * @param reason reason to enter maintenance mode
+   * @param timeout time in milliseconds after which maintenance mode should be exited automatically.
+   *                Only applicable when enabled is true. Set to -1 for no automatic exit.
+   * @param customFields user-specified KV mappings to be stored in the ZNode
+   */
+  void manuallyEnableMaintenanceModeWithTimeout(String clusterName, boolean enabled, String reason,
+      long timeout, Map<String, String> customFields);
+
   /**
    * Check specific cluster is in maintenance mode or not
    * @param clusterName the cluster name

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

@@ -63,17 +63,33 @@ public void execute(final ClusterEvent event) throws Exception {
     // Check for the maintenance signal
     // If it was entered manually or the signal is null (which shouldn't happen), skip this stage
     MaintenanceSignal maintenanceSignal = cache.getMaintenanceSignal();
-    if (maintenanceSignal == null || maintenanceSignal
-        .getTriggeringEntity() != MaintenanceSignal.TriggeringEntity.CONTROLLER) {
+    if (maintenanceSignal == null) {
       return;
     }
-
     HelixManager manager = event.getAttribute(AttributeName.helixmanager.name());
     if (manager == null || !manager.isConnected()) {
       LogUtil.logInfo(LOG, _eventId,
           "MaintenanceRecoveryStage failed due to HelixManager being null or not connected!");
       return;
     }
+    // Check if this is a user-triggered maintenance mode with an end time
+    if (maintenanceSignal.getTriggeringEntity() == MaintenanceSignal.TriggeringEntity.USER) {
+      long endTime = maintenanceSignal.getEndTime();
+      // If endTime is set and the current time has passed the end time, exit maintenance mode
+      if (endTime > 0 && System.currentTimeMillis() >= endTime) {
+        String reason = String.format(
+            "Timeout-based exit from maintenance mode for cluster %s; End time %d has passed.",
+            event.getClusterName(), endTime);
+
+        manager.getClusterManagmentTool().manuallyEnableMaintenanceMode(manager.getClusterName(), false,
+            reason, null);
+        cache.setMaintenanceSignalChanged(); // Set the flag so we do not double enable/disable
+        LogUtil.logInfo(LOG, _eventId, reason);
+        return;
+      }
+      // Not yet time to exit, or no end time set
+      return;
+    }
 
     // At this point, the cluster entered maintenance mode automatically. Retrieve the
     // auto-triggering reason

helix-core/src/main/java/org/apache/helix/manager/zk/ZKHelixAdmin.java

@@ -1179,7 +1179,15 @@ public void manuallyEnableMaintenanceMode(String clusterName, boolean enabled, S
       Map<String, String> customFields) {
     processMaintenanceMode(clusterName, enabled, reason,
         MaintenanceSignal.AutoTriggerReason.NOT_APPLICABLE, customFields,
-        MaintenanceSignal.TriggeringEntity.USER);
+        MaintenanceSignal.TriggeringEntity.USER, -1);
+  }
+
+  @Override
+  public void manuallyEnableMaintenanceModeWithTimeout(String clusterName, boolean enabled,
+      String reason, long timeout, Map<String, String> customFields) {
+    processMaintenanceMode(clusterName, enabled, reason,
+        MaintenanceSignal.AutoTriggerReason.NOT_APPLICABLE, customFields,
+        MaintenanceSignal.TriggeringEntity.USER, timeout);
   }
 
   /**
@@ -1195,6 +1203,26 @@ private void processMaintenanceMode(String clusterName, final boolean enabled,
       final String reason, final MaintenanceSignal.AutoTriggerReason internalReason,
       final Map<String, String> customFields,
       final MaintenanceSignal.TriggeringEntity triggeringEntity) {
+    processMaintenanceMode(clusterName, enabled, reason, internalReason, customFields,
+        triggeringEntity, -1);
+  }
+
+  /**
+   * Helper method for enabling/disabling maintenance mode.
+   * @param clusterName
+   * @param enabled
+   * @param reason
+   * @param internalReason
+   * @param customFields
+   * @param triggeringEntity
+   * @param timeout time in milliseconds after which maintenance mode should be exited automatically.
+   *                Only applicable when enabled is true. Set to -1 for no automatic exit.
+   */
+  private void processMaintenanceMode(String clusterName, final boolean enabled,
+      final String reason, final MaintenanceSignal.AutoTriggerReason internalReason,
+      final Map<String, String> customFields,
+      final MaintenanceSignal.TriggeringEntity triggeringEntity,
+      final long timeout) {
     HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseDataAccessor);
     PropertyKey.Builder keyBuilder = accessor.keyBuilder();
     logger.info("Cluster {} {} {} maintenance mode for reason {}.", clusterName,
@@ -1212,6 +1240,15 @@ private void processMaintenanceMode(String clusterName, final boolean enabled,
       }
       maintenanceSignal.setTimestamp(currentTime);
       maintenanceSignal.setTriggeringEntity(triggeringEntity);
+
+      // Set end time if timeout is provided
+      if (timeout > 0) {
+        long endTime = currentTime + timeout;
+        maintenanceSignal.setEndTime(endTime);
+      } else {
+        maintenanceSignal.setEndTime(-1); // No automatic exit
+      }
+
       switch (triggeringEntity) {
         case CONTROLLER:
           // autoEnable
@@ -1234,6 +1271,23 @@ private void processMaintenanceMode(String clusterName, final boolean enabled,
       if (!accessor.createMaintenance(maintenanceSignal)) {
         throw new HelixException("Failed to create maintenance signal!");
       }
+
+      // If timeout is provided, create a /CONTROLLER/MAINTENANCE_TTL that is a PERSISTENT_WITH_TTL znode
+      if (timeout > 0) {
+        try {
+          String maintenanceTTLPath = "/" + clusterName + "/CONTROLLER/MAINTENANCE_TTL";
+          ZNRecord record = new ZNRecord("MAINTENANCE_TTL");
+
+          boolean success = ((ZkBaseDataAccessor<ZNRecord>) accessor.getBaseDataAccessor())
+              .create(maintenanceTTLPath, record, AccessOption.PERSISTENT_WITH_TTL, (int)(timeout / 1000));
+
+          if (!success) {
+            logger.warn("Failed to create TTL znode for maintenance mode. Auto exit may not work.");
+          }
+        } catch (Exception e) {
+          logger.warn("Failed to create TTL znode for maintenance mode. Auto exit may not work. ", e);
+        }
+      }
     }
 
     // Record a MaintenanceSignal history

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

@@ -32,7 +32,8 @@ public class MaintenanceSignal extends PauseSignal {
   public enum MaintenanceSignalProperty {
     TRIGGERED_BY,
     TIMESTAMP,
-    AUTO_TRIGGER_REASON
+    AUTO_TRIGGER_REASON,
+    END_TIME
   }
 
   /**
@@ -112,4 +113,21 @@ public void setTimestamp(long timestamp) {
   public long getTimestamp() {
     return _record.getLongField(MaintenanceSignalProperty.TIMESTAMP.name(), -1);
   }
+
+  /**
+   * Sets the end time for maintenance mode.
+   * @param endTime the time (in milliseconds) when maintenance mode should end. A value of -1 means 
+   *                no automatic exit.
+   */
+  public void setEndTime(long endTime) {
+    _record.setLongField(MaintenanceSignalProperty.END_TIME.name(), endTime);
+  }
+
+  /**
+   * Returns the end time for maintenance mode.
+   * @return the time (in milliseconds) when maintenance mode should end. Returns -1 if no end time is set.
+   */
+  public long getEndTime() {
+    return _record.getLongField(MaintenanceSignalProperty.END_TIME.name(), -1);
+  }
 }

helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ClusterAccessor.java

@@ -332,9 +332,27 @@ public Response updateCluster(@PathParam("clusterId") String clusterId,
         } catch (Exception e) {
           // NOP
         }
-        helixAdmin
-            .manuallyEnableMaintenanceMode(clusterId, command == Command.enableMaintenanceMode,
-                content, customFieldsMap);
+
+        // Check if a timeout is specified
+        long timeout = -1;
+        if (customFieldsMap != null) {
+          try {
+            String timeoutStr = customFieldsMap.get("timeout");
+            if (timeoutStr != null && !timeoutStr.isEmpty()) {
+              timeout = Long.parseLong(timeoutStr);
+            }
+          } catch (NumberFormatException nfe) {
+            LOG.warn("Invalid timeout value specified", nfe);
+          }
+        }
+
+        if (timeout > 0 && command == Command.enableMaintenanceMode) {
+          helixAdmin.manuallyEnableMaintenanceModeWithTimeout(clusterId, true, content, timeout,
+              customFieldsMap);
+        } else {
+          helixAdmin.manuallyEnableMaintenanceMode(clusterId, command == Command.enableMaintenanceMode,
+              content, customFieldsMap);
+        }
         break;
       case enableWagedRebalanceForAllResources:
         // Enable WAGED rebalance for all resources in the cluster