From 7801f53cd621788b4714b74836b87a9052c1ba13 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:47:20 +0800 Subject: [PATCH] Fix interrupted snapshot subscription delivery --- .../broker/SubscriptionPrefetchingQueue.java | 25 ++- .../subscription/event/SubscriptionEvent.java | 6 + ...scriptionPrefetchingQueueRecoveryTest.java | 178 ++++++++++++++++++ 3 files changed, 205 insertions(+), 4 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecoveryTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java index bd0b193c5b08c..cb45fc785ec1d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java @@ -316,7 +316,8 @@ public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer) return null; } - private SubscriptionEvent pollInternalV2(final String consumerId) { + private synchronized SubscriptionEvent pollInternalV2(final String consumerId) { + // Keep the queue-to-in-flight transition atomic with the snapshot completion check. if (prefetchingQueue.isEmpty()) { onEvent(); } @@ -825,7 +826,10 @@ private synchronized boolean tryCommitCurrentTerminateEvent() { return false; } - if (!prefetchingQueue.isEmpty() || !inFlightEvents.isEmpty()) { + // Requeue adds to prefetchingQueue before removing from inFlightEvents. Check in-flight first + // so a concurrent requeue cannot make both collections appear empty across that transition. + // Polling takes this monitor while moving events in the opposite direction. + if (!inFlightEvents.isEmpty() || !prefetchingQueue.isEmpty()) { return false; } @@ -881,6 +885,13 @@ private boolean refreshInFlightEventLeaseInternal( *

This is used when the server polled the event but cannot fit it in the current response. */ public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) { + return requeue(consumerId, commitContext, false); + } + + private boolean requeue( + final String consumerId, + final SubscriptionCommitContext commitContext, + final boolean resetResponse) { acquireReadLock(); try { if (isClosed()) { @@ -897,7 +908,13 @@ public boolean requeue(final String consumerId, final SubscriptionCommitContext ev.cleanUp(false); return null; } - ev.resetLastPolledTimestamp(); + if (resetResponse) { + // A new connection starts a fresh transfer and cannot interpret FILE_PIECE/FILE_SEAL + // as its initial poll response. Local response-size control must preserve that state. + ev.resetForRedelivery(); + } else { + ev.resetLastPolledTimestamp(); + } prefetchEvent(ev); requeued.set(true); return null; @@ -912,7 +929,7 @@ public int requeueInFlightEvents(final String consumerId) { int requeuedCount = 0; for (final Pair key : ImmutableSet.copyOf(inFlightEvents.keySet())) { - if (Objects.equals(consumerId, key.left) && requeue(consumerId, key.right)) { + if (Objects.equals(consumerId, key.left) && requeue(consumerId, key.right, true)) { requeuedCount++; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java index 0c5fd9327f24f..98273407a7218 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java @@ -327,6 +327,12 @@ public synchronized void nack() { } } + /** Rewinds an interrupted delivery without counting it as a poison-message retry. */ + public synchronized void resetForRedelivery() { + response.nack(); + lastPolledTimestamp.set(INVALID_TIMESTAMP); + } + /** Makes this event pollable again without treating local response-size control as a nack. */ public void resetLastPolledTimestamp() { lastPolledTimestamp.set(INVALID_TIMESTAMP); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecoveryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecoveryTest.java new file mode 100644 index 0000000000000..3e10aff883dbb --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecoveryTest.java @@ -0,0 +1,178 @@ +/* + * 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.db.subscription.broker; + +import org.apache.iotdb.db.pipe.event.common.terminate.PipeTerminateEvent; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.db.subscription.event.pipe.SubscriptionPipeEvents; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType; +import org.apache.iotdb.rpc.subscription.payload.poll.TerminationPayload; +import org.apache.iotdb.session.subscription.util.PollTimer; + +import org.junit.Test; + +import java.lang.reflect.Field; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class SubscriptionPrefetchingQueueRecoveryTest { + + @Test + public void testDisconnectRestartsTsFileWithoutPoisoning() throws Exception { + final SubscriptionPrefetchingQueue queue = newQueue(); + final Path file = Files.createTempFile("subscription-recovery-", ".tsfile"); + try { + Files.write(file, new byte[] {1, 2, 3, 4}); + final SubscriptionEvent event = + new SubscriptionEvent( + mock(SubscriptionPipeEvents.class), file.toFile(), null, newContext()); + queue.prefetchEvent(event); + assertSame(event, poll(queue)); + event.fetchNextResponse(Files.size(file)); + assertEquals( + SubscriptionPollResponseType.FILE_SEAL.getType(), + event.getCurrentResponse().getResponseType()); + + // Response-size rejection retries the same response on the same transfer. + assertTrue(queue.requeue("consumer", event.getCommitContext())); + assertSame(event, poll(queue)); + assertEquals( + SubscriptionPollResponseType.FILE_SEAL.getType(), + event.getCurrentResponse().getResponseType()); + + // Disconnects restart the protocol, even after enough retries to poison a nacked event. + for (int retry = 0; retry < 12; retry++) { + assertEquals(1, queue.requeueInFlightEvents("consumer")); + assertSame(event, poll(queue)); + assertEquals( + SubscriptionPollResponseType.FILE_INIT.getType(), + event.getCurrentResponse().getResponseType()); + assertFalse(event.isCommittable()); + assertEquals(0L, event.getNackCount()); + event.fetchNextResponse(Files.size(file)); + } + assertTrue(queue.ack("consumer", event.getCommitContext())); + assertEquals(0L, queue.getSubscriptionUncommittedEventCount()); + } finally { + queue.cleanUp(); + Files.deleteIfExists(file); + } + } + + @Test + public void testSnapshotCannotCompleteDuringPollV2Transfer() throws Exception { + final SubscriptionPrefetchingQueue queue = newQueue(); + final CountDownLatch transferStarted = new CountDownLatch(1); + final CountDownLatch resumeTransfer = new CountDownLatch(1); + final CountDownLatch prefetchStarted = new CountDownLatch(1); + final CountDownLatch completion = new CountDownLatch(1); + final ExecutorService executor = Executors.newFixedThreadPool(2); + try { + final SubscriptionEvent event = + new SubscriptionEvent( + SubscriptionPollResponseType.TERMINATION.getType(), + new TerminationPayload(), + newContext()) { + @Override + public void recordLastPolledTimestamp() { + transferStarted.countDown(); + try { + assertTrue(resumeTransfer.await(5, TimeUnit.SECONDS)); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + super.recordLastPolledTimestamp(); + } + }; + queue.prefetchEvent(event); + final PipeTerminateEvent terminate = mock(PipeTerminateEvent.class); + doAnswer( + invocation -> { + completion.countDown(); + return null; + }) + .when(terminate) + .addOnCommittedHook(any(Runnable.class)); + final Field terminateField = + SubscriptionPrefetchingQueue.class.getDeclaredField("currentTerminateEvent"); + terminateField.setAccessible(true); + terminateField.set(queue, terminate); + + final Future polled = executor.submit(() -> poll(queue)); + assertTrue(transferStarted.await(5, TimeUnit.SECONDS)); + // The event has left the prefetch queue but is not yet in the in-flight map. + final Future prefetched = + executor.submit( + () -> { + prefetchStarted.countDown(); + queue.executePrefetchV2(); + }); + assertTrue(prefetchStarted.await(5, TimeUnit.SECONDS)); + assertFalse( + "Snapshot completion must wait for the unacknowledged delivery", + completion.await(200, TimeUnit.MILLISECONDS)); + resumeTransfer.countDown(); + assertSame(event, polled.get(5, TimeUnit.SECONDS)); + prefetched.get(5, TimeUnit.SECONDS); + assertEquals(1L, completion.getCount()); + assertTrue(queue.ack("consumer", event.getCommitContext())); + queue.executePrefetchV2(); + assertEquals(0L, completion.getCount()); + } finally { + resumeTransfer.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + queue.cleanUp(); + } + } + + private static SubscriptionPrefetchingQueue newQueue() { + final SubscriptionBlockingPendingQueue inputQueue = + mock(SubscriptionBlockingPendingQueue.class); + when(inputQueue.isEmpty()).thenReturn(true); + return new SubscriptionPrefetchingTabletQueue( + "recovery_group", "recovery_topic", inputQueue, new AtomicLong()); + } + + private static SubscriptionCommitContext newContext() { + return new SubscriptionCommitContext(1, 1, "recovery_topic", "recovery_group", 1); + } + + private static SubscriptionEvent poll(final SubscriptionPrefetchingQueue queue) { + return queue.pollV2("consumer", new PollTimer(System.currentTimeMillis(), 0L)); + } +}