diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java index 31f4dd46a0e0..d866f50749ce 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java @@ -247,6 +247,9 @@ void heartbeat(final AbstractSubscriptionConsumer consumer) { acquireWriteLock(); try { + if (consumer.isClosed()) { + return; + } heartbeatInternal(consumer); } finally { releaseWriteLock(); @@ -293,6 +296,9 @@ void sync(final AbstractSubscriptionConsumer consumer) { acquireWriteLock(); try { + if (consumer.isClosed()) { + return; + } syncInternal(consumer); } finally { releaseWriteLock(); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java index 98e16ec8cb41..811b4228c1e8 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java @@ -136,10 +136,14 @@ protected synchronized void open() throws SubscriptionException { return; } - super.open(); - // set isClosed to false before submitting workers isClosed.set(false); + try { + super.open(); + } catch (final SubscriptionException e) { + isClosed.set(true); + throw e; + } emptyPollLogThrottler.reset(); // submit auto poll worker if enabling auto commit @@ -199,8 +203,8 @@ public synchronized void close() { commitAllUncommittedMessages(); } - super.close(); isClosed.set(true); + super.close(); } /////////////////////////////// poll & commit /////////////////////////////// diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java index 89cd3a25b044..f3c7d38c4dc7 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java @@ -127,10 +127,14 @@ protected synchronized void open() throws SubscriptionException { return; } - super.open(); - // set isClosed to false before submitting workers isClosed.set(false); + try { + super.open(); + } catch (final SubscriptionException e) { + isClosed.set(true); + throw e; + } emptyPollLogThrottler.reset(); // submit auto poll worker @@ -143,8 +147,8 @@ public synchronized void close() { return; } - super.close(); isClosed.set(true); + super.close(); } @Override diff --git a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java new file mode 100644 index 000000000000..68f714f2ce52 --- /dev/null +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java @@ -0,0 +1,267 @@ +/* + * 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.session.subscription.consumer.base; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.rpc.subscription.exception.SubscriptionException; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHeartbeatResp; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder; + +import org.junit.Assert; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.function.BooleanSupplier; + +public class SubscriptionConsumerLifecycleTest { + + private static final String HOST = "127.0.0.1"; + private static final int PORT = 6667; + private static final String CONSUMER_ID = "test_consumer"; + private static final String CONSUMER_GROUP_ID = "test_consumer_group"; + private static final long LONG_INTERVAL_MS = 86_400_000L; + + @Test + public void testPushConsumerIsOpenBeforeProviderHandshake() throws SubscriptionException { + final TestPushConsumer consumer = new TestPushConsumer(); + + try { + consumer.open(); + + Assert.assertEquals(1, consumer.closedStatesDuringHandshake.size()); + Assert.assertFalse(consumer.closedStatesDuringHandshake.get(0)); + } finally { + consumer.close(); + } + } + + @Test + public void testPushConsumerIsClosedBeforeProviderClose() throws SubscriptionException { + final TestPushConsumer consumer = new TestPushConsumer(); + + consumer.open(); + consumer.close(); + + Assert.assertEquals(1, consumer.closedStatesDuringClose.size()); + Assert.assertTrue(consumer.closedStatesDuringClose.get(0)); + } + + @Test + public void testPullConsumerIsOpenBeforeProviderHandshake() throws SubscriptionException { + final TestPullConsumer consumer = new TestPullConsumer(); + + try { + consumer.open(); + + Assert.assertEquals(1, consumer.closedStatesDuringHandshake.size()); + Assert.assertFalse(consumer.closedStatesDuringHandshake.get(0)); + } finally { + consumer.close(); + } + } + + @Test + public void testPullConsumerIsClosedBeforeProviderClose() throws SubscriptionException { + final TestPullConsumer consumer = new TestPullConsumer(); + + consumer.open(); + consumer.close(); + + Assert.assertEquals(1, consumer.closedStatesDuringClose.size()); + Assert.assertTrue(consumer.closedStatesDuringClose.get(0)); + } + + private static class TestPushConsumer extends AbstractSubscriptionPushConsumer { + + private final List closedStatesDuringHandshake = new ArrayList<>(); + private final List closedStatesDuringClose = new ArrayList<>(); + + private TestPushConsumer() { + super( + new AbstractSubscriptionPushConsumerBuilder() + .host(HOST) + .port(PORT) + .consumerId(CONSUMER_ID) + .consumerGroupId(CONSUMER_GROUP_ID) + .heartbeatIntervalMs(LONG_INTERVAL_MS) + .endpointsSyncIntervalMs(LONG_INTERVAL_MS) + .autoPollIntervalMs(LONG_INTERVAL_MS)); + } + + @Override + protected AbstractSubscriptionProvider constructSubscriptionProvider( + final TEndPoint endPoint, + final String username, + final String password, + final String encryptedPassword, + final String consumerId, + final String consumerGroupId, + final String ownerId, + final Long ownerEpoch, + final int thriftMaxFrameSize, + final long heartbeatIntervalMs, + final int connectionTimeoutInMs) { + return new TestSubscriptionProvider( + endPoint, + username, + password, + encryptedPassword, + consumerId, + consumerGroupId, + ownerId, + ownerEpoch, + thriftMaxFrameSize, + heartbeatIntervalMs, + connectionTimeoutInMs, + this::isClosed, + closedStatesDuringHandshake, + closedStatesDuringClose); + } + } + + private static class TestPullConsumer extends AbstractSubscriptionPullConsumer { + + private final List closedStatesDuringHandshake = new ArrayList<>(); + private final List closedStatesDuringClose = new ArrayList<>(); + + private TestPullConsumer() { + super( + new AbstractSubscriptionPullConsumerBuilder() + .host(HOST) + .port(PORT) + .consumerId(CONSUMER_ID) + .consumerGroupId(CONSUMER_GROUP_ID) + .heartbeatIntervalMs(LONG_INTERVAL_MS) + .endpointsSyncIntervalMs(LONG_INTERVAL_MS) + .autoCommit(false)); + } + + @Override + protected AbstractSubscriptionProvider constructSubscriptionProvider( + final TEndPoint endPoint, + final String username, + final String password, + final String encryptedPassword, + final String consumerId, + final String consumerGroupId, + final String ownerId, + final Long ownerEpoch, + final int thriftMaxFrameSize, + final long heartbeatIntervalMs, + final int connectionTimeoutInMs) { + return new TestSubscriptionProvider( + endPoint, + username, + password, + encryptedPassword, + consumerId, + consumerGroupId, + ownerId, + ownerEpoch, + thriftMaxFrameSize, + heartbeatIntervalMs, + connectionTimeoutInMs, + this::isClosed, + closedStatesDuringHandshake, + closedStatesDuringClose); + } + } + + private static class TestSubscriptionProvider extends AbstractSubscriptionProvider { + + private final BooleanSupplier consumerClosedSupplier; + private final List closedStatesDuringHandshake; + private final List closedStatesDuringClose; + + private TestSubscriptionProvider( + final TEndPoint endPoint, + final String username, + final String password, + final String encryptedPassword, + final String consumerId, + final String consumerGroupId, + final String ownerId, + final Long ownerEpoch, + final int thriftMaxFrameSize, + final long heartbeatIntervalMs, + final int connectionTimeoutInMs, + final BooleanSupplier consumerClosedSupplier, + final List closedStatesDuringHandshake, + final List closedStatesDuringClose) { + super( + endPoint, + username, + password, + encryptedPassword, + consumerId, + consumerGroupId, + ownerId, + ownerEpoch, + thriftMaxFrameSize, + heartbeatIntervalMs, + connectionTimeoutInMs); + this.consumerClosedSupplier = consumerClosedSupplier; + this.closedStatesDuringHandshake = closedStatesDuringHandshake; + this.closedStatesDuringClose = closedStatesDuringClose; + } + + @Override + protected AbstractSessionBuilder constructSubscriptionSessionBuilder( + final String host, + final int port, + final String username, + final String password, + final String encryptedPassword, + final int thriftMaxFrameSize, + final int connectionTimeoutInMs) { + final boolean useEncryptedPassword = Objects.nonNull(encryptedPassword); + return new SubscriptionTreeSessionBuilder() + .host(host) + .port(port) + .username(username) + .password(useEncryptedPassword ? encryptedPassword : password) + .useEncryptedPassword(useEncryptedPassword) + .thriftMaxFrameSize(thriftMaxFrameSize) + .connectionTimeoutInMs(connectionTimeoutInMs); + } + + @Override + synchronized void handshake() { + closedStatesDuringHandshake.add(consumerClosedSupplier.getAsBoolean()); + setAvailable(); + } + + @Override + synchronized void close() { + closedStatesDuringClose.add(consumerClosedSupplier.getAsBoolean()); + setUnavailable(); + } + + @Override + PipeSubscribeHeartbeatResp heartbeat( + final List processorBufferedCommitContexts) { + return new PipeSubscribeHeartbeatResp(); + } + } +}