[Subscription] Restore consensus subscription brokers after restart

integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/IoTDBConsensusSubscriptionRecoveryTableIT.java

@@ -0,0 +1,165 @@
+/*
+ * 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.
+ */
+
+package org.apache.iotdb.subscription.it.consensus.local.tablemodel;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.env.cluster.config.MppCommonConfig;
+import org.apache.iotdb.it.env.cluster.env.AbstractEnv;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer;
+import org.apache.iotdb.session.subscription.payload.SubscriptionMessage;
+import org.apache.iotdb.session.subscription.payload.SubscriptionRecordHandler;
+import org.apache.iotdb.subscription.it.consensus.local.AbstractSubscriptionConsensusLocalIT;
+
+import org.apache.tsfile.read.query.dataset.ResultSet;
+import org.awaitility.Awaitility;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.time.Duration;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({LocalStandaloneIT.class})
+public class IoTDBConsensusSubscriptionRecoveryTableIT
+    extends AbstractSubscriptionConsensusLocalIT {
+
+  @Override
+  @Before
+  public void setUp() throws Exception {
+    // Isolate subscription recovery from data loss caused by ASYNC WAL during a forced shutdown.
+    EnvFactory.getEnv().getConfig().getCommonConfig().setWalMode("SYNC");
+    // Persist each acknowledged batch before killing the servers.
+    final Properties properties = new Properties();
+    properties.setProperty("subscription_consensus_commit_persist_interval", "1");
+    properties.setProperty("subscription_consensus_commit_fsync_enabled", "true");
+    ((MppCommonConfig) EnvFactory.getEnv().getConfig().getDataNodeCommonConfig())
+        .updateProperties(properties);
+    super.setUp();
+  }
+
+  @Test
+  public void testRebuiltConsumersResumeAfterFullClusterRestart() throws Exception {
+    final ConsensusSubscriptionTableITSupport.TestIdentifiers ids =
+        ConsensusSubscriptionTableITSupport.newIdentifiers("table_full_cluster_restart");
+    final String tableName = "t1";
+    final Set<String> expectedRows = new LinkedHashSet<>();
+    final Set<String> consumedRows = new LinkedHashSet<>();
+    SubscriptionTablePullConsumer firstConsumer = null;
+    SubscriptionTablePullConsumer secondConsumer = null;
+
+    try {
+      ConsensusSubscriptionTableITSupport.bootstrapDatabaseAndTable(
+          ids.getDatabase(), tableName, ConsensusSubscriptionTableITSupport.DEFAULT_TABLE_SCHEMA);
+      ConsensusSubscriptionTableITSupport.createConsensusTopic(ids.getTopic(), ".*", ".*");
+      firstConsumer =
+          ConsensusSubscriptionTableITSupport.createConsumer(
+              ids.consumer("first"), ids.getConsumerGroupId());
+      secondConsumer =
+          ConsensusSubscriptionTableITSupport.createConsumer(
+              ids.consumer("second"), ids.getConsumerGroupId());
+      firstConsumer.subscribe(ids.getTopic());
+      secondConsumer.subscribe(ids.getTopic());
+
+      final Set<String> committedRows =
+          ConsensusSubscriptionTableITSupport.insertRows(
+              ids.getDatabase(), tableName, 1_000L, 16, true);
+      expectedRows.addAll(committedRows);
+      assertConsumedRows(expectedRows, consumedRows, firstConsumer, secondConsumer);
+
+      final Set<String> pendingRows =
+          ConsensusSubscriptionTableITSupport.insertRows(
+              ids.getDatabase(), tableName, 2_000L, 64, true);
+
+      EnvFactory.getEnv().shutdownForciblyAllDataNodes();
+      EnvFactory.getEnv().shutdownForciblyAllConfigNodes();
+      // Close while the servers are down so the old clients cannot remove persisted subscriptions.
+      firstConsumer.close();
+      firstConsumer = null;
+      secondConsumer.close();
+      secondConsumer = null;
+
+      EnvFactory.getEnv().startAllConfigNodes();
+      EnvFactory.getEnv().startAllDataNodes();
+      ((AbstractEnv) EnvFactory.getEnv()).checkClusterStatusWithoutUnknown();
+
+      firstConsumer =
+          ConsensusSubscriptionTableITSupport.createConsumer(
+              ids.consumer("first"), ids.getConsumerGroupId());
+      secondConsumer =
+          ConsensusSubscriptionTableITSupport.createConsumer(
+              ids.consumer("second"), ids.getConsumerGroupId());
+      firstConsumer.subscribe(ids.getTopic());
+      secondConsumer.subscribe(ids.getTopic());
+
+      // Do not seek, recreate the topic, or insert new rows to trigger subscription setup.
+      // Track unique rows as a receiver would, allowing redelivery around the forced shutdown.
+      expectedRows.addAll(pendingRows);
+      assertConsumedRows(expectedRows, consumedRows, firstConsumer, secondConsumer);
+
+      final Set<String> newRows =
+          ConsensusSubscriptionTableITSupport.insertRows(
+              ids.getDatabase(), tableName, 3_000L, 16, true);
+      expectedRows.addAll(newRows);
+      assertConsumedRows(expectedRows, consumedRows, firstConsumer, secondConsumer);
+    } finally {
+      ConsensusSubscriptionTableITSupport.cleanup(firstConsumer, ids.getTopic(), ids.getDatabase());
+      ConsensusSubscriptionTableITSupport.cleanup(
+          secondConsumer, ids.getTopic(), ids.getDatabase());
+    }
+  }
+
+  private static void assertConsumedRows(
+      final Set<String> expectedRows,
+      final Set<String> consumedRows,
+      final SubscriptionTablePullConsumer... consumers) {
+    Awaitility.await()
+        .pollInSameThread()
+        .atMost(Duration.ofMinutes(2))
+        .untilAsserted(
+            () -> {
+              for (final SubscriptionTablePullConsumer consumer : consumers) {
+                final List<SubscriptionMessage> messages = consumer.poll(Duration.ofSeconds(1));
+                for (final SubscriptionMessage message : messages) {
+                  for (final ResultSet resultSet : message.getResultSets()) {
+                    final SubscriptionRecordHandler.SubscriptionResultSet records =
+                        (SubscriptionRecordHandler.SubscriptionResultSet) resultSet;
+                    while (records.hasNext()) {
+                      consumedRows.add(
+                          ConsensusSubscriptionTableITSupport.rowKey(
+                              records.getDatabaseName(),
+                              records.getTableName(),
+                              records.nextRecord().getTimestamp()));
+                    }
+                  }
+                }
+                consumer.commitSync(messages);
+              }
+              Assert.assertEquals(expectedRows, consumedRows);
+            });
+  }
+}

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensus.java

@@ -111,8 +111,9 @@ public class IoTConsensus implements IConsensus {
   private IoTConsensusConfig config;
 
   /**
-   * Optional callback invoked after a new local peer is created via {@link #createLocalPeer}. Used
-   * by the subscription system to auto-bind prefetching queues to new DataRegions.
+   * Optional callback invoked after a local peer is started, either during recovery or via {@link
+   * #createLocalPeer}. Used by the subscription system to auto-bind prefetching queues to available
+   * DataRegions.
    */
   public static volatile BiConsumer<ConsensusGroupId, IoTConsensusServerImpl> onNewPeerCreated;
 
@@ -257,6 +258,10 @@ private void initAndRecover() throws IOException {
             callback.accept(consensusGroupId, consensus);
           }
           consensus.start();
+          // Subscription metadata may have been loaded before any local consensus peers existed.
+          // Recovered peers must trigger the same binding as newly created peers, after startup
+          // has restored their indices and installed the pre-start WAL-retention guards.
+          notifyPeerCreated(consensusGroupId, consensus);
         });
   }
 
@@ -362,11 +367,15 @@ public void createLocalPeer(ConsensusGroupId groupId, List<Peer> peers)
       throw new ConsensusGroupAlreadyExistException(groupId);
     }
 
-    // Notify subscription system about new peer creation for auto-binding
+    notifyPeerCreated(groupId, stateMachineMap.get(groupId));
+  }
+
+  private void notifyPeerCreated(
+      final ConsensusGroupId groupId, final IoTConsensusServerImpl serverImpl) {
     final BiConsumer<ConsensusGroupId, IoTConsensusServerImpl> callback = onNewPeerCreated;
     if (callback != null) {
       try {
-        callback.accept(groupId, stateMachineMap.get(groupId));
+        callback.accept(groupId, serverImpl);
       } catch (final Exception e) {
         logger.warn(
             IoTConsensusMessages.LOG_ONNEWPEERCREATED_CALLBACK_FAILED_GROUP_ARG_2671FCDA,

iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/IoTConsensusRecoveryTest.java

@@ -0,0 +1,181 @@
+/*
+ * 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.
+ */
+
+package org.apache.iotdb.consensus.iot;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.consensus.ConsensusGroupId;
+import org.apache.iotdb.commons.consensus.DataRegionId;
+import org.apache.iotdb.consensus.common.Peer;
+import org.apache.iotdb.consensus.config.ConsensusConfig;
+import org.apache.iotdb.consensus.iot.util.TestStateMachine;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiConsumer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class IoTConsensusRecoveryTest {
+
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+  private final Map<ConsensusGroupId, LifecycleStateMachine> stateMachines = new HashMap<>();
+  private final Map<ConsensusGroupId, List<Peer>> peerLists = new HashMap<>();
+  private BiConsumer<ConsensusGroupId, IoTConsensusServerImpl> originalNewPeerCallback;
+  private BiConsumer<ConsensusGroupId, IoTConsensusServerImpl> originalRecoveringCallback;
+  private ConsensusConfig config;
+  private IoTConsensus consensus;
+
+  @Before
+  public void setUp() throws Exception {
+    originalNewPeerCallback = IoTConsensus.onNewPeerCreated;
+    originalRecoveringCallback = IoTConsensus.onPeerRecovering;
+    IoTConsensus.onNewPeerCreated = null;
+    IoTConsensus.onPeerRecovering = null;
+    config =
+        ConsensusConfig.newBuilder()
+            .setThisNodeId(1)
+            .setThisNode(new TEndPoint("127.0.0.1", 0))
+            .setStorageDir(temporaryFolder.newFolder("consensus").getAbsolutePath())
+            .setRecvSnapshotDirs(
+                Collections.singletonList(temporaryFolder.newFolder("snapshots").getAbsolutePath()))
+            .setConsensusGroupType(TConsensusGroupType.DataRegion)
+            .build();
+    consensus = newConsensus();
+    consensus.start();
+  }
+
+  @After
+  public void tearDown() {
+    try {
+      if (consensus != null) {
+        consensus.stop();
+      }
+    } finally {
+      IoTConsensus.onNewPeerCreated = originalNewPeerCallback;
+      IoTConsensus.onPeerRecovering = originalRecoveringCallback;
+    }
+  }
+
+  @Test
+  public void testRecoveredPeersNotifyAfterStarting() throws Exception {
+    final DataRegionId firstRegion = new DataRegionId(1);
+    final DataRegionId secondRegion = new DataRegionId(2);
+    createPeer(firstRegion);
+    createPeer(secondRegion);
+    consensus.stop();
+    consensus = newConsensus();
+
+    final Map<ConsensusGroupId, List<String>> notifications = new HashMap<>();
+    IoTConsensus.onPeerRecovering =
+        (groupId, server) -> {
+          assertSame(server, consensus.getImpl(groupId));
+          assertFalse(stateMachines.get(groupId).started);
+          notifications.computeIfAbsent(groupId, ignored -> new ArrayList<>()).add("recovering");
+        };
+    IoTConsensus.onNewPeerCreated =
+        (groupId, server) -> {
+          assertSame(server, consensus.getImpl(groupId));
+          assertTrue(stateMachines.get(groupId).started);
+          notifications.computeIfAbsent(groupId, ignored -> new ArrayList<>()).add("available");
+        };
+
+    consensus.start();
+
+    assertEquals(2, notifications.size());
+    assertEquals(Arrays.asList("recovering", "available"), notifications.get(firstRegion));
+    assertEquals(Arrays.asList("recovering", "available"), notifications.get(secondRegion));
+  }
+
+  @Test
+  public void testRecoveryContinuesAfterPeerNotificationFailure() throws Exception {
+    final DataRegionId firstRegion = new DataRegionId(1);
+    final DataRegionId secondRegion = new DataRegionId(2);
+    createPeer(firstRegion);
+    createPeer(secondRegion);
+    consensus.stop();
+    consensus = newConsensus();
+
+    final List<ConsensusGroupId> notifiedRegions = new ArrayList<>();
+    IoTConsensus.onNewPeerCreated =
+        (groupId, server) -> {
+          assertTrue(stateMachines.get(groupId).started);
+          notifiedRegions.add(groupId);
+          throw new IllegalStateException("subscription binding failed");
+        };
+
+    consensus.start();
+
+    assertEquals(2, notifiedRegions.size());
+    assertTrue(notifiedRegions.containsAll(Arrays.asList(firstRegion, secondRegion)));
+    assertTrue(stateMachines.get(firstRegion).started);
+    assertTrue(stateMachines.get(secondRegion).started);
+  }
+
+  private IoTConsensus newConsensus() {
+    final IoTConsensus instance =
+        new IoTConsensus(
+            config,
+            groupId -> {
+              final LifecycleStateMachine stateMachine = new LifecycleStateMachine();
+              stateMachines.put(groupId, stateMachine);
+              return stateMachine;
+            });
+    instance.recordCorrectPeerListBeforeStarting(peerLists);
+    return instance;
+  }
+
+  private void createPeer(final DataRegionId regionId) throws Exception {
+    final List<Peer> peers =
+        Collections.singletonList(
+            new Peer(regionId, config.getThisNodeId(), config.getThisNodeEndPoint()));
+    consensus.createLocalPeer(regionId, peers);
+    peerLists.put(regionId, peers);
+  }
+
+  private static final class LifecycleStateMachine extends TestStateMachine {
+
+    private boolean started;
+
+    @Override
+    public void start() {
+      started = true;
+    }
+
+    @Override
+    public void stop() {
+      started = false;
+    }
+  }
+}