Skip to content

Subscription: support unsubscribe from completed topics under client heartbeat thread#15595

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SteveYurongSu merged 4 commits into from
May 29, 2025
Merged

Subscription: support unsubscribe from completed topics under client heartbeat thread #15595

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e5c5db9
setup
VGalaxies May 28, 2025
53dfd04
minor improve
VGalaxies May 28, 2025
f55e05c
Merge branch 'master' into sub-heartbeat-drop
VGalaxies May 28, 2025
d8f4b7f
empty commit
VGalaxies May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions ...iple/treemodel/regression/pullconsumer/mode/IoTDBSnapshotDevicePullConsumerDataSetIT.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -148,5 +148,8 @@ public void do_test()
consume_data(consumer, session_dest);
check_count(8, "select count(s_0) from " + device, "Consume data again:" + pattern);
check_count(8, "select count(s_1) from " + device, "Consumption data: s_1");
while (!consumer.allTopicMessagesHaveBeenConsumed()) {
Comment on lines 150 to +151

CopilotAIMay 29, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This busy-wait loop lacks a timeout or failure condition, which can hang the test indefinitely. Add a maximum wait time or a timeout assertion to prevent CI stalls.

Suggested change
check_count(8, "select count(s_1) from " + device, "Consumption data: s_1");
while (!consumer.allTopicMessagesHaveBeenConsumed()) {
check_count(8, "select count(s_1) from " + device, "Consumption data: s_1");
longstartTime = System.currentTimeMillis();
longtimeout = 30000; // 30 seconds timeout
while (!consumer.allTopicMessagesHaveBeenConsumed()) {
if (System.currentTimeMillis() - startTime > timeout) {
thrownewInterruptedException("Timeout reached while waiting for all topic messages to be consumed.");
}

Copilot uses AI. Check for mistakes.
Thread.sleep(1000);
}
}
}
17 changes: 15 additions & 2 deletions ...n/java/org/apache/iotdb/rpc/subscription/payload/response/PipeSubscribeHeartbeatResp.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -30,8 +30,10 @@
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

Expand All@@ -41,6 +43,9 @@ public class PipeSubscribeHeartbeatResp extends TPipeSubscribeResp {

private transient Map<Integer, TEndPoint> endPoints = new HashMap<>(); // available endpoints

private transient List<String> topicNamesToUnsubscribe =
new ArrayList<>(); // topics should be unsubscribed

public Map<String, TopicConfig> getTopics() {
return topics;
}
Expand All@@ -49,6 +54,10 @@ public Map<Integer, TEndPoint> getEndPoints() {
return endPoints;
}

public List<String> getTopicNamesToUnsubscribe() {
return topicNamesToUnsubscribe;
}

/////////////////////////////// Thrift ///////////////////////////////

/**
Expand All@@ -72,7 +81,8 @@ public static PipeSubscribeHeartbeatResp toTPipeSubscribeResp(final TSStatus sta
public static PipeSubscribeHeartbeatResp toTPipeSubscribeResp(
final TSStatus status,
final Map<String, TopicConfig> topics,
final Map<Integer, TEndPoint> endPoints)
final Map<Integer, TEndPoint> endPoints,
final List<String> topicNamesToUnsubscribe)
throws IOException {
final PipeSubscribeHeartbeatResp resp = toTPipeSubscribeResp(status);

Expand All@@ -89,6 +99,7 @@ public static PipeSubscribeHeartbeatResp toTPipeSubscribeResp(
ReadWriteIOUtils.write(entry.getValue().getIp(), outputStream);
ReadWriteIOUtils.write(entry.getValue().getPort(), outputStream);
}
ReadWriteIOUtils.writeStringList(topicNamesToUnsubscribe, outputStream);
resp.body =
Collections.singletonList(
ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()));
Expand DownExpand Up@@ -126,6 +137,7 @@ public static PipeSubscribeHeartbeatResp fromTPipeSubscribeResp(
}
resp.endPoints = endPoints;
}
resp.topicNamesToUnsubscribe = ReadWriteIOUtils.readStringList(byteBuffer);
break;
}
}
Expand All@@ -152,6 +164,7 @@ public boolean equals(final Object obj) {
final PipeSubscribeHeartbeatResp that = (PipeSubscribeHeartbeatResp) obj;
return Objects.equals(this.topics, that.topics)
&& Objects.equals(this.endPoints, that.endPoints)
&& Objects.equals(this.topicNamesToUnsubscribe, that.topicNamesToUnsubscribe)
&& Objects.equals(this.status, that.status)
&& this.version == that.version
&& this.type == that.type
Expand All@@ -160,6 +173,6 @@ public boolean equals(final Object obj) {

@Override
public int hashCode() {
return Objects.hash(topics, endPoints, status, version, type, body);
return Objects.hash(topics, endPoints, topicNamesToUnsubscribe, status, version, type, body);
}
}
13 changes: 12 additions & 1 deletion ...va/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -23,6 +23,7 @@
import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHeartbeatResp;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand DownExpand Up@@ -244,7 +245,17 @@ void heartbeat(final AbstractSubscriptionConsumer consumer) {
private void heartbeatInternal(final AbstractSubscriptionConsumer consumer) {
for (final AbstractSubscriptionProvider provider : getAllProviders()) {
try {
consumer.subscribedTopics = provider.heartbeat().getTopics();
final PipeSubscribeHeartbeatResp resp = provider.heartbeat();
// update subscribed topics
consumer.subscribedTopics = resp.getTopics();
// unsubscribe completed topics
for (final String topicName : resp.getTopicNamesToUnsubscribe()) {
LOGGER.info(
"Termination occurred when SubscriptionConsumer {} polling topics, unsubscribe topic {} automatically",
consumer.coreReportMessage(),
topicName);
consumer.unsubscribe(topicName);
}
provider.setAvailable();
} catch (final Exception e) {
LOGGER.warn(
Expand Down
11 changes: 11 additions & 0 deletions ...atanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
Expand DownExpand Up@@ -125,6 +126,16 @@ public boolean isCommitContextOutdated(final SubscriptionCommitContext commitCon
return broker.isCommitContextOutdated(commitContext);
}

public List<String> fetchTopicNamesToUnsubscribe(
final ConsumerConfig consumerConfig, final Set<String> topicNames) {
final String consumerGroupId = consumerConfig.getConsumerGroupId();
final SubscriptionBroker broker = consumerGroupIdToSubscriptionBroker.get(consumerGroupId);
if (Objects.isNull(broker)) {
return Collections.emptyList();
}
return broker.fetchTopicNamesToUnsubscribe(topicNames);
}

/////////////////////////////// broker ///////////////////////////////

public boolean isBrokerExist(final String consumerGroupId) {
Expand Down
19 changes: 19 additions & 0 deletions ...re/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -339,6 +339,25 @@ public boolean isCommitContextOutdated(final SubscriptionCommitContext commitCon
return prefetchingQueue.isCommitContextOutdated(commitContext);
}

public List<String> fetchTopicNamesToUnsubscribe(final Set<String> topicNames) {
final List<String> topicNamesToUnsubscribe = new ArrayList<>();

for (final String topicName : topicNames) {
final SubscriptionPrefetchingQueue prefetchingQueue =
topicNameToPrefetchingQueue.get(topicName);
// If there is no prefetching queue for the topic, check if it's completed
if (Objects.isNull(prefetchingQueue) && completedTopicNames.containsKey(topicName)) {
LOGGER.info(

CopilotAIMay 29, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] Logging each completed topic at INFO level on every heartbeat may flood the logs under high heartbeat frequency. Consider lowering this to DEBUG or batching these messages.

Suggested change
LOGGER.info(
LOGGER.debug(

Copilot uses AI. Check for mistakes.
"Subscription: prefetching queue bound to topic [{}] for consumer group [{}] is completed, reply to client heartbeat request",
topicName,
brokerId);
topicNamesToUnsubscribe.add(topicName);
}
}

return topicNamesToUnsubscribe;
}

/////////////////////////////// prefetching queue ///////////////////////////////

public void bindPrefetchingQueue(
Expand Down
6 changes: 5 additions & 1 deletion ...anode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -308,8 +308,12 @@ private TPipeSubscribeResp handlePipeSubscribeHeartbeatInternal(
throw new SubscriptionException(exceptionMessage);
}

// fetch topics should be unsubscribed
final List<String> topicNamesToUnsubscribe =
SubscriptionAgent.broker().fetchTopicNamesToUnsubscribe(consumerConfig, topics.keySet());

return PipeSubscribeHeartbeatResp.toTPipeSubscribeResp(
RpcUtils.SUCCESS_STATUS, topics, endPoints);
RpcUtils.SUCCESS_STATUS, topics, endPoints, topicNamesToUnsubscribe);
}

private TPipeSubscribeResp handlePipeSubscribeSubscribe(final PipeSubscribeSubscribeReq req) {
Expand Down