Subscription: add topic owner epoch fencing

Author: codex

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java

@@ -316,6 +316,10 @@ public enum TSStatusCode {
   SHOW_SUBSCRIPTION_ERROR(1910),
   SUBSCRIPTION_PIPE_TIMEOUT_ERROR(1911),
   SUBSCRIPTION_NOT_ENABLED_ERROR(1912),
+  SUBSCRIPTION_OWNER_FENCED(1913),
+  SUBSCRIPTION_OWNER_REQUIRED(1914),
+  SUBSCRIPTION_OWNER_EPOCH_REQUIRED(1915),
+  SUBSCRIPTION_OWNER_LEASE_EXPIRED(1916),
 
   // Topic
   CREATE_TOPIC_ERROR(2000),

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

@@ -68,6 +68,16 @@ public String getConsumerGroupId() {
     return getString(ConsumerConstant.CONSUMER_GROUP_ID_KEY);
   }
 
+  public String getOwnerId() {
+    return getString(ConsumerConstant.OWNER_ID_KEY);
+  }
+
+  public Long getOwnerEpoch() {
+    return hasAttribute(ConsumerConstant.OWNER_EPOCH_KEY)
+        ? getLong(ConsumerConstant.OWNER_EPOCH_KEY)
+        : null;
+  }
+
   public String getUsername() {
     return getString(ConsumerConstant.USERNAME_KEY);
   }

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

@@ -40,6 +40,8 @@ public class ConsumerConstant {
 
   public static final String CONSUMER_ID_KEY = "consumer-id";
   public static final String CONSUMER_GROUP_ID_KEY = "group-id";
+  public static final String OWNER_ID_KEY = "owner-id";
+  public static final String OWNER_EPOCH_KEY = "owner-epoch";
 
   public static final String HEARTBEAT_INTERVAL_MS_KEY = "heartbeat-interval-ms";
   public static final long HEARTBEAT_INTERVAL_MS_DEFAULT_VALUE = 30_000L;

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

@@ -61,6 +61,10 @@ public class TopicConstant {
   public static final String STRICT_KEY = "strict";
   public static final String STRICT_DEFAULT_VALUE = "true";
 
+  public static final String OWNER_ID_KEY = "owner-id";
+  public static final String OWNER_EPOCH_KEY = "owner-epoch";
+  public static final String OWNER_LEASE_EXPIRE_TIME_MS_KEY = "owner-lease-expire-time-ms";
+
   private TopicConstant() {
     throw new IllegalStateException(SubscriptionMessages.UTILITY_CLASS);
   }

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

@@ -105,6 +105,8 @@ abstract class AbstractSubscriptionConsumer implements AutoCloseable {
 
   protected String consumerId;
   protected String consumerGroupId;
+  protected String ownerId;
+  protected Long ownerEpoch;
 
   private final long heartbeatIntervalMs;
   private final long endpointsSyncIntervalMs;
@@ -154,6 +156,14 @@ public String getConsumerGroupId() {
     return consumerGroupId;
   }
 
+  public String getOwnerId() {
+    return ownerId;
+  }
+
+  public Long getOwnerEpoch() {
+    return ownerEpoch;
+  }
+
   /////////////////////////////// ctor ///////////////////////////////
 
   protected AbstractSubscriptionConsumer(final AbstractSubscriptionConsumerBuilder builder) {
@@ -183,6 +193,8 @@ protected AbstractSubscriptionConsumer(final AbstractSubscriptionConsumerBuilder
 
     this.consumerId = builder.consumerId;
     this.consumerGroupId = builder.consumerGroupId;
+    this.ownerId = builder.ownerId;
+    this.ownerEpoch = builder.ownerEpoch;
 
     this.heartbeatIntervalMs = builder.heartbeatIntervalMs;
     this.endpointsSyncIntervalMs = builder.endpointsSyncIntervalMs;
@@ -213,6 +225,8 @@ protected AbstractSubscriptionConsumer(
             .encryptedPassword((String) properties.get(ConsumerConstant.ENCRYPTED_PASSWORD_KEY))
             .consumerId((String) properties.get(ConsumerConstant.CONSUMER_ID_KEY))
             .consumerGroupId((String) properties.get(ConsumerConstant.CONSUMER_GROUP_ID_KEY))
+            .ownerId((String) properties.get(ConsumerConstant.OWNER_ID_KEY))
+            .ownerEpoch((Long) properties.get(ConsumerConstant.OWNER_EPOCH_KEY))
             .heartbeatIntervalMs(
                 (Long)
                     properties.getOrDefault(
@@ -394,6 +408,8 @@ protected abstract AbstractSubscriptionProvider constructSubscriptionProvider(
       final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
+      final String ownerId,
+      final Long ownerEpoch,
       final int thriftMaxFrameSize,
       final long heartbeatIntervalMs,
       final int connectionTimeoutInMs);
@@ -408,6 +424,8 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo
             this.encryptedPassword,
             this.consumerId,
             this.consumerGroupId,
+            this.ownerId,
+            this.ownerEpoch,
             this.thriftMaxFrameSize,
             this.heartbeatIntervalMs,
             this.connectionTimeoutInMs);
@@ -1429,6 +1447,8 @@ protected Map<String, String> coreReportMessage() {
     final Map<String, String> result = new HashMap<>();
     result.put("consumerId", consumerId);
     result.put("consumerGroupId", consumerGroupId);
+    result.put("ownerId", ownerId);
+    result.put("ownerEpoch", String.valueOf(ownerEpoch));
     result.put("isClosed", isClosed.toString());
     result.put("fileSaveDir", fileSaveDir);
     result.put(
@@ -1443,6 +1463,8 @@ protected Map<String, String> allReportMessage() {
     final Map<String, String> result = new HashMap<>();
     result.put("consumerId", consumerId);
     result.put("consumerGroupId", consumerGroupId);
+    result.put("ownerId", ownerId);
+    result.put("ownerEpoch", String.valueOf(ownerEpoch));
     result.put("heartbeatIntervalMs", String.valueOf(heartbeatIntervalMs));
     result.put("endpointsSyncIntervalMs", String.valueOf(endpointsSyncIntervalMs));
     result.put("providers", providers.toString());

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

@@ -40,6 +40,8 @@ public class AbstractSubscriptionConsumerBuilder {
 
   protected String consumerId;
   protected String consumerGroupId;
+  protected String ownerId;
+  protected Long ownerEpoch;
 
   protected long heartbeatIntervalMs = ConsumerConstant.HEARTBEAT_INTERVAL_MS_DEFAULT_VALUE;
   protected long endpointsSyncIntervalMs =
@@ -111,6 +113,27 @@ public AbstractSubscriptionConsumerBuilder consumerGroupId(
     return this;
   }
 
+  public AbstractSubscriptionConsumerBuilder ownerId(@Nullable final String ownerId) {
+    if (Objects.isNull(ownerId)) {
+      return this;
+    }
+    this.ownerId = ownerId;
+    return this;
+  }
+
+  public AbstractSubscriptionConsumerBuilder ownerEpoch(final long ownerEpoch) {
+    this.ownerEpoch = ownerEpoch;
+    return this;
+  }
+
+  public AbstractSubscriptionConsumerBuilder ownerEpoch(@Nullable final Long ownerEpoch) {
+    if (Objects.isNull(ownerEpoch)) {
+      return this;
+    }
+    this.ownerEpoch = ownerEpoch;
+    return this;
+  }
+
   public AbstractSubscriptionConsumerBuilder heartbeatIntervalMs(final long heartbeatIntervalMs) {
     this.heartbeatIntervalMs =
         Math.max(heartbeatIntervalMs, ConsumerConstant.HEARTBEAT_INTERVAL_MS_MIN_VALUE);

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

@@ -81,6 +81,8 @@ public abstract class AbstractSubscriptionProvider {
 
   private String consumerId;
   private String consumerGroupId;
+  private final String ownerId;
+  private final Long ownerEpoch;
 
   private final AtomicBoolean isClosed = new AtomicBoolean(true);
   private final AtomicBoolean isAvailable = new AtomicBoolean(false);
@@ -109,6 +111,8 @@ protected AbstractSubscriptionProvider(
       final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
+      final String ownerId,
+      final Long ownerEpoch,
       final int thriftMaxFrameSize,
       final long heartbeatIntervalMs,
       final int connectionTimeoutInMs) {
@@ -124,6 +128,8 @@ protected AbstractSubscriptionProvider(
     this.endPoint = endPoint;
     this.consumerId = consumerId;
     this.consumerGroupId = consumerGroupId;
+    this.ownerId = ownerId;
+    this.ownerEpoch = ownerEpoch;
     this.username = username;
     this.password = password;
     this.encryptedPassword = encryptedPassword;
@@ -176,6 +182,12 @@ synchronized void handshake() throws SubscriptionException, IoTDBConnectionExcep
     final Map<String, String> consumerAttributes = new HashMap<>();
     consumerAttributes.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, consumerGroupId);
     consumerAttributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
+    if (ownerId != null) {
+      consumerAttributes.put(ConsumerConstant.OWNER_ID_KEY, ownerId);
+    }
+    if (ownerEpoch != null) {
+      consumerAttributes.put(ConsumerConstant.OWNER_EPOCH_KEY, String.valueOf(ownerEpoch));
+    }
     consumerAttributes.put(ConsumerConstant.USERNAME_KEY, username);
     consumerAttributes.put(ConsumerConstant.PASSWORD_KEY, password);
     if (encryptedPassword != null) {

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

@@ -33,6 +33,8 @@ final class SubscriptionTableProvider extends AbstractSubscriptionProvider {
       final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
+      final String ownerId,
+      final Long ownerEpoch,
       final int thriftMaxFrameSize,
       final long heartbeatIntervalMs,
       final int connectionTimeoutInMs) {
@@ -43,6 +45,8 @@ final class SubscriptionTableProvider extends AbstractSubscriptionProvider {
         encryptedPassword,
         consumerId,
         consumerGroupId,
+        ownerId,
+        ownerEpoch,
         thriftMaxFrameSize,
         heartbeatIntervalMs,
         connectionTimeoutInMs);

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

@@ -45,6 +45,8 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
       final String encryptedPassword,
       final String consumerId,
       final String consumerGroupId,
+      final String ownerId,
+      final Long ownerEpoch,
       final int thriftMaxFrameSize,
       final long heartbeatIntervalMs,
       final int connectionTimeoutInMs) {
@@ -55,6 +57,8 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider(
         encryptedPassword,
         consumerId,
         consumerGroupId,
+        ownerId,
+        ownerEpoch,
         thriftMaxFrameSize,
         heartbeatIntervalMs,
         connectionTimeoutInMs);

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

@@ -74,6 +74,24 @@ public SubscriptionTablePullConsumerBuilder consumerGroupId(final String consume
     return this;
   }
 
+  @Override
+  public SubscriptionTablePullConsumerBuilder ownerId(final String ownerId) {
+    super.ownerId(ownerId);
+    return this;
+  }
+
+  @Override
+  public SubscriptionTablePullConsumerBuilder ownerEpoch(final long ownerEpoch) {
+    super.ownerEpoch(ownerEpoch);
+    return this;
+  }
+
+  @Override
+  public SubscriptionTablePullConsumerBuilder ownerEpoch(final Long ownerEpoch) {
+    super.ownerEpoch(ownerEpoch);
+    return this;
+  }
+
   @Override
   public SubscriptionTablePullConsumerBuilder heartbeatIntervalMs(final long heartbeatIntervalMs) {
     super.heartbeatIntervalMs(heartbeatIntervalMs);