From 55e66a8393d6d20e7931807e088390ab2ad45550 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:06:10 +0800 Subject: [PATCH] Fix subscription consumer lifecycle state ordering (#18151) --- iotdb-client/subscription/pom.xml | 5 + .../consumer/SubscriptionProviders.java | 6 + .../consumer/SubscriptionPullConsumer.java | 10 +- .../consumer/SubscriptionPushConsumer.java | 10 +- .../SubscriptionConsumerLifecycleTest.java | 214 ++++++++++++++++++ 5 files changed, 239 insertions(+), 6 deletions(-) create mode 100644 iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/SubscriptionConsumerLifecycleTest.java diff --git a/iotdb-client/subscription/pom.xml b/iotdb-client/subscription/pom.xml index 675e2678fa70d..c6cca865533b0 100644 --- a/iotdb-client/subscription/pom.xml +++ b/iotdb-client/subscription/pom.xml @@ -82,5 +82,10 @@ junit test + + org.mockito + mockito-core + test + diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionProviders.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionProviders.java index 6c0b7d03b184f..d67d355b24978 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionProviders.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionProviders.java @@ -235,6 +235,9 @@ void heartbeat(final SubscriptionConsumer consumer) { acquireWriteLock(); try { + if (consumer.isClosed()) { + return; + } heartbeatInternal(consumer); } finally { releaseWriteLock(); @@ -277,6 +280,9 @@ void sync(final SubscriptionConsumer 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/SubscriptionPullConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPullConsumer.java index 80eaf34b868a0..dbb79d9827a25 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPullConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPullConsumer.java @@ -114,10 +114,14 @@ public 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 @@ -138,8 +142,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/SubscriptionPushConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPushConsumer.java index fd28e81660960..c1f7283940777 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPushConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionPushConsumer.java @@ -120,10 +120,14 @@ public 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 @@ -136,8 +140,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/SubscriptionConsumerLifecycleTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/SubscriptionConsumerLifecycleTest.java new file mode 100644 index 0000000000000..bf94129ad8062 --- /dev/null +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/SubscriptionConsumerLifecycleTest.java @@ -0,0 +1,214 @@ +/* + * 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; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException; +import org.apache.iotdb.rpc.subscription.exception.SubscriptionException; +import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq; +import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.subscription.SubscriptionSessionConnection; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +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; + private static final int THRIFT_MAX_FRAME_SIZE = 1024; + private static final int CONNECTION_TIMEOUT_MS = 1000; + + @Test + public void testPushConsumerIsOpenBeforeProviderHandshake() { + final TestPushConsumer consumer = new TestPushConsumer(); + + try { + consumer.open(); + Assert.fail("Expected provider creation failure."); + } catch (final SubscriptionException e) { + Assert.assertTrue(e instanceof SubscriptionConnectionException); + Assert.assertEquals(1, consumer.closedStatesDuringHandshake.size()); + Assert.assertFalse(consumer.closedStatesDuringHandshake.get(0)); + Assert.assertTrue(consumer.isClosed()); + } + } + + @Test + public void testPushConsumerIsClosedBeforeProviderClose() throws Exception { + final TestPushConsumer consumer = new TestPushConsumer(); + final List closedStatesDuringClose = new ArrayList<>(); + + prepareOpenedConsumerWithProvider( + consumer, SubscriptionPushConsumer.class, closedStatesDuringClose); + consumer.close(); + + Assert.assertEquals(1, closedStatesDuringClose.size()); + Assert.assertTrue(closedStatesDuringClose.get(0)); + } + + @Test + public void testPullConsumerIsOpenBeforeProviderHandshake() { + final TestPullConsumer consumer = new TestPullConsumer(); + + try { + consumer.open(); + Assert.fail("Expected provider creation failure."); + } catch (final SubscriptionException e) { + Assert.assertTrue(e instanceof SubscriptionConnectionException); + Assert.assertEquals(1, consumer.closedStatesDuringHandshake.size()); + Assert.assertFalse(consumer.closedStatesDuringHandshake.get(0)); + Assert.assertTrue(consumer.isClosed()); + } + } + + @Test + public void testPullConsumerIsClosedBeforeProviderClose() throws Exception { + final TestPullConsumer consumer = new TestPullConsumer(); + final List closedStatesDuringClose = new ArrayList<>(); + + prepareOpenedConsumerWithProvider( + consumer, SubscriptionPullConsumer.class, closedStatesDuringClose); + consumer.close(); + + Assert.assertEquals(1, closedStatesDuringClose.size()); + Assert.assertTrue(closedStatesDuringClose.get(0)); + } + + private static void prepareOpenedConsumerWithProvider( + final SubscriptionConsumer consumer, + final Class concreteConsumerClass, + final List closedStatesDuringClose) + throws Exception { + setAtomicBooleanField(concreteConsumerClass, consumer, "isClosed", false); + setAtomicBooleanField(SubscriptionConsumer.class, consumer, "isClosed", false); + + final SubscriptionProvider provider = + new SubscriptionProvider( + new TEndPoint(HOST, PORT), + "root", + "root", + null, + CONSUMER_ID, + CONSUMER_GROUP_ID, + THRIFT_MAX_FRAME_SIZE, + LONG_INTERVAL_MS, + CONNECTION_TIMEOUT_MS); + setAtomicBooleanField(SubscriptionProvider.class, provider, "isClosed", false); + + final SubscriptionSessionConnection connection = + Mockito.mock(SubscriptionSessionConnection.class); + Mockito.when(connection.pipeSubscribe(Mockito.any(TPipeSubscribeReq.class))) + .thenAnswer( + invocation -> { + closedStatesDuringClose.add(consumer.isClosed()); + return successResponse(); + }); + setField(Session.class, provider, "defaultSessionConnection", connection); + + getProviders(consumer).addProvider(0, provider); + } + + private static TPipeSubscribeResp successResponse() { + return new TPipeSubscribeResp().setStatus(new TSStatus(200)); + } + + private static SubscriptionProviders getProviders(final SubscriptionConsumer consumer) + throws Exception { + final Field field = SubscriptionConsumer.class.getDeclaredField("providers"); + field.setAccessible(true); + return (SubscriptionProviders) field.get(consumer); + } + + private static void setAtomicBooleanField( + final Class clazz, final Object target, final String fieldName, final boolean value) + throws Exception { + final Field field = clazz.getDeclaredField(fieldName); + field.setAccessible(true); + ((AtomicBoolean) field.get(target)).set(value); + } + + private static void setField( + final Class clazz, final Object target, final String fieldName, final Object value) + throws Exception { + final Field field = clazz.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(target, value); + } + + private static class TestPushConsumer extends SubscriptionPushConsumer { + + private final List closedStatesDuringHandshake = new ArrayList<>(); + + private TestPushConsumer() { + super( + new SubscriptionPushConsumer.Builder() + .host(HOST) + .port(PORT) + .consumerId(CONSUMER_ID) + .consumerGroupId(CONSUMER_GROUP_ID) + .heartbeatIntervalMs(LONG_INTERVAL_MS) + .endpointsSyncIntervalMs(LONG_INTERVAL_MS) + .autoPollIntervalMs(LONG_INTERVAL_MS)); + } + + @Override + SubscriptionProvider constructProviderAndHandshake(final TEndPoint endPoint) + throws SubscriptionException { + closedStatesDuringHandshake.add(isClosed()); + throw new SubscriptionException("intentional provider creation failure"); + } + } + + private static class TestPullConsumer extends SubscriptionPullConsumer { + + private final List closedStatesDuringHandshake = new ArrayList<>(); + + private TestPullConsumer() { + super( + new SubscriptionPullConsumer.Builder() + .host(HOST) + .port(PORT) + .consumerId(CONSUMER_ID) + .consumerGroupId(CONSUMER_GROUP_ID) + .heartbeatIntervalMs(LONG_INTERVAL_MS) + .endpointsSyncIntervalMs(LONG_INTERVAL_MS) + .autoCommit(false)); + } + + @Override + SubscriptionProvider constructProviderAndHandshake(final TEndPoint endPoint) + throws SubscriptionException { + closedStatesDuringHandshake.add(isClosed()); + throw new SubscriptionException("intentional provider creation failure"); + } + } +}