Fixed TTL problems

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/TTLManager.java

@@ -127,6 +127,10 @@
     return ((ShowTTLResp) showTTL(new ShowTTLPlan())).getPathTTLMap();
   }
 
+  public long getTTL(final String[] pathPattern) {
+    return ttlInfo.getTTL(pathPattern);
+  }
+
   public int getTTLCount() {
     return ttlInfo.getTTLCount();
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/TTLInfo.java

@@ -70,13 +70,15 @@
     try {
       // check ttl rule capacity
       final int tTlRuleCapacity = CommonDescriptor.getInstance().getConfig().getTTlRuleCapacity();
-      if (getTTLCount() >= tTlRuleCapacity) {
+      final int newTTLRuleCount = calculateNewTTLRuleCount(plan);
+      final int requestedTTLRuleCount = ttlCache.getTtlCount() + newTTLRuleCount;
+      if (newTTLRuleCount > 0 && requestedTTLRuleCount > tTlRuleCapacity) {
         TSStatus errorStatus = new TSStatus(TSStatusCode.OVERSIZE_TTL.getStatusCode());
         errorStatus.setMessage(
             String.format(
-                "The number of TTL rules has reached the limit (%d). Please delete "
-                    + "some existing rules first.",
-                tTlRuleCapacity));
+                "The number of TTL rules has reached the limit "
+                    + "(capacity: %d, requested total: %d). Please delete some existing rules first.",
+                tTlRuleCapacity, requestedTTLRuleCount));
         return errorStatus;
       }
       ttlCache.setTTL(plan.getPathPattern(), plan.getTTL());
@@ -92,8 +94,22 @@
     return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
   }
 
+  private int calculateNewTTLRuleCount(SetTTLPlan plan) {
+    int newTTLRuleCount = isNewTTLRule(plan.getPathPattern()) ? 1 : 0;
+    if (plan.isDataBase()) {
+      String[] pathNodes = Arrays.copyOf(plan.getPathPattern(), plan.getPathPattern().length + 1);
+      pathNodes[pathNodes.length - 1] = IoTDBConstant.MULTI_LEVEL_PATH_WILDCARD;
+      newTTLRuleCount += isNewTTLRule(pathNodes) ? 1 : 0;
+    }
+    return newTTLRuleCount;
+  }
+
+  private boolean isNewTTLRule(String[] pathNodes) {
+    return ttlCache.getLastNodeTTL(pathNodes) == TTLCache.NULL_TTL;
+  }
+
   /** Only used for upgrading from database level ttl to device level ttl. */
   public void setTTL(Map<String, Long> databaseTTLMap) throws IllegalPathException {
     lock.writeLock().lock();
     try {
       for (Map.Entry<String, Long> entry : databaseTTLMap.entrySet()) {
@@ -159,6 +175,15 @@
     }
   }
 
+  public long getTTL(final String[] pathPattern) {
+    lock.readLock().lock();
+    try {
+      return ttlCache.getLastNodeTTL(pathPattern);
+    } finally {
+      lock.readLock().unlock();
+    }
+  }
+
   /**
    * Get the maximum ttl of the corresponding database level.
    *

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/schema/SetTTLProcedure.java

@@ -22,8 +22,10 @@
 import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.common.rpc.thrift.TSetTTLReq;
+import org.apache.iotdb.commons.conf.IoTDBConstant;
 import org.apache.iotdb.commons.exception.IoTDBException;
 import org.apache.iotdb.commons.exception.MetadataException;
+import org.apache.iotdb.commons.schema.ttl.TTLCache;
 import org.apache.iotdb.confignode.client.async.CnToDnAsyncRequestType;
 import org.apache.iotdb.confignode.client.async.CnToDnInternalServiceAsyncRequestManager;
 import org.apache.iotdb.confignode.client.async.handlers.DataNodeAsyncRequestContext;
@@ -47,14 +49,21 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Objects;
 
 public class SetTTLProcedure extends StateMachineProcedure<ConfigNodeProcedureEnv, SetTTLState> {
   private static final Logger LOGGER = LoggerFactory.getLogger(SetTTLProcedure.class);
+  // Distinguishes no previous TTL from TTLCache.NULL_TTL, the explicit unset marker for rollback.
+  private static final long TTL_NOT_EXIST = Long.MIN_VALUE;
+  private static final int ROLLBACK_STATE_BYTES = Byte.BYTES + Long.BYTES * 2;
 
   private SetTTLPlan plan;
+  private long previousTTL = TTL_NOT_EXIST;
+  private long previousDatabaseWildcardTTL = TTL_NOT_EXIST;
+  private boolean previousTTLStateCaptured = false;
 
   public SetTTLProcedure(final boolean isGeneratedByPipe) {
     super(isGeneratedByPipe);
@@ -71,6 +80,10 @@
     long startTime = System.currentTimeMillis();
     try {
       switch (state) {
+        case CAPTURE_PREVIOUS_TTL:
+          capturePreviousTTLState(env);
+          setNextState(SetTTLState.SET_CONFIGNODE_TTL);
+          return Flow.HAS_MORE_STATE;
         case SET_CONFIGNODE_TTL:
           setConfigNodeTTL(env);
           return Flow.HAS_MORE_STATE;
@@ -86,18 +99,13 @@
     }
   }
 
-  private void setConfigNodeTTL(ConfigNodeProcedureEnv env) {
-    TSStatus res;
-    try {
-      res =
-          env.getConfigManager()
-              .getConsensusManager()
-              .write(isGeneratedByPipe ? new PipeEnrichedPlan(this.plan) : this.plan);
-    } catch (ConsensusException e) {
-      LOGGER.warn(ConfigNodeMessages.FAILED_IN_THE_WRITE_API_EXECUTING_THE_CONSENSUS_LAYER_DUE, e);
-      res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode());
-      res.setMessage(e.getMessage());
+  void setConfigNodeTTL(final ConfigNodeProcedureEnv env) {
+    if (!previousTTLStateCaptured) {
+      capturePreviousTTLState(env);
+      setNextState(SetTTLState.SET_CONFIGNODE_TTL);
+      return;
     }
+    final TSStatus res = writeConfigNodePlan(env, plan);
     if (res.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
       LOGGER.info(ProcedureMessages.FAILED_TO_EXECUTE_PLAN_BECAUSE, plan, res.message);
       setFailure(new ProcedureException(new IoTDBException(res)));
@@ -106,35 +114,177 @@
     }
   }
 
-  private void updateDataNodeTTL(ConfigNodeProcedureEnv env) {
-    Map<Integer, TDataNodeLocation> dataNodeLocationMap =
+  void updateDataNodeTTL(final ConfigNodeProcedureEnv env) {
+    final Map<Integer, TDataNodeLocation> dataNodeLocationMap =
         env.getConfigManager().getNodeManager().getRegisteredDataNodeLocations();
-    DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> clientHandler =
-        new DataNodeAsyncRequestContext<>(
-            CnToDnAsyncRequestType.SET_TTL,
-            new TSetTTLReq(
-                Collections.singletonList(String.join(".", plan.getPathPattern())),
-                plan.getTTL(),
-                plan.isDataBase()),
-            dataNodeLocationMap);
+    final DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> clientHandler =
+        sendTTLRequest(
+            dataNodeLocationMap,
+            buildSetTTLReq(plan.getPathPattern(), plan.getTTL(), plan.isDataBase()));
+    if (hasFailedDataNode(clientHandler)) {
+      LOGGER.error(ProcedureMessages.FAILED_TO_UPDATE_TTL_CACHE_OF_DATANODE);
+      setFailure(
+          new ProcedureException(
+              new MetadataException(ProcedureMessages.UPDATE_DATANODE_TTL_CACHE_FAILED)));
+    }
+  }
+
+  private void capturePreviousTTLState(final ConfigNodeProcedureEnv env) {
+    if (previousTTLStateCaptured) {
+      return;
+    }
+    previousTTL = getTTLOrDefault(env, plan.getPathPattern());
+    if (plan.isDataBase()) {
+      previousDatabaseWildcardTTL =
+          getTTLOrDefault(env, getDatabaseWildcardPathPattern(plan.getPathPattern()));
+    }
+    previousTTLStateCaptured = true;
+  }
+
+  TSStatus writeConfigNodePlan(final ConfigNodeProcedureEnv env, final SetTTLPlan setTTLPlan) {
+    try {
+      return env.getConfigManager()
+          .getConsensusManager()
+          .write(isGeneratedByPipe ? new PipeEnrichedPlan(setTTLPlan) : setTTLPlan);
+    } catch (ConsensusException e) {
+      LOGGER.warn(ConfigNodeMessages.FAILED_IN_THE_WRITE_API_EXECUTING_THE_CONSENSUS_LAYER_DUE, e);
+      final TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode());
+      res.setMessage(e.getMessage());
+      return res;
+    }
+  }
+
+  DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> sendTTLRequest(
+      final Map<Integer, TDataNodeLocation> dataNodeLocationMap, final TSetTTLReq req) {
+    final DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> clientHandler =
+        new DataNodeAsyncRequestContext<>(CnToDnAsyncRequestType.SET_TTL, req, dataNodeLocationMap);
     CnToDnInternalServiceAsyncRequestManager.getInstance().sendAsyncRequestWithRetry(clientHandler);
-    Map<Integer, TSStatus> statusMap = clientHandler.getResponseMap();
-    for (TSStatus status : statusMap.values()) {
-      // all dataNodes must clear the related schemaengine cache
+    return clientHandler;
+  }
+
+  private TSetTTLReq buildSetTTLReq(
+      final String[] pathPattern, final long ttl, final boolean isDataBase) {
+    return new TSetTTLReq(
+        Collections.singletonList(String.join(".", pathPattern)), ttl, isDataBase);
+  }
+
+  private boolean hasFailedDataNode(
+      final DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> clientHandler) {
+    if (!clientHandler.getRequestIndices().isEmpty()) {
+      return true;
+    }
+    for (TSStatus status : clientHandler.getResponseMap().values()) {
       if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-        LOGGER.error(ProcedureMessages.FAILED_TO_UPDATE_TTL_CACHE_OF_DATANODE);
-        setFailure(
-            new ProcedureException(
-                new MetadataException(ProcedureMessages.UPDATE_DATANODE_TTL_CACHE_FAILED)));
-        return;
+        return true;
       }
     }
+    return false;
   }
 
+  private long getTTLOrDefault(final ConfigNodeProcedureEnv env, final String[] pathPattern) {
+    final long ttl = env.getConfigManager().getTTLManager().getTTL(pathPattern);
+    return ttl == TTLCache.NULL_TTL ? TTL_NOT_EXIST : ttl;
+  }
+
+  private String[] getDatabaseWildcardPathPattern(final String[] pathPattern) {
+    final String[] pathNodes = Arrays.copyOf(pathPattern, pathPattern.length + 1);
+    pathNodes[pathNodes.length - 1] = IoTDBConstant.MULTI_LEVEL_PATH_WILDCARD;
+    return pathNodes;
+  }
+
+  private void rollbackConfigNodeTTL(final ConfigNodeProcedureEnv env) throws ProcedureException {
+    restoreTTLOnConfigNode(env, plan.getPathPattern(), previousTTL);
+    if (plan.isDataBase()) {
+      restoreTTLOnConfigNode(
+          env, getDatabaseWildcardPathPattern(plan.getPathPattern()), previousDatabaseWildcardTTL);
+    }
+  }
+
+  private void restoreTTLOnConfigNode(
+      final ConfigNodeProcedureEnv env, final String[] pathPattern, final long ttl)
+      throws ProcedureException {
+    // TTL_NOT_EXIST means the original ttl was absent; NULL_TTL asks the executor to unset it.
+    final SetTTLPlan rollbackPlan =
+        new SetTTLPlan(pathPattern, ttl == TTL_NOT_EXIST ? TTLCache.NULL_TTL : ttl);
+    // Database rollback restores the database path and db.** separately, so avoid auto-expansion.
+    rollbackPlan.setDataBase(false);
+    final TSStatus status = writeConfigNodePlan(env, rollbackPlan);
+    if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+      throw new ProcedureException(
+          new MetadataException(
+              "Rollback ConfigNode ttl failed for "
+                  + String.join(".", pathPattern)
+                  + ": "
+                  + status.getMessage()));
+    }
+  }
+
+  private void rollbackDataNodeTTL(final ConfigNodeProcedureEnv env) throws ProcedureException {
+    final Map<Integer, TDataNodeLocation> dataNodeLocationMap =
+        env.getConfigManager().getNodeManager().getRegisteredDataNodeLocations();
+    restoreTTLOnDataNodes(dataNodeLocationMap, plan.getPathPattern(), previousTTL);
+    if (plan.isDataBase()) {
+      restoreTTLOnDataNodes(
+          dataNodeLocationMap,
+          getDatabaseWildcardPathPattern(plan.getPathPattern()),
+          previousDatabaseWildcardTTL);
+    }
+  }
+
+  private void restoreTTLOnDataNodes(
+      final Map<Integer, TDataNodeLocation> dataNodeLocationMap,
+      final String[] pathPattern,
+      final long ttl)
+      throws ProcedureException {
+    if (dataNodeLocationMap.isEmpty()) {
+      return;
+    }
+    final DataNodeAsyncRequestContext<TSetTTLReq, TSStatus> clientHandler =
+        sendTTLRequest(
+            dataNodeLocationMap,
+            buildSetTTLReq(pathPattern, ttl == TTL_NOT_EXIST ? TTLCache.NULL_TTL : ttl, false));
+    if (hasFailedDataNode(clientHandler)) {
+      throw new ProcedureException(
+          new MetadataException(
+              "Rollback dataNode ttl cache failed for " + String.join(".", pathPattern)));
+    }
+  }
+
+  /**
+   * Best-effort rollback: restore both sides, throw the earliest failure, and suppress later ones.
+   */
   @Override
-  protected void rollbackState(
-      ConfigNodeProcedureEnv configNodeProcedureEnv, SetTTLState setTTLState)
-      throws IOException, InterruptedException, ProcedureException {}
+  protected void rollbackState(final ConfigNodeProcedureEnv env, final SetTTLState setTTLState)
+      throws IOException, InterruptedException, ProcedureException {
+    if (setTTLState != SetTTLState.UPDATE_DATANODE_CACHE || !previousTTLStateCaptured) {
+      return;
+    }
+    ProcedureException rollbackFailure = null;
+    try {
+      rollbackConfigNodeTTL(env);
+    } catch (ProcedureException e) {
+      LOGGER.error("Failed to rollback ConfigNode ttl state.", e);
+      rollbackFailure = e;
+    }
+    try {
+      rollbackDataNodeTTL(env);
+    } catch (ProcedureException e) {
+      LOGGER.error("Failed to rollback DataNode ttl cache.", e);
+      if (rollbackFailure == null) {
+        rollbackFailure = e;
+      } else {
+        rollbackFailure.addSuppressed(e);
+      }
+    }
+    if (rollbackFailure != null) {
+      throw rollbackFailure;
+    }
+  }
+
+  @Override
+  protected boolean isRollbackSupported(final SetTTLState state) {
+    return state == SetTTLState.UPDATE_DATANODE_CACHE;
+  }
 
   @Override
   protected SetTTLState getState(int stateId) {
@@ -148,7 +298,7 @@
 
   @Override
   protected SetTTLState getInitialState() {
-    return SetTTLState.SET_CONFIGNODE_TTL;
+    return SetTTLState.CAPTURE_PREVIOUS_TTL;
   }
 
   @Override
@@ -159,14 +309,25 @@
             : ProcedureType.SET_TTL_PROCEDURE.getTypeCode());
     super.serialize(stream);
     ReadWriteIOUtils.write(plan.serializeToByteBuffer(), stream);
+    stream.writeBoolean(previousTTLStateCaptured);
+    stream.writeLong(previousTTL);
+    stream.writeLong(previousDatabaseWildcardTTL);
   }
 
   @Override
   public void deserialize(ByteBuffer byteBuffer) {
     super.deserialize(byteBuffer);
     try {
-      ReadWriteIOUtils.readInt(byteBuffer);
+      final int length = ReadWriteIOUtils.readInt(byteBuffer);
+      final int position = byteBuffer.position();
       this.plan = (SetTTLPlan) ConfigPhysicalPlan.Factory.create(byteBuffer);
+      // The serialized plan buffer may include padding; skip to the actual payload end.
+      byteBuffer.position(position + length);
+      if (byteBuffer.remaining() >= ROLLBACK_STATE_BYTES) {
+        this.previousTTLStateCaptured = byteBuffer.get() != 0;
+        this.previousTTL = byteBuffer.getLong();
+        this.previousDatabaseWildcardTTL = byteBuffer.getLong();
+      }
     } catch (IOException e) {
       LOGGER.error(ProcedureMessages.IO_ERROR_WHEN_DESERIALIZE_SETTTL_PLAN, e);
     }
@@ -180,12 +341,21 @@
     if (o == null || getClass() != o.getClass()) {
       return false;
     }
-    return this.plan.equals(((SetTTLProcedure) o).plan)
-        && this.isGeneratedByPipe == (((SetTTLProcedure) o).isGeneratedByPipe);
+    final SetTTLProcedure that = (SetTTLProcedure) o;
+    return this.isGeneratedByPipe == that.isGeneratedByPipe
+        && this.previousTTL == that.previousTTL
+        && this.previousDatabaseWildcardTTL == that.previousDatabaseWildcardTTL
+        && this.previousTTLStateCaptured == that.previousTTLStateCaptured
+        && this.plan.equals(that.plan);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(plan, isGeneratedByPipe);
+    return Objects.hash(
+        plan,
+        isGeneratedByPipe,
+        previousTTL,
+        previousDatabaseWildcardTTL,
+        previousTTLStateCaptured);
   }
 }