Skip to content

Subscription: clean up consumers after abrupt disconnect #18434

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

Open
Caideyipi wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Subscription: clean up consumers after abrupt disconnect #18434

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
242 changes: 214 additions & 28 deletions ...anode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
import org.apache.iotdb.db.subscription.receiver.SubscriptionReceiverV1;
import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.rpc.subscription.config.ConsumerConfig;
import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion;
import org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseType;
import org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseVersion;
Expand All @@ -37,9 +39,10 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
Expand All @@ -49,9 +52,6 @@ public class SubscriptionReceiverAgent {

private static final Logger LOGGER = LoggerFactory.getLogger(SubscriptionReceiverAgent.class);

private static final Map<Byte, Supplier<SubscriptionReceiver>> RECEIVER_CONSTRUCTORS =
new HashMap<>();

private static final TPipeSubscribeResp SUBSCRIPTION_NOT_ENABLED_ERROR_RESP =
new TPipeSubscribeResp(
RpcUtils.getStatus(
Expand All @@ -60,21 +60,44 @@ public class SubscriptionReceiverAgent {
PipeSubscribeResponseVersion.VERSION_1.getVersion(),
PipeSubscribeResponseType.ACK.getType());

private final Map<Byte, Supplier<SubscriptionReceiver>> receiverConstructors = new HashMap<>();
private final ThreadLocal<SubscriptionReceiver> receiverThreadLocal = new ThreadLocal<>();
private final Set<SubscriptionReceiver> activeReceivers = ConcurrentHashMap.newKeySet();
private final ScheduledExecutorService receiverTimeoutChecker =
IoTDBThreadPoolFactory.newSingleThreadScheduledExecutor(
SubscriptionReceiverAgent.class.getSimpleName() + "-Timeout-Checker");

/**
* The receiver currently serving each consumer identity. A disconnected receiver deliberately
* remains in this map until its inactivity timeout closes the consumer, while a reconnecting
* receiver replaces it atomically through {@link ConcurrentHashMap#compute(Object,
* java.util.function.BiFunction)}.
*/
private final ConcurrentHashMap<ConsumerIdentity, SubscriptionReceiver> consumerReceivers =
new ConcurrentHashMap<>();

private final ScheduledExecutorService receiverTimeoutChecker;

SubscriptionReceiverAgent() {
RECEIVER_CONSTRUCTORS.put(
PipeSubscribeRequestVersion.VERSION_1.getVersion(), SubscriptionReceiverV1::new);
ScheduledExecutorUtil.safelyScheduleWithFixedDelay(
receiverTimeoutChecker,
this::checkReceiverTimeouts,
Math.max(1_000L, SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
Math.max(1_000L, SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
TimeUnit.MILLISECONDS);
this(SubscriptionReceiverV1::new, true);
}

SubscriptionReceiverAgent(
final Supplier<SubscriptionReceiver> receiverConstructor,
final boolean scheduleTimeoutChecker) {
receiverConstructors.put(
PipeSubscribeRequestVersion.VERSION_1.getVersion(), receiverConstructor);
if (scheduleTimeoutChecker) {
receiverTimeoutChecker =
IoTDBThreadPoolFactory.newSingleThreadScheduledExecutor(
SubscriptionReceiverAgent.class.getSimpleName() + "-Timeout-Checker");
ScheduledExecutorUtil.safelyScheduleWithFixedDelay(
receiverTimeoutChecker,
this::checkReceiverTimeouts,
Math.max(
1_000L, SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
Math.max(
1_000L, SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
TimeUnit.MILLISECONDS);
} else {
receiverTimeoutChecker = null;
}
}

public TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
Expand All @@ -93,12 +116,48 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req, final String usern
}

final byte reqVersion = req.getVersion();
if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
if (receiverConstructors.containsKey(reqVersion)) {
final SubscriptionReceiver receiver = getReceiver(reqVersion);
receiver.setAuthenticatedUsername(username);
activeReceivers.add(receiver);
receiver.handleTimeout();
return receiver.handle(req);
final ConsumerIdentity consumerIdentity = getConsumerIdentity(req, receiver);
final RequestResult requestResult = new RequestResult();

if (Objects.isNull(consumerIdentity)) {
requestResult.response = handleRequest(receiver, req, null);
} else {
consumerReceivers.compute(
consumerIdentity,
(identity, currentReceiver) -> {
requestResult.response = handleRequest(receiver, req, currentReceiver);

if (isHandshake(req)) {
if (isSuccessful(requestResult.response)) {
if (currentReceiver != null && currentReceiver != receiver) {
currentReceiver.invalidateConsumer();
}
return receiver;
}
return currentReceiver;
}

if (currentReceiver != null && currentReceiver != receiver) {
return currentReceiver;
}
return receiver.hasActiveConsumer() ? receiver : null;
});
}

if (isHandshake(req) && isSuccessful(requestResult.response)) {
final ConsumerIdentity activeIdentity = getConsumerIdentity(receiver);
if (!Objects.equals(consumerIdentity, activeIdentity)) {
registerReceiver(receiver, activeIdentity);
} else {
removeReceiverMappingsExcept(receiver, activeIdentity);
}
} else if (isClose(req) && isSuccessful(requestResult.response)) {
removeReceiverMappings(receiver);
}
return requestResult.response;
} else {
final TSStatus status =
RpcUtils.getStatus(
Expand All @@ -120,7 +179,7 @@ public long remainingMs() {
}

public long remainingMs(final byte reqVersion) {
if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
if (receiverConstructors.containsKey(reqVersion)) {
return getReceiver(reqVersion).remainingMs();
} else {
return SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs();
Expand All @@ -147,8 +206,8 @@ private SubscriptionReceiver getReceiver(final byte reqVersion) {
}

private SubscriptionReceiver setAndGetReceiver(final byte reqVersion) {
if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
receiverThreadLocal.set(RECEIVER_CONSTRUCTORS.get(reqVersion).get());
if (receiverConstructors.containsKey(reqVersion)) {
receiverThreadLocal.set(receiverConstructors.get(reqVersion).get());
} else {
throw new UnsupportedOperationException(
String.format(
Expand All @@ -162,13 +221,140 @@ private SubscriptionReceiver setAndGetReceiver(final byte reqVersion) {
public final void handleClientExit() {
final SubscriptionReceiver receiver = receiverThreadLocal.get();
if (receiver != null) {
activeReceivers.remove(receiver);
receiver.handleExit();
receiverThreadLocal.remove();
try {
final ConsumerIdentity consumerIdentity = getConsumerIdentity(receiver);
if (Objects.isNull(consumerIdentity)) {
receiver.handleExit();
} else {
consumerReceivers.compute(
consumerIdentity,
(identity, currentReceiver) -> {
if (currentReceiver != null && currentReceiver != receiver) {
// A newer connection has already taken over this consumer. Do not let the old
// connection's exit cleanup touch the new owner's subscription state.
receiver.invalidateConsumer();
receiver.handleExit();
return currentReceiver;
}
receiver.handleExit();
return receiver.hasActiveConsumer() ? receiver : null;
});
}
} finally {
receiverThreadLocal.remove();
}
}
}

void checkReceiverTimeouts() {
consumerReceivers.forEach(
(identity, receiver) ->
consumerReceivers.computeIfPresent(
identity,
(currentIdentity, currentReceiver) -> {
if (currentReceiver != receiver) {
return currentReceiver;
}
if (!identity.equals(getConsumerIdentity(receiver))) {
return null;
}
receiver.handleTimeout();
return receiver.hasActiveConsumer() ? receiver : null;
}));
}

private TPipeSubscribeResp handleRequest(
final SubscriptionReceiver receiver,
final TPipeSubscribeReq req,
final SubscriptionReceiver currentReceiver) {
if (!isHandshake(req) && currentReceiver != null && currentReceiver != receiver) {
receiver.invalidateConsumer();
}
return receiver.handle(req);
}

private void registerReceiver(
final SubscriptionReceiver receiver, final ConsumerIdentity identity) {
if (Objects.isNull(identity)) {
removeReceiverMappings(receiver);
return;
}
consumerReceivers.compute(
identity,
(key, currentReceiver) -> {
if (currentReceiver != null && currentReceiver != receiver) {
currentReceiver.invalidateConsumer();
}
return receiver;
});
removeReceiverMappingsExcept(receiver, identity);
}

private void removeReceiverMappingsExcept(
final SubscriptionReceiver receiver, final ConsumerIdentity retainedIdentity) {
consumerReceivers.forEach(
(registeredIdentity, currentReceiver) -> {
if (currentReceiver == receiver
&& !Objects.equals(retainedIdentity, registeredIdentity)) {
consumerReceivers.remove(registeredIdentity, receiver);
}
});
}

private void removeReceiverMappings(final SubscriptionReceiver receiver) {
consumerReceivers.forEach(
(identity, currentReceiver) -> consumerReceivers.remove(identity, receiver));
}

private static ConsumerIdentity getConsumerIdentity(
final TPipeSubscribeReq req, final SubscriptionReceiver receiver) {
if (isHandshake(req) && req.isSetBody()) {
try {
final ByteBuffer body = req.bufferForBody();
if (body.hasRemaining()) {
final ConsumerConfig consumerConfig = ConsumerConfig.deserialize(body);
final ConsumerIdentity identity =
ConsumerIdentity.of(
consumerConfig.getConsumerGroupId(), consumerConfig.getConsumerId());
if (Objects.nonNull(identity)) {
return identity;
}
}
} catch (final RuntimeException ignored) {
// Let the receiver report the malformed handshake request. It still needs to see the
// original buffer, so parsing is intentionally done on a duplicate above.
}
}
return getConsumerIdentity(receiver);
}

private static ConsumerIdentity getConsumerIdentity(final SubscriptionReceiver receiver) {
return ConsumerIdentity.of(receiver.getConsumerGroupId(), receiver.getConsumerId());
}

private void checkReceiverTimeouts() {
activeReceivers.forEach(SubscriptionReceiver::handleTimeout);
private static boolean isHandshake(final TPipeSubscribeReq req) {
return req.getType() == PipeSubscribeRequestType.HANDSHAKE.getType();
}

private static boolean isClose(final TPipeSubscribeReq req) {
return req.getType() == PipeSubscribeRequestType.CLOSE.getType();
}

private static boolean isSuccessful(final TPipeSubscribeResp response) {
return response != null
&& response.getStatus() != null
&& response.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode();
}

private static final class RequestResult {
private TPipeSubscribeResp response;
}

private record ConsumerIdentity(String consumerGroupId, String consumerId) {
private static ConsumerIdentity of(final String consumerGroupId, final String consumerId) {
return Objects.isNull(consumerGroupId) || Objects.isNull(consumerId)
? null
: new ConsumerIdentity(consumerGroupId, consumerId);
}
}
}
21 changes: 21 additions & 0 deletions ...atanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,5 +35,26 @@ public interface SubscriptionReceiver {

void handleTimeout();

/**
* Returns the identity of the consumer currently associated with this receiver, or {@code null}
* if the receiver has not completed a handshake (or has already been invalidated).
*/
String getConsumerId();

/**
* Returns the consumer group currently associated with this receiver, or {@code null} if the
* receiver has not completed a handshake (or has already been invalidated).
*/
String getConsumerGroupId();

/**
* Invalidates this receiver so that requests from an obsolete connection cannot affect a new
* owner.
*/
void invalidateConsumer();

/** Returns whether this receiver still owns an active consumer timeout state. */
boolean hasActiveConsumer();

long remainingMs();
}
Loading
Loading