Subscription: support consumer timeout and idle disconnect

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java

@@ -80,6 +80,18 @@ public String getSqlDialect() {
     return getString(ConsumerConstant.SQL_DIALECT_KEY);
   }
 
+  public long getHeartbeatIntervalMs() {
+    return getLongOrDefault(
+        ConsumerConstant.HEARTBEAT_INTERVAL_MS_KEY,
+        ConsumerConstant.HEARTBEAT_INTERVAL_MS_DEFAULT_VALUE);
+  }
+
+  public int getConnectionTimeoutInMs() {
+    return getIntOrDefault(
+        ConsumerConstant.CONNECTION_TIMEOUT_MS_KEY,
+        ConsumerConstant.CONNECTION_TIMEOUT_MS_DEFAULT_VALUE);
+  }
+
   public void setConsumerId(final String consumerId) {
     attributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
   }

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java

@@ -55,6 +55,9 @@ public class ConsumerConstant {
 
   public static final String THRIFT_MAX_FRAME_SIZE_KEY = "thrift-max-frame-size";
 
+  public static final String CONNECTION_TIMEOUT_MS_KEY = "connection-timeout-ms";
+  public static final int CONNECTION_TIMEOUT_MS_DEFAULT_VALUE = 0;
+
   public static final String MAX_POLL_PARALLELISM_KEY = "max-poll-parallelism";
   public static final int MAX_POLL_PARALLELISM_DEFAULT_VALUE = 1;
 

iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java

@@ -41,6 +41,7 @@
   public ZoneId zoneId = null;
   public int thriftDefaultBufferSize = SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY;
   public int thriftMaxFrameSize = SessionConfig.DEFAULT_MAX_FRAME_SIZE;
+  public int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS;
   // this field only take effect in write request, nothing to do with any other type requests,
   // like query, load and so on.
   // if set to true, it means that we may redirect the write request to its corresponding leader

iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java

@@ -470,6 +470,7 @@ public Session(AbstractSessionBuilder builder) {
     this.zoneId = builder.zoneId;
     this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize;
     this.thriftMaxFrameSize = builder.thriftMaxFrameSize;
+    this.connectionTimeoutInMs = builder.connectionTimeoutInMs;
     this.version = builder.version;
     this.useSSL = builder.useSSL;
     this.trustStore = builder.trustStore;

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java

@@ -36,7 +36,7 @@ public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) {
   }
 
   public void open() throws IoTDBConnectionException {
-    super.open();
+    super.open(enableThriftRpcCompaction, connectionTimeoutInMs);
   }
 
   public void close() throws IoTDBConnectionException {

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java

@@ -58,6 +58,11 @@ public SubscriptionTableSessionBuilder thriftMaxFrameSize(final int thriftMaxFra
     return this;
   }
 
+  public SubscriptionTableSessionBuilder connectionTimeoutInMs(final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs = connectionTimeoutInMs;
+    return this;
+  }
+
   public ISubscriptionTableSession build() throws IoTDBConnectionException {
     final ISubscriptionTableSession session = new SubscriptionTableSession(this);
     session.open();

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTreeSessionBuilder.java

@@ -57,6 +57,11 @@ public SubscriptionTreeSessionBuilder thriftMaxFrameSize(final int thriftMaxFram
     return this;
   }
 
+  public SubscriptionTreeSessionBuilder connectionTimeoutInMs(final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs = connectionTimeoutInMs;
+    return this;
+  }
+
   public ISubscriptionTreeSession build() {
     return new SubscriptionTreeSession(this);
   }

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java

@@ -118,6 +118,7 @@
   private final Set<SubscriptionCommitContext> inFlightFilesCommitContextSet = new HashSet<>();
 
   private final int thriftMaxFrameSize;
+  private final int connectionTimeoutInMs;
   private final int maxPollParallelism;
 
   @SuppressWarnings("java:S3077")
@@ -187,6 +188,7 @@
     this.fileSaveFsync = builder.fileSaveFsync;
 
     this.thriftMaxFrameSize = builder.thriftMaxFrameSize;
+    this.connectionTimeoutInMs = builder.connectionTimeoutInMs;
     this.maxPollParallelism = builder.maxPollParallelism;
   }
 
@@ -232,6 +234,11 @@
                     properties.getOrDefault(
                         ConsumerConstant.THRIFT_MAX_FRAME_SIZE_KEY,
                         SessionConfig.DEFAULT_MAX_FRAME_SIZE))
+            .connectionTimeoutInMs(
+                (Integer)
+                    properties.getOrDefault(
+                        ConsumerConstant.CONNECTION_TIMEOUT_MS_KEY,
+                        SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS))
             .maxPollParallelism(
                 (Integer)
                     properties.getOrDefault(
@@ -376,13 +383,15 @@
 
   /////////////////////////////// subscription provider ///////////////////////////////
 
   protected abstract AbstractSubscriptionProvider constructSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize);
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs);
 
   AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPoint)
       throws SubscriptionException {
@@ -393,7 +402,9 @@
             this.password,
             this.consumerId,
             this.consumerGroupId,
-            this.thriftMaxFrameSize);
+            this.thriftMaxFrameSize,
+            this.heartbeatIntervalMs,
+            this.connectionTimeoutInMs);
     try {
       provider.handshake();
     } catch (final Exception e) {
@@ -1428,6 +1439,7 @@
     result.put("fileSaveFsync", String.valueOf(fileSaveFsync));
     result.put("inFlightFilesCommitContextSet", inFlightFilesCommitContextSet.toString());
     result.put("thriftMaxFrameSize", String.valueOf(thriftMaxFrameSize));
+    result.put("connectionTimeoutInMs", String.valueOf(connectionTimeoutInMs));
     result.put("maxPollParallelism", String.valueOf(maxPollParallelism));
     result.put("subscribedTopics", subscribedTopics.toString());
     return result;

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java

@@ -48,6 +48,7 @@ public class AbstractSubscriptionConsumerBuilder {
   protected boolean fileSaveFsync = ConsumerConstant.FILE_SAVE_FSYNC_DEFAULT_VALUE;
 
   protected int thriftMaxFrameSize = SessionConfig.DEFAULT_MAX_FRAME_SIZE;
+  protected int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS;
   protected int maxPollParallelism = ConsumerConstant.MAX_POLL_PARALLELISM_DEFAULT_VALUE;
 
   public AbstractSubscriptionConsumerBuilder host(final String host) {
@@ -120,6 +121,12 @@ public AbstractSubscriptionConsumerBuilder thriftMaxFrameSize(final int thriftMa
     return this;
   }
 
+  public AbstractSubscriptionConsumerBuilder connectionTimeoutInMs(
+      final int connectionTimeoutInMs) {
+    this.connectionTimeoutInMs = Math.max(connectionTimeoutInMs, 0);
+    return this;
+  }
+
   public AbstractSubscriptionConsumerBuilder maxPollParallelism(final int maxPollParallelism) {
     // Here the minimum value of max poll parallelism is set to 1 instead of 0, in order to use a
     // single thread to execute poll whenever there are idle resources available, thereby

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java

@@ -90,30 +90,42 @@
 
   private final String username;
   private final String password;
+  private final long heartbeatIntervalMs;
+  private final int connectionTimeoutInMs;
 
   protected abstract AbstractSessionBuilder constructSubscriptionSessionBuilder(
       final String host,
       final int port,
       final String username,
       final String password,
-      final int thriftMaxFrameSize);
+      final int thriftMaxFrameSize,
+      final int connectionTimeoutInMs);
 
   protected AbstractSubscriptionProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize) {
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs) {
     this.session =
         new SubscriptionSessionWrapper(
             constructSubscriptionSessionBuilder(
-                endPoint.ip, endPoint.port, username, password, thriftMaxFrameSize));
+                endPoint.ip,
+                endPoint.port,
+                username,
+                password,
+                thriftMaxFrameSize,
+                connectionTimeoutInMs));
     this.endPoint = endPoint;
     this.consumerId = consumerId;
     this.consumerGroupId = consumerGroupId;
     this.username = username;
     this.password = password;
+    this.heartbeatIntervalMs = heartbeatIntervalMs;
+    this.connectionTimeoutInMs = connectionTimeoutInMs;
   }
 
   SubscriptionSessionConnection getSessionConnection() throws IoTDBConnectionException {
@@ -164,6 +176,10 @@
     consumerAttributes.put(ConsumerConstant.USERNAME_KEY, username);
     consumerAttributes.put(ConsumerConstant.PASSWORD_KEY, password);
     consumerAttributes.put(ConsumerConstant.SQL_DIALECT_KEY, session.getSqlDialect());
+    consumerAttributes.put(
+        ConsumerConstant.HEARTBEAT_INTERVAL_MS_KEY, String.valueOf(heartbeatIntervalMs));
+    consumerAttributes.put(
+        ConsumerConstant.CONNECTION_TIMEOUT_MS_KEY, String.valueOf(connectionTimeoutInMs));
 
     final PipeSubscribeHandshakeResp resp =
         handshake(new ConsumerConfig(consumerAttributes)); // throw SubscriptionException

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumerBuilder.java

@@ -102,6 +102,13 @@ public AbstractSubscriptionPullConsumerBuilder thriftMaxFrameSize(final int thri
     return this;
   }
 
+  @Override
+  public AbstractSubscriptionPullConsumerBuilder connectionTimeoutInMs(
+      final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs(connectionTimeoutInMs);
+    return this;
+  }
+
   @Override
   public AbstractSubscriptionPullConsumerBuilder maxPollParallelism(final int maxPollParallelism) {
     super.maxPollParallelism(maxPollParallelism);

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumerBuilder.java

@@ -108,6 +108,13 @@ public AbstractSubscriptionPushConsumerBuilder thriftMaxFrameSize(final int thri
     return this;
   }
 
+  @Override
+  public AbstractSubscriptionPushConsumerBuilder connectionTimeoutInMs(
+      final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs(connectionTimeoutInMs);
+    return this;
+  }
+
   @Override
   public AbstractSubscriptionPushConsumerBuilder maxPollParallelism(final int maxPollParallelism) {
     super.maxPollParallelism(maxPollParallelism);

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTableProvider.java

@@ -26,14 +26,24 @@
 
 final class SubscriptionTableProvider extends AbstractSubscriptionProvider {
 
   SubscriptionTableProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize) {
-    super(endPoint, username, password, consumerId, consumerGroupId, thriftMaxFrameSize);
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs) {
+    super(
+        endPoint,
+        username,
+        password,
+        consumerId,
+        consumerGroupId,
+        thriftMaxFrameSize,
+        heartbeatIntervalMs,
+        connectionTimeoutInMs);
   }
 
   @Override
@@ -42,12 +52,14 @@
       final int port,
       final String username,
       final String password,
-      final int thriftMaxFrameSize) {
+      final int thriftMaxFrameSize,
+      final int connectionTimeoutInMs) {
     return new SubscriptionTableSessionBuilder()
         .host(host)
         .port(port)
         .username(username)
         .password(password)
-        .thriftMaxFrameSize(thriftMaxFrameSize);
+        .thriftMaxFrameSize(thriftMaxFrameSize)
+        .connectionTimeoutInMs(connectionTimeoutInMs);
   }
 }

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumer.java

@@ -44,9 +44,18 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize) {
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs) {
     return new SubscriptionTableProvider(
-        endPoint, username, password, consumerId, consumerGroupId, thriftMaxFrameSize);
+        endPoint,
+        username,
+        password,
+        consumerId,
+        consumerGroupId,
+        thriftMaxFrameSize,
+        heartbeatIntervalMs,
+        connectionTimeoutInMs);
   }
 
   /////////////////////////////// ctor ///////////////////////////////

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java

@@ -99,6 +99,13 @@ public SubscriptionTablePullConsumerBuilder thriftMaxFrameSize(final int thriftM
     return this;
   }
 
+  @Override
+  public SubscriptionTablePullConsumerBuilder connectionTimeoutInMs(
+      final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs(connectionTimeoutInMs);
+    return this;
+  }
+
   @Override
   public SubscriptionTablePullConsumerBuilder maxPollParallelism(final int maxPollParallelism) {
     super.maxPollParallelism(maxPollParallelism);

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumer.java

@@ -40,9 +40,18 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize) {
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs) {
     return new SubscriptionTableProvider(
-        endPoint, username, password, consumerId, consumerGroupId, thriftMaxFrameSize);
+        endPoint,
+        username,
+        password,
+        consumerId,
+        consumerGroupId,
+        thriftMaxFrameSize,
+        heartbeatIntervalMs,
+        connectionTimeoutInMs);
   }
 
   /////////////////////////////// ctor ///////////////////////////////

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java

@@ -101,6 +101,13 @@ public SubscriptionTablePushConsumerBuilder thriftMaxFrameSize(final int thriftM
     return this;
   }
 
+  @Override
+  public SubscriptionTablePushConsumerBuilder connectionTimeoutInMs(
+      final int connectionTimeoutInMs) {
+    super.connectionTimeoutInMs(connectionTimeoutInMs);
+    return this;
+  }
+
   @Override
   public SubscriptionTablePushConsumerBuilder maxPollParallelism(final int maxPollParallelism) {
     super.maxPollParallelism(maxPollParallelism);

iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreeProvider.java

@@ -26,14 +26,24 @@
 
 final class SubscriptionTreeProvider extends AbstractSubscriptionProvider {
 
   SubscriptionTreeProvider(
       final TEndPoint endPoint,
       final String username,
       final String password,
       final String consumerId,
       final String consumerGroupId,
-      final int thriftMaxFrameSize) {
-    super(endPoint, username, password, consumerId, consumerGroupId, thriftMaxFrameSize);
+      final int thriftMaxFrameSize,
+      final long heartbeatIntervalMs,
+      final int connectionTimeoutInMs) {
+    super(
+        endPoint,
+        username,
+        password,
+        consumerId,
+        consumerGroupId,
+        thriftMaxFrameSize,
+        heartbeatIntervalMs,
+        connectionTimeoutInMs);
   }
 
   @Override
@@ -42,12 +52,14 @@
       final int port,
       final String username,
       final String password,
-      final int thriftMaxFrameSize) {
+      final int thriftMaxFrameSize,
+      final int connectionTimeoutInMs) {
     return new SubscriptionTreeSessionBuilder()
         .host(host)
         .port(port)
         .username(username)
         .password(password)
-        .thriftMaxFrameSize(thriftMaxFrameSize);
+        .thriftMaxFrameSize(thriftMaxFrameSize)
+        .connectionTimeoutInMs(connectionTimeoutInMs);
   }
 }