Subscription: preserve topic updates during owner transfer

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

@@ -59,6 +59,7 @@
 import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
 import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan;
 import org.apache.iotdb.confignode.i18n.ManagerMessages;
+import org.apache.iotdb.confignode.manager.subscription.SubscriptionCoordinator;
 import org.apache.iotdb.confignode.persistence.ProcedureInfo;
 import org.apache.iotdb.confignode.procedure.PartitionTableAutoCleaner;
 import org.apache.iotdb.confignode.procedure.Procedure;
@@ -1915,21 +1916,18 @@ public TSStatus createTopic(TCreateTopicReq req) {
   }
 
   public TSStatus alterTopic(TAlterTopicReq req) {
+    final SubscriptionCoordinator subscriptionCoordinator =
+        configManager.getSubscriptionManager().getSubscriptionCoordinator();
+    subscriptionCoordinator.lockTopicAlteration(req.getTopicName());
     boolean isOwnerLeaseRenewalBlocked = false;
     try {
       isOwnerLeaseRenewalBlocked =
-          configManager
-              .getSubscriptionManager()
-              .getSubscriptionCoordinator()
-              .blockOwnerLeaseRenewalIfOwnerTransfer(req);
+          subscriptionCoordinator.blockOwnerLeaseRenewalIfOwnerTransfer(req);
       // Owner transfers wait for the previous owner's lease to drain (lease duration + one
       // heartbeat interval, measured on the ConfigNode clock) inside the call below before the
       // updated meta is built; epoch fencing on DataNodes guarantees correctness in the meantime.
       final TopicMeta updatedTopicMeta =
-          configManager
-              .getSubscriptionManager()
-              .getSubscriptionCoordinator()
-              .buildAlteredTopicMetaAfterOwnerLeaseExpired(req);
+          subscriptionCoordinator.buildAlteredTopicMetaAfterOwnerLeaseExpired(req);
       if (updatedTopicMeta == null) {
         return new TSStatus(TSStatusCode.ALTER_TOPIC_ERROR.getStatusCode())
             .setMessage(
@@ -1938,9 +1936,27 @@ public TSStatus alterTopic(TAlterTopicReq req) {
                     req.getTopicName()));
       }
 
+      final Map<String, String> attributesBeforeInjection =
+          new HashMap<>(updatedTopicMeta.getConfig().getAttribute());
       injectTreeViewSourceAttributes(updatedTopicMeta.getConfig().getAttribute());
 
-      AlterTopicProcedure procedure = new AlterTopicProcedure(updatedTopicMeta);
+      final Map<String, String> updatedTopicAttributes = new HashMap<>();
+      if (Objects.nonNull(req.getTopicAttributes())) {
+        updatedTopicAttributes.putAll(req.getTopicAttributes());
+      }
+      updatedTopicMeta
+          .getConfig()
+          .getAttribute()
+          .forEach(
+              (key, value) -> {
+                if (!attributesBeforeInjection.containsKey(key)
+                    || !Objects.equals(attributesBeforeInjection.get(key), value)) {
+                  updatedTopicAttributes.put(key, value);
+                }
+              });
+
+      AlterTopicProcedure procedure =
+          new AlterTopicProcedure(updatedTopicMeta, updatedTopicAttributes);
       executor.submitProcedure(procedure);
       TSStatus status = waitingProcedureFinished(procedure);
       if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
@@ -1953,11 +1969,9 @@ public TSStatus alterTopic(TAlterTopicReq req) {
           .setMessage(e.getMessage());
     } finally {
       if (isOwnerLeaseRenewalBlocked) {
-        configManager
-            .getSubscriptionManager()
-            .getSubscriptionCoordinator()
-            .unblockOwnerLeaseRenewal(req.getTopicName());
+        subscriptionCoordinator.unblockOwnerLeaseRenewal(req.getTopicName());
       }
+      subscriptionCoordinator.unlockTopicAlteration(req.getTopicName());
     }
   }
 

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

@@ -65,6 +65,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 
 public class SubscriptionCoordinator {
 
@@ -80,11 +81,23 @@ public class SubscriptionCoordinator {
 
   private final SubscriptionMetaSyncer subscriptionMetaSyncer;
   private final SubscriptionOwnerLeaseSyncer subscriptionOwnerLeaseSyncer;
+
+  // Serialize client ALTER TOPIC requests per topic. Besides preventing ordinary partial updates
+  // from being built from the same snapshot, this also ensures that only one owner-transfer
+  // request can own the renewal-block entry for a topic at a time.
+  private final Map<String, TopicAlterationLock> topicAlterationLocks = new HashMap<>();
+
   // topicName -> blockSinceMs (ConfigNode local clock when owner-lease renewal was stopped for an
   // in-flight owner transfer). Used to skip renewal and to bound the admission wait.
   private final Map<String, Long> blockedOwnerLeaseRenewalTopics =
       Collections.synchronizedMap(new HashMap<>());
 
+  private static class TopicAlterationLock {
+
+    private final ReentrantLock lock = new ReentrantLock(true);
+    private int referenceCount;
+  }
+
   public SubscriptionCoordinator(ConfigManager configManager, SubscriptionInfo subscriptionInfo) {
     this.configManager = configManager;
     this.subscriptionInfo = subscriptionInfo;
@@ -134,6 +147,30 @@ public boolean isLocked() {
     return coordinatorLock.isLocked();
   }
 
+  public void lockTopicAlteration(final String topicName) {
+    final TopicAlterationLock topicAlterationLock;
+    synchronized (topicAlterationLocks) {
+      topicAlterationLock =
+          topicAlterationLocks.computeIfAbsent(topicName, ignored -> new TopicAlterationLock());
+      topicAlterationLock.referenceCount++;
+    }
+    topicAlterationLock.lock.lock();
+  }
+
+  public void unlockTopicAlteration(final String topicName) {
+    final TopicAlterationLock topicAlterationLock;
+    synchronized (topicAlterationLocks) {
+      topicAlterationLock = topicAlterationLocks.get(topicName);
+    }
+
+    topicAlterationLock.lock.unlock();
+    synchronized (topicAlterationLocks) {
+      if (--topicAlterationLock.referenceCount == 0) {
+        topicAlterationLocks.remove(topicName, topicAlterationLock);
+      }
+    }
+  }
+
   /////////////////////////////// Meta sync ///////////////////////////////
 
   public void startSubscriptionMetaSync() {
@@ -185,7 +222,10 @@ public TSStatus alterTopic(TAlterTopicReq req) {
 
   public boolean blockOwnerLeaseRenewalIfOwnerTransfer(TAlterTopicReq req) {
     final TopicMeta currentTopicMeta = subscriptionInfo.deepCopyTopicMeta(req.getTopicName());
-    final TopicMeta updatedTopicMeta = buildAlteredTopicMeta(req);
+    final TopicMeta updatedTopicMeta =
+        Objects.isNull(currentTopicMeta)
+            ? null
+            : currentTopicMeta.deepCopyWithUpdatedAttributes(req.getTopicAttributes());
     if (Objects.isNull(currentTopicMeta)
         || Objects.isNull(updatedTopicMeta)
         || Objects.equals(currentTopicMeta.getOwnerId(), updatedTopicMeta.getOwnerId())) {
@@ -211,7 +251,10 @@ public void unblockOwnerLeaseRenewal(String topicName) {
   public TopicMeta buildAlteredTopicMetaAfterOwnerLeaseExpired(TAlterTopicReq req)
       throws InterruptedException {
     final TopicMeta currentTopicMeta = subscriptionInfo.deepCopyTopicMeta(req.getTopicName());
-    final TopicMeta updatedTopicMeta = buildAlteredTopicMeta(req);
+    final TopicMeta updatedTopicMeta =
+        Objects.isNull(currentTopicMeta)
+            ? null
+            : currentTopicMeta.deepCopyWithUpdatedAttributes(req.getTopicAttributes());
     if (Objects.isNull(currentTopicMeta)
         || Objects.isNull(updatedTopicMeta)
         || Objects.equals(currentTopicMeta.getOwnerId(), updatedTopicMeta.getOwnerId())) {
@@ -220,20 +263,32 @@ public TopicMeta buildAlteredTopicMetaAfterOwnerLeaseExpired(TAlterTopicReq req)
     }
 
     final Long leaseDurationMs = currentTopicMeta.getOwnerLeaseDurationMs();
-    final Long blockSinceMs = blockedOwnerLeaseRenewalTopics.get(req.getTopicName());
-    if (Objects.isNull(leaseDurationMs) || Objects.isNull(blockSinceMs)) {
-      // No lease configured (no drain to wait for) or renewal not blocked: epoch fencing applies on
-      // reachable DataNodes; nothing further to wait on here.
+    if (Objects.isNull(leaseDurationMs)) {
+      // No lease configured, so there is nothing to drain.
       return updatedTopicMeta;
     }
 
+    waitForOwnerLeaseExpiration(req.getTopicName(), leaseDurationMs);
+
+    // Another alteration may have completed while this owner transfer was waiting for the old
+    // lease to drain. Rebuild from the latest TopicMeta so that this request only applies its own
+    // attributes instead of restoring the stale snapshot captured before the wait.
+    return buildAlteredTopicMeta(req);
+  }
+
+  void waitForOwnerLeaseExpiration(final String topicName, final long leaseDurationMs)
+      throws InterruptedException {
+    final Long blockSinceMs = blockedOwnerLeaseRenewalTopics.get(topicName);
+    if (Objects.isNull(blockSinceMs)) {
+      return;
+    }
+
     final long drainDeadlineMs =
         blockSinceMs + leaseDurationMs + SubscriptionOwnerLeaseSyncer.getHeartbeatIntervalMs();
     long remainingMs;
     while ((remainingMs = drainDeadlineMs - System.currentTimeMillis()) > 0) {
       Thread.sleep(Math.min(remainingMs, 1000L));
     }
-    return updatedTopicMeta;
   }
 
   public TopicMeta buildAlteredTopicMeta(TAlterTopicReq req) {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/subscription/topic/AlterTopicProcedure.java

@@ -41,7 +41,9 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -53,15 +55,35 @@ public class AlterTopicProcedure extends AbstractOperateSubscriptionProcedure {
 
   private TopicMeta existedTopicMeta;
 
+  // Non-null only for client ALTER TOPIC requests. These are merged again after the procedure has
+  // acquired the subscription lock, so a request cannot overwrite attributes committed after its
+  // initial TopicMeta snapshot was built.
+  private Map<String, String> updatedTopicAttributes;
+
   public AlterTopicProcedure() {
     super();
   }
 
+  public AlterTopicProcedure(final boolean shouldMergeUpdatedTopicAttributes) {
+    super();
+    if (shouldMergeUpdatedTopicAttributes) {
+      updatedTopicAttributes = new HashMap<>();
+    }
+  }
+
   public AlterTopicProcedure(TopicMeta updatedTopicMeta) {
     super();
     this.updatedTopicMeta = updatedTopicMeta;
   }
 
+  public AlterTopicProcedure(
+      TopicMeta updatedTopicMeta, Map<String, String> updatedTopicAttributes) {
+    super();
+    this.updatedTopicMeta = updatedTopicMeta;
+    this.updatedTopicAttributes =
+        Objects.isNull(updatedTopicAttributes) ? null : new HashMap<>(updatedTopicAttributes);
+  }
+
   /** This is only used when the subscription info lock is held by another procedure. */
   public AlterTopicProcedure(
       TopicMeta updatedTopicMeta, AtomicReference<SubscriptionInfo> subscriptionInfo) {
@@ -70,6 +92,15 @@ public AlterTopicProcedure(
     this.subscriptionInfo = subscriptionInfo;
   }
 
+  /** This is only used when the subscription info lock is held by another procedure. */
+  public AlterTopicProcedure(
+      TopicMeta updatedTopicMeta,
+      Map<String, String> updatedTopicAttributes,
+      AtomicReference<SubscriptionInfo> subscriptionInfo) {
+    this(updatedTopicMeta, updatedTopicAttributes);
+    this.subscriptionInfo = subscriptionInfo;
+  }
+
   /** This should be called after {@link #executeFromValidate}. */
   public TopicMeta getExistedTopicMeta() {
     return existedTopicMeta;
@@ -80,6 +111,10 @@ public TopicMeta getUpdatedTopicMeta() {
     return updatedTopicMeta;
   }
 
+  public boolean shouldMergeUpdatedTopicAttributes() {
+    return Objects.nonNull(updatedTopicAttributes);
+  }
+
   @Override
   protected SubscriptionOperation getOperation() {
     return SubscriptionOperation.ALTER_TOPIC;
@@ -89,9 +124,12 @@ protected SubscriptionOperation getOperation() {
   public boolean executeFromValidate(ConfigNodeProcedureEnv env) throws SubscriptionException {
     LOGGER.info(ProcedureMessages.ALTERTOPICPROCEDURE_EXECUTEFROMVALIDATE);
 
-    subscriptionInfo.get().validateBeforeAlteringTopic(updatedTopicMeta);
+    existedTopicMeta = subscriptionInfo.get().deepCopyTopicMeta(updatedTopicMeta.getTopicName());
+    if (Objects.nonNull(updatedTopicAttributes) && Objects.nonNull(existedTopicMeta)) {
+      updatedTopicMeta = existedTopicMeta.deepCopyWithUpdatedAttributes(updatedTopicAttributes);
+    }
 
-    existedTopicMeta = subscriptionInfo.get().getTopicMeta(updatedTopicMeta.getTopicName());
+    subscriptionInfo.get().validateBeforeAlteringTopic(updatedTopicMeta);
 
     return true;
   }
@@ -196,7 +234,11 @@ public void rollbackFromOperateOnDataNodes(ConfigNodeProcedureEnv env)
 
   @Override
   public void serialize(DataOutputStream stream) throws IOException {
-    stream.writeShort(ProcedureType.ALTER_TOPIC_PROCEDURE.getTypeCode());
+    stream.writeShort(
+        (shouldMergeUpdatedTopicAttributes()
+                ? ProcedureType.ALTER_TOPIC_WITH_ATTRIBUTES_PROCEDURE
+                : ProcedureType.ALTER_TOPIC_PROCEDURE)
+            .getTypeCode());
     super.serialize(stream);
 
     ReadWriteIOUtils.write(updatedTopicMeta != null, stream);
@@ -208,6 +250,10 @@ public void serialize(DataOutputStream stream) throws IOException {
     if (existedTopicMeta != null) {
       existedTopicMeta.serialize(stream);
     }
+
+    if (shouldMergeUpdatedTopicAttributes()) {
+      ReadWriteIOUtils.write(updatedTopicAttributes, stream);
+    }
   }
 
   @Override
@@ -221,6 +267,10 @@ public void deserialize(ByteBuffer byteBuffer) {
     if (ReadWriteIOUtils.readBool(byteBuffer)) {
       existedTopicMeta = TopicMeta.deserialize(byteBuffer);
     }
+
+    if (shouldMergeUpdatedTopicAttributes()) {
+      updatedTopicAttributes = ReadWriteIOUtils.readMap(byteBuffer);
+    }
   }
 
   @Override
@@ -236,12 +286,18 @@ public boolean equals(Object o) {
         && Objects.equals(getCurrentState(), that.getCurrentState())
         && getCycles() == that.getCycles()
         && Objects.equals(updatedTopicMeta, that.updatedTopicMeta)
-        && Objects.equals(existedTopicMeta, that.existedTopicMeta);
+        && Objects.equals(existedTopicMeta, that.existedTopicMeta)
+        && Objects.equals(updatedTopicAttributes, that.updatedTopicAttributes);
   }
 
   @Override
   public int hashCode() {
     return Objects.hash(
-        getProcId(), getCurrentState(), getCycles(), updatedTopicMeta, existedTopicMeta);
+        getProcId(),
+        getCurrentState(),
+        getCycles(),
+        updatedTopicMeta,
+        existedTopicMeta,
+        updatedTopicAttributes);
   }
 }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/store/ProcedureFactory.java

@@ -382,6 +382,9 @@ public Procedure create(ByteBuffer buffer) throws IOException {
       case ALTER_TOPIC_PROCEDURE:
         procedure = new AlterTopicProcedure();
         break;
+      case ALTER_TOPIC_WITH_ATTRIBUTES_PROCEDURE:
+        procedure = new AlterTopicProcedure(true);
+        break;
       case TOPIC_META_SYNC_PROCEDURE:
         procedure = new TopicMetaSyncProcedure();
         break;
@@ -546,7 +549,9 @@ public static ProcedureType getProcedureType(final Procedure<?> procedure) {
     } else if (procedure instanceof DropTopicProcedure) {
       return ProcedureType.DROP_TOPIC_PROCEDURE;
     } else if (procedure instanceof AlterTopicProcedure) {
-      return ProcedureType.ALTER_TOPIC_PROCEDURE;
+      return ((AlterTopicProcedure) procedure).shouldMergeUpdatedTopicAttributes()
+          ? ProcedureType.ALTER_TOPIC_WITH_ATTRIBUTES_PROCEDURE
+          : ProcedureType.ALTER_TOPIC_PROCEDURE;
     } else if (procedure instanceof TopicMetaSyncProcedure) {
       return ProcedureType.TOPIC_META_SYNC_PROCEDURE;
     } else if (procedure instanceof CreateSubscriptionProcedure) {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/store/ProcedureType.java

@@ -170,6 +170,7 @@ public enum ProcedureType {
   CONSUMER_GROUP_META_SYNC_PROCEDURE((short) 1509),
   COMMIT_PROGRESS_SYNC_PROCEDURE((short) 1510),
   SUBSCRIPTION_HANDLE_LEADER_CHANGE_PROCEDURE((short) 1511),
+  ALTER_TOPIC_WITH_ATTRIBUTES_PROCEDURE((short) 1512),
 
   /** Other */
   @TestOnly