Subscription: implement IoTConsensus-based subscription

example/subscription/src/main/java/org/apache/iotdb/ConsensusSubscriptionSessionExample.java

@@ -0,0 +1,198 @@
+/*
+ * 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;
+
+import org.apache.iotdb.isession.ISession;
+import org.apache.iotdb.isession.util.Version;
+import org.apache.iotdb.rpc.subscription.config.TopicConstant;
+import org.apache.iotdb.session.Session;
+import org.apache.iotdb.session.subscription.ISubscriptionTreeSession;
+import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder;
+import org.apache.iotdb.session.subscription.consumer.ISubscriptionTreePullConsumer;
+import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumerBuilder;
+import org.apache.iotdb.session.subscription.payload.SubscriptionMessage;
+import org.apache.iotdb.session.subscription.payload.SubscriptionRecordHandler;
+
+import org.apache.tsfile.read.query.dataset.ResultSet;
+
+import java.util.List;
+import java.util.Properties;
+
+public class ConsensusSubscriptionSessionExample {
+
+  private static final String HOST = "127.0.0.1";
+  private static final int PORT = 6667;
+  private static final String USER = "root";
+  private static final String PASSWORD = "root";
+
+  private static final long POLL_TIMEOUT_MS = 1_000L;
+  private static final int MAX_POLL_ROUNDS = 20;
+  private static final int EXPECTED_ROWS = 5;
+
+  public static void main(final String[] args) throws Exception {
+    final long runId = System.currentTimeMillis();
+    final String database = "root.db_consensus_example_" + runId;
+    final String device = database + ".d0";
+    final String topic = "topic_consensus_example_" + runId;
+    final String consumerGroup = "cg_consensus_example_" + runId;
+    final String consumerId = "consumer_consensus_example_" + runId;
+
+    System.out.println("=== Consensus Subscription Tree Example ===");
+    System.out.println("database = " + database);
+    System.out.println("topic = " + topic);
+    System.out.println("consumerGroup = " + consumerGroup);
+
+    prepareBootstrapData(database, device);
+    createConsensusTopic(topic, database + ".**");
+
+    try (final ISubscriptionTreePullConsumer consumer =
+        new SubscriptionTreePullConsumerBuilder()
+            .host(HOST)
+            .port(PORT)
+            .username(USER)
+            .password(PASSWORD)
+            .consumerId(consumerId)
+            .consumerGroupId(consumerGroup)
+            .autoCommit(false)
+            .build()) {
+      consumer.open();
+      consumer.subscribe(topic);
+
+      writeRealtimeDataAfterSubscribe(device);
+      pollAndCommit(consumer, EXPECTED_ROWS);
+      consumer.unsubscribe(topic);
+    } finally {
+      dropTopic(topic);
+    }
+  }
+
+  private static void prepareBootstrapData(final String database, final String device)
+      throws Exception {
+    try (final ISession session = openSession()) {
+      session.executeNonQueryStatement("CREATE DATABASE " + database);
+      session.executeNonQueryStatement(
+          "CREATE TIMESERIES "
+              + device
+              + ".s1 with datatype=INT64, encoding=RLE, compressor=SNAPPY");
+      session.executeNonQueryStatement(
+          String.format("insert into %s(time, s1) values (%d, %d)", device, 0L, 0L));
+      session.executeNonQueryStatement("flush");
+    }
+  }
+
+  private static void createConsensusTopic(final String topicName, final String path)
+      throws Exception {
+    try (final ISubscriptionTreeSession session =
+        new SubscriptionTreeSessionBuilder()
+            .host(HOST)
+            .port(PORT)
+            .username(USER)
+            .password(PASSWORD)
+            .build()) {
+      session.open();
+
+      final Properties config = new Properties();
+      config.put(TopicConstant.MODE_KEY, TopicConstant.MODE_CONSENSUS_VALUE);
+      config.put(TopicConstant.FORMAT_KEY, TopicConstant.FORMAT_RECORD_HANDLER_VALUE);
+      config.put(TopicConstant.PATH_KEY, path);
+      config.put(TopicConstant.ORDER_MODE_KEY, TopicConstant.ORDER_MODE_PER_WRITER_VALUE);
+      session.createTopicIfNotExists(topicName, config);
+    }
+  }
+
+  private static void writeRealtimeDataAfterSubscribe(final String device) throws Exception {
+    try (final ISession session = openSession()) {
+      for (int i = 1; i <= EXPECTED_ROWS; i++) {
+        session.executeNonQueryStatement(
+            String.format("insert into %s(time, s1) values (%d, %d)", device, i, i * 10L));
+      }
+      session.executeNonQueryStatement("flush");
+    }
+  }
+
+  private static void pollAndCommit(
+      final ISubscriptionTreePullConsumer consumer, final int expectedRows) throws Exception {
+    int totalRows = 0;
+    int consecutiveEmptyPolls = 0;
+
+    for (int round = 1; round <= MAX_POLL_ROUNDS; round++) {
+      final List<SubscriptionMessage> messages = consumer.poll(POLL_TIMEOUT_MS);
+      if (messages.isEmpty()) {
+        consecutiveEmptyPolls++;
+        if (totalRows >= expectedRows && consecutiveEmptyPolls >= 3) {
+          break;
+        }
+        continue;
+      }
+
+      consecutiveEmptyPolls = 0;
+
+      for (final SubscriptionMessage message : messages) {
+        for (final ResultSet resultSet : message.getResultSets()) {
+          final SubscriptionRecordHandler.SubscriptionResultSet subscriptionResultSet =
+              (SubscriptionRecordHandler.SubscriptionResultSet) resultSet;
+          System.out.println("Columns = " + subscriptionResultSet.getColumnNames());
+          while (subscriptionResultSet.hasNext()) {
+            System.out.println(subscriptionResultSet.nextRecord());
+            totalRows++;
+          }
+        }
+      }
+
+      consumer.commitSync(messages);
+      System.out.println("poll round " + round + ", totalRows = " + totalRows);
+    }
+
+    if (totalRows != expectedRows) {
+      throw new IllegalStateException(
+          "Expected "
+              + expectedRows
+              + " realtime rows, but consumed "
+              + totalRows
+              + ". Please check whether consensus subscription is enabled on the server.");
+    }
+  }
+
+  private static void dropTopic(final String topicName) throws Exception {
+    try (final ISubscriptionTreeSession session =
+        new SubscriptionTreeSessionBuilder()
+            .host(HOST)
+            .port(PORT)
+            .username(USER)
+            .password(PASSWORD)
+            .build()) {
+      session.open();
+      session.dropTopicIfExists(topicName);
+    }
+  }
+
+  private static ISession openSession() throws Exception {
+    final ISession session =
+        new Session.Builder()
+            .host(HOST)
+            .port(PORT)
+            .username(USER)
+            .password(PASSWORD)
+            .version(Version.V_1_0)
+            .build();
+    session.open(false);
+    return session;
+  }
+}