From d271d4aa8c06dd9ed55e3966d4d4c84710368b8b Mon Sep 17 00:00:00 2001 From: Andrii Landiak Date: Thu, 11 Jun 2026 14:56:10 +0300 Subject: [PATCH 1/4] Add an opt-in readinessCheck gate so the edge consumer pauses polling instead of polling and dropping events while not ready (sync/high-priority/disconnected). --- .../edge/rpc/KafkaEdgeGrpcSession.java | 11 +- .../common/consumer/QueueConsumerManager.java | 31 ++- .../consumer/QueueConsumerManagerTest.java | 247 ++++++++++++++++++ 3 files changed, 286 insertions(+), 3 deletions(-) create mode 100644 common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java diff --git a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java index f7e729c6b7..082452e944 100644 --- a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java +++ b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java @@ -70,7 +70,9 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { private void processMsgs(List> msgs, TbQueueConsumer> consumer) { log.trace("[{}][{}] starting processing edge events", tenantId, edge.getId()); - if (!isConnected() || isSyncInProgress() || isHighPriorityProcessing) { + // Defensive backstop: the loop already gates polling on readiness; this only fires on the narrow race + // where readiness flips during poll(), and that already-polled batch is dropped here (can't rewind). + if (!isReadyToProcessGeneralEvents()) { log.debug("[{}][{}] edge not connected, edge sync is not completed or high priority processing in progress, " + "connected = {}, sync in progress = {}, high priority in progress = {}. Skipping iteration", tenantId, edge.getId(), isConnected(), isSyncInProgress(), isHighPriorityProcessing); @@ -96,6 +98,10 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { } } + private boolean isReadyToProcessGeneralEvents() { + return isConnected() && !isSyncInProgress() && !isHighPriorityProcessing; + } + @Override public ListenableFuture migrateEdgeEvents() throws Exception { return super.processEdgeEvents(); @@ -103,7 +109,7 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { @Override public ListenableFuture processEdgeEvents() { - if (!isConnected() || isSyncInProgress() || isHighPriorityProcessing) { + if (!isReadyToProcessGeneralEvents()) { log.warn("[{}][{}] Session is not ready (connected={}, syncInProgress={}, highPriority={}), skip starting edge event consumer", tenantId, edge != null ? edge.getId() : null, isConnected(), isSyncInProgress(), isHighPriorityProcessing); return Futures.immediateFuture(Boolean.FALSE); @@ -126,6 +132,7 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { .consumerCreator(() -> tbCoreQueueFactory.createEdgeEventMsgConsumer(tenantId, edge.getId())) .consumerExecutor(consumerExecutor) .threadPrefix("edge-events-" + edge.getId()) + .readinessCheck(this::isReadyToProcessGeneralEvents) .build(); consumer.subscribe(); consumer.launch(); diff --git a/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java b/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java index 4adc0354c4..9ee6bc0b50 100644 --- a/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java +++ b/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java @@ -30,6 +30,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.BooleanSupplier; import java.util.function.Supplier; @Slf4j @@ -40,6 +41,8 @@ public class QueueConsumerManager { private final long pollInterval; private final ExecutorService consumerExecutor; private final String threadPrefix; + /** Optional poll gate: while {@code false} the loop skips polling so the position doesn't advance; {@code null} = always ready (default). */ + private final BooleanSupplier readinessCheck; @Getter private final TbQueueConsumer consumer; @@ -49,12 +52,13 @@ public class QueueConsumerManager { @Builder public QueueConsumerManager(String name, MsgPackProcessor msgPackProcessor, long pollInterval, Supplier> consumerCreator, - ExecutorService consumerExecutor, String threadPrefix) { + ExecutorService consumerExecutor, String threadPrefix, BooleanSupplier readinessCheck) { this.name = name; this.pollInterval = pollInterval; this.msgPackProcessor = msgPackProcessor; this.consumerExecutor = consumerExecutor; this.threadPrefix = threadPrefix; + this.readinessCheck = readinessCheck; this.consumer = consumerCreator.get(); } @@ -84,6 +88,12 @@ public class QueueConsumerManager { private void consumerLoop(TbQueueConsumer consumer) { while (!stopped && !consumer.isStopped()) { try { + if (!isReadyToProcess()) { + if (!awaitNextReadinessCheck()) { + return; + } + continue; + } List msgs = consumer.poll(pollInterval); if (msgs.isEmpty()) { continue; @@ -102,6 +112,25 @@ public class QueueConsumerManager { } } + private boolean isReadyToProcess() { + return readinessCheck == null || readinessCheck.getAsBoolean(); + } + + /** + * Waits one poll interval before readiness is re-checked. Returns {@code false} if interrupted, which is treated as + * a stop signal so the consumer loop exits. + */ + private boolean awaitNextReadinessCheck() { + log.trace("[{}] Consumer is not ready to process messages yet, skipping poll iteration", name); + try { + Thread.sleep(pollInterval); + return true; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + } + public void stop() { log.debug("[{}] Stopping consumer", name); stopped = true; diff --git a/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java b/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java new file mode 100644 index 0000000000..6cf2d8bfe0 --- /dev/null +++ b/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java @@ -0,0 +1,247 @@ +/** + * Copyright © 2016-2026 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.queue.common.consumer; + +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; +import org.thingsboard.server.queue.TbQueueConsumer; +import org.thingsboard.server.queue.TbQueueMsg; + +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.Mockito.mock; + +@Slf4j +class QueueConsumerManagerTest { + + private static final long POLL_INTERVAL_MS = 20L; + // Before asserting the consumer never polled we wait until the loop has evaluated the readiness gate at least + // this many times. That proves the consumer thread is actually running and deciding not to poll, rather than the + // assertion passing vacuously because the thread simply has not started yet. + private static final int MIN_READINESS_CHECKS = 3; + + private final AtomicBoolean readyToProcess = new AtomicBoolean(false); + private final AtomicInteger readinessChecks = new AtomicInteger(); + private final TestQueueConsumer consumer = new TestQueueConsumer(); + + private ExecutorService consumerExecutor; + private QueueConsumerManager manager; + + @AfterEach + void tearDown() { + if (manager != null) { + manager.stop(); + } + if (consumerExecutor != null) { + consumerExecutor.shutdownNow(); + } + } + + @Test + void eventQueuedWhileNotReadyIsDeliveredAfterReadinessGateOpensInsteadOfBeingDropped() { + List delivered = new CopyOnWriteArrayList<>(); + + consumer.enqueue(List.of(mock(TbQueueMsg.class))); + + // The processor is unconditional: only the readiness gate may hold the event back, so delivery proves the + // gate (not the processor) is what kept the event queued while not ready. + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> { + delivered.addAll(msgs); + c.commit(); + }); + + // The loop is running and repeatedly evaluating the gate during the not-ready (sync) window... + awaitReadinessGateEvaluated(readinessChecks); + // ...yet the queued event is neither polled nor delivered - it stays in the queue rather than being dropped. + assertThat(consumer.getPollCount()) + .as("consumer must not poll while not ready") + .isZero(); + assertThat(delivered) + .as("event must not be delivered while not ready") + .isEmpty(); + + // Sync completes - the processor becomes ready. + readyToProcess.set(true); + + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(delivered) + .as("event queued during the not-ready window must be delivered, not dropped") + .hasSize(1)); + } + + @Test + void consumerIsNotPolledWhileNotReadyToProcess() { + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> c.commit()); + + awaitReadinessGateEvaluated(readinessChecks); + assertThat(consumer.getPollCount()) + .as("consumer must not be polled while not ready to process") + .isZero(); + + readyToProcess.set(true); + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(consumer.getPollCount()) + .as("consumer resumes polling once ready") + .isPositive()); + } + + @Test + void consumerWithoutReadinessCheckPollsAndDeliversImmediately() { + List delivered = new CopyOnWriteArrayList<>(); + + consumer.enqueue(List.of(mock(TbQueueMsg.class))); + + // No readiness gate configured - the consumer must default to "always ready", preserving the behaviour every + // consumer that does not opt in relies on. + manager = launchManager(consumer, null, (msgs, c) -> { + delivered.addAll(msgs); + c.commit(); + }); + + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(delivered) + .as("consumer without a readiness gate must poll and deliver immediately") + .hasSize(1)); + } + + @Test + void consumerLoopExitsWhenInterruptedWhileNotReady() throws Exception { + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> c.commit()); + + // The loop is parked in the not-ready wait... + awaitReadinessGateEvaluated(readinessChecks); + + // ...interrupting the worker (as shutdownNow does on stop) must end the loop, not spin or hang. + consumerExecutor.shutdownNow(); + assertThat(consumerExecutor.awaitTermination(5, TimeUnit.SECONDS)) + .as("consumer loop must exit when interrupted while waiting to become ready") + .isTrue(); + } + + private static void awaitReadinessGateEvaluated(AtomicInteger readinessChecks) { + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(readinessChecks.get()) + .as("consumer loop must be running and repeatedly evaluating the readiness gate") + .isGreaterThanOrEqualTo(MIN_READINESS_CHECKS)); + } + + private static BooleanSupplier countingReadiness(AtomicBoolean ready, AtomicInteger readinessChecks) { + return () -> { + readinessChecks.incrementAndGet(); + return ready.get(); + }; + } + + private QueueConsumerManager launchManager(TestQueueConsumer consumer, BooleanSupplier readinessCheck, + QueueConsumerManager.MsgPackProcessor processor) { + consumerExecutor = Executors.newSingleThreadExecutor(); + QueueConsumerManager queueConsumerManager = QueueConsumerManager.builder() + .name("test-consumer") + .pollInterval(POLL_INTERVAL_MS) + .consumerCreator(() -> consumer) + .consumerExecutor(consumerExecutor) + .readinessCheck(readinessCheck) + .msgPackProcessor(processor) + .build(); + queueConsumerManager.subscribe(); + queueConsumerManager.launch(); + return queueConsumerManager; + } + + private static class TestQueueConsumer implements TbQueueConsumer { + + private final Queue> batches = new ConcurrentLinkedQueue<>(); + private final AtomicInteger pollCount = new AtomicInteger(); + private volatile boolean stopped; + + void enqueue(List batch) { + batches.add(batch); + } + + int getPollCount() { + return pollCount.get(); + } + + @Override + public List poll(long durationInMillis) { + pollCount.incrementAndGet(); + List batch = batches.poll(); + if (batch != null) { + return batch; + } + try { + Thread.sleep(durationInMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return Collections.emptyList(); + } + + @Override + public String getTopic() { + return "test-topic"; + } + + @Override + public void subscribe() { + } + + @Override + public void subscribe(Set partitions) { + } + + @Override + public void stop() { + stopped = true; + } + + @Override + public void unsubscribe() { + stopped = true; + } + + @Override + public void commit() { + } + + @Override + public boolean isStopped() { + return stopped; + } + + @Override + public List getFullTopicNames() { + return Collections.emptyList(); + } + + } + +} From 2186aebf90e04547d61e8388e0d60ec9e6c5ecd9 Mon Sep 17 00:00:00 2001 From: Andrii Landiak Date: Thu, 11 Jun 2026 15:43:30 +0300 Subject: [PATCH 2/4] KafkaEdgeGrpcSessionTest: added --- .../edge/rpc/KafkaEdgeGrpcSession.java | 9 +- .../edge/rpc/KafkaEdgeGrpcSessionTest.java | 294 ++++++++++++++++++ 2 files changed, 300 insertions(+), 3 deletions(-) create mode 100644 application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java diff --git a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java index 082452e944..67c0a19e66 100644 --- a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java +++ b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java @@ -74,7 +74,7 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { // where readiness flips during poll(), and that already-polled batch is dropped here (can't rewind). if (!isReadyToProcessGeneralEvents()) { log.debug("[{}][{}] edge not connected, edge sync is not completed or high priority processing in progress, " + - "connected = {}, sync in progress = {}, high priority in progress = {}. Skipping iteration", + "connected = {}, sync in progress = {}, high priority in progress = {}. Skipping iteration", tenantId, edge.getId(), isConnected(), isSyncInProgress(), isHighPriorityProcessing); return; } @@ -147,8 +147,11 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { @Override public void processHighPriorityEvents() { isHighPriorityProcessing = true; - super.processHighPriorityEvents(); - isHighPriorityProcessing = false; + try { + super.processHighPriorityEvents(); + } finally { + isHighPriorityProcessing = false; + } } @Override diff --git a/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java b/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java new file mode 100644 index 0000000000..98931df813 --- /dev/null +++ b/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java @@ -0,0 +1,294 @@ +/** + * Copyright © 2016-2026 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.service.edge.rpc; + +import io.grpc.stub.StreamObserver; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.test.util.ReflectionTestUtils; +import org.thingsboard.server.common.data.edge.Edge; +import org.thingsboard.server.common.data.id.EdgeId; +import org.thingsboard.server.common.data.id.TenantId; +import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; +import org.thingsboard.server.gen.edge.v1.ResponseMsg; +import org.thingsboard.server.gen.transport.TransportProtos.ToEdgeEventNotificationMsg; +import org.thingsboard.server.queue.TbQueueConsumer; +import org.thingsboard.server.queue.common.TbProtoQueueMsg; +import org.thingsboard.server.queue.common.consumer.QueueConsumerManager; +import org.thingsboard.server.queue.discovery.TopicService; +import org.thingsboard.server.queue.kafka.KafkaAdmin; +import org.thingsboard.server.queue.provider.TbCoreQueueFactory; +import org.thingsboard.server.service.edge.EdgeContextComponent; + +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class KafkaEdgeGrpcSessionTest { + + private static final long POLL_INTERVAL_MS = 20L; + + private EdgeContextComponent ctx; + private TbCoreQueueFactory tbCoreQueueFactory; + private KafkaEdgeGrpcSession session; + + @BeforeEach + void setUp() { + ctx = mock(EdgeContextComponent.class); + TopicService topicService = mock(TopicService.class); + tbCoreQueueFactory = mock(TbCoreQueueFactory.class); + KafkaAdmin kafkaAdmin = mock(KafkaAdmin.class); + @SuppressWarnings("unchecked") + StreamObserver outputStream = mock(StreamObserver.class); + + session = new KafkaEdgeGrpcSession(ctx, topicService, tbCoreQueueFactory, kafkaAdmin, outputStream, + (edgeId, s) -> {}, (edge, uuid) -> {}, null, 0, 0); + + ReflectionTestUtils.setField(session, "edge", new Edge(new EdgeId(UUID.randomUUID()))); + ReflectionTestUtils.setField(session, "tenantId", TenantId.fromUUID(UUID.randomUUID())); + } + + @AfterEach + void tearDown() { + if (session != null) { + session.destroy(); + } + } + + @Test + void readyOnlyWhenConnectedNotSyncingNotHighPriority() { + setReadinessFlags(true, false, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("connected, not syncing, no high-priority work -> ready") + .isTrue(); + } + + @Test + void notReadyWhenDisconnected() { + setReadinessFlags(false, false, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("disconnected -> not ready") + .isFalse(); + } + + @Test + void notReadyWhileSyncInProgress() { + setReadinessFlags(true, true, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("sync in progress -> not ready (this is the window where events were being dropped)") + .isFalse(); + } + + @Test + void notReadyWhileHighPriorityProcessing() { + setReadinessFlags(true, false, true); + assertThat(isReadyToProcessGeneralEvents()) + .as("high-priority processing -> not ready") + .isFalse(); + } + + @Test + void processEdgeEventsWiresReadinessPredicateIntoConsumerGate() { + // processEdgeEvents() builds the consumer lazily; stub just enough of that path. + EdgeEventStorageSettings storageSettings = new EdgeEventStorageSettings(); + storageSettings.setNoRecordsSleepInterval(1000L); + when(ctx.getEdgeEventStorageSettings()).thenReturn(storageSettings); + + @SuppressWarnings("unchecked") + TbQueueConsumer> queueConsumer = mock(TbQueueConsumer.class); + // Report stopped so the launched consumer loop exits immediately - this test asserts on wiring, not polling. + when(queueConsumer.isStopped()).thenReturn(true); + when(tbCoreQueueFactory.createEdgeEventMsgConsumer(any(), any())).thenReturn(queueConsumer); + + // The consumer is only started when the session is ready. + setReadinessFlags(true, false, false); + session.processEdgeEvents(); + + QueueConsumerManager> manager = session.getConsumer(); + assertThat(manager).as("processEdgeEvents must build the consumer when ready").isNotNull(); + + BooleanSupplier readinessCheck = (BooleanSupplier) ReflectionTestUtils.getField(manager, "readinessCheck"); + assertThat(readinessCheck) + .as("the edge consumer must be wired with a readinessCheck (the .readinessCheck(...) builder line)") + .isNotNull(); + + // It must be the live predicate, not a snapshot: flipping the session's state must flip the gate. + assertThat(readinessCheck.getAsBoolean()).as("ready session -> gate open").isTrue(); + setReadinessFlags(true, true, false); + assertThat(readinessCheck.getAsBoolean()).as("sync starts -> gate closes, consumer pauses polling").isFalse(); + } + + @Test + void eventArrivingDuringSyncIsHeldByTheEdgeConsumerUntilSyncCompletes() { + EdgeEventStorageSettings storageSettings = new EdgeEventStorageSettings(); + storageSettings.setNoRecordsSleepInterval(POLL_INTERVAL_MS); + when(ctx.getEdgeEventStorageSettings()).thenReturn(storageSettings); + + RecordingEdgeEventConsumer queueConsumer = new RecordingEdgeEventConsumer(); + when(tbCoreQueueFactory.createEdgeEventMsgConsumer(any(), any())).thenReturn(queueConsumer); + + // The consumer is launched only while the session is ready - that is how it starts in production. + setReadinessFlags(true, false, false); + session.processEdgeEvents(); + + // Sync starts: the gate closes. Wait until the loop has actually parked on it (poll count stops advancing) + // before enqueuing - otherwise we would race an in-flight poll() and the test would be non-deterministic. + setReadinessFlags(true, true, false); + awaitParkedOnClosedGate(queueConsumer); + + // An event lands in the edge-event topic during the sync window - exactly the case that used to be dropped. + @SuppressWarnings("unchecked") + TbProtoQueueMsg event = mock(TbProtoQueueMsg.class); + int pollsBeforeEvent = queueConsumer.getPollCount(); + queueConsumer.enqueue(List.of(event)); + + // While sync is in progress the consumer stays parked: it neither polls nor consumes the event, + // so the event survives in the queue instead of being read-and-skipped. + sleepQuietly(POLL_INTERVAL_MS * 5); + assertThat(queueConsumer.getPolledEvents()) + .as("event must not be polled while sync is in progress (it must stay queued, not be dropped)") + .isEmpty(); + assertThat(queueConsumer.getPollCount()) + .as("consumer must not poll at all while the gate is closed") + .isEqualTo(pollsBeforeEvent); + + // Sync completes: the gate opens and the held event is finally picked up by the consumer. + // (We assert at the poll boundary - the actual drop point - since the downlink-send path that + // processMsgs drives afterwards is not reachable from a unit test.) + setReadinessFlags(true, false, false); + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(queueConsumer.getPolledEvents()) + .as("event held during sync must be picked up once sync completes, not lost") + .hasSize(1)); + } + + private boolean isReadyToProcessGeneralEvents() { + return Boolean.TRUE.equals(ReflectionTestUtils.invokeMethod(session, "isReadyToProcessGeneralEvents")); + } + + private void setReadinessFlags(boolean connected, boolean syncInProgress, boolean highPriorityProcessing) { + ReflectionTestUtils.setField(session, "connected", connected); + ReflectionTestUtils.setField(session, "syncInProgress", syncInProgress); + ReflectionTestUtils.setField(session, "isHighPriorityProcessing", highPriorityProcessing); + } + + private static void awaitParkedOnClosedGate(RecordingEdgeEventConsumer consumer) { + await().atMost(5, TimeUnit.SECONDS).until(() -> { + int before = consumer.getPollCount(); + // Several poll intervals with no new poll means the loop is parked on the readiness gate rather than + // blocked inside poll(), so it is safe to enqueue without racing an in-flight read. + sleepQuietly(POLL_INTERVAL_MS * 5); + return consumer.getPollCount() == before; + }); + } + + private static void sleepQuietly(long millis) { + try { + Thread.sleep(millis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + /** + * Fake edge-event consumer that records what it was actually polled, so a test can prove an event stayed in the + * queue while the readiness gate was closed and was only read once it opened. + */ + private static class RecordingEdgeEventConsumer implements TbQueueConsumer> { + + private final Queue>> pending = new ConcurrentLinkedQueue<>(); + private final List> polledEvents = new CopyOnWriteArrayList<>(); + private final AtomicInteger pollCount = new AtomicInteger(); + private volatile boolean stopped; + + void enqueue(List> batch) { + pending.add(batch); + } + + int getPollCount() { + return pollCount.get(); + } + + List> getPolledEvents() { + return polledEvents; + } + + @Override + public List> poll(long durationInMillis) { + pollCount.incrementAndGet(); + List> batch = pending.poll(); + if (batch != null) { + polledEvents.addAll(batch); + return batch; + } + sleepQuietly(durationInMillis); + return Collections.emptyList(); + } + + @Override + public String getTopic() { + return "test-edge-event-topic"; + } + + @Override + public void subscribe() { + } + + @Override + public void subscribe(Set partitions) { + } + + @Override + public void stop() { + stopped = true; + } + + @Override + public void unsubscribe() { + stopped = true; + } + + @Override + public void commit() { + } + + @Override + public boolean isStopped() { + return stopped; + } + + @Override + public List getFullTopicNames() { + return Collections.emptyList(); + } + + } + +} From ad270c81cb0410ba2c8aaf792936422ba0f04f64 Mon Sep 17 00:00:00 2001 From: Andrii Landiak Date: Thu, 11 Jun 2026 16:07:25 +0300 Subject: [PATCH 3/4] Add an opt-in readinessCheck gate so the edge consumer pauses polling instead of polling and dropping events while not ready (sync/high-priority/disconnected). --- .../edge/rpc/KafkaEdgeGrpcSession.java | 11 +- .../common/consumer/QueueConsumerManager.java | 31 ++- .../consumer/QueueConsumerManagerTest.java | 251 ++++++++++++++++++ 3 files changed, 290 insertions(+), 3 deletions(-) create mode 100644 common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java diff --git a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java index f7e729c6b7..082452e944 100644 --- a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java +++ b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java @@ -70,7 +70,9 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { private void processMsgs(List> msgs, TbQueueConsumer> consumer) { log.trace("[{}][{}] starting processing edge events", tenantId, edge.getId()); - if (!isConnected() || isSyncInProgress() || isHighPriorityProcessing) { + // Defensive backstop: the loop already gates polling on readiness; this only fires on the narrow race + // where readiness flips during poll(), and that already-polled batch is dropped here (can't rewind). + if (!isReadyToProcessGeneralEvents()) { log.debug("[{}][{}] edge not connected, edge sync is not completed or high priority processing in progress, " + "connected = {}, sync in progress = {}, high priority in progress = {}. Skipping iteration", tenantId, edge.getId(), isConnected(), isSyncInProgress(), isHighPriorityProcessing); @@ -96,6 +98,10 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { } } + private boolean isReadyToProcessGeneralEvents() { + return isConnected() && !isSyncInProgress() && !isHighPriorityProcessing; + } + @Override public ListenableFuture migrateEdgeEvents() throws Exception { return super.processEdgeEvents(); @@ -103,7 +109,7 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { @Override public ListenableFuture processEdgeEvents() { - if (!isConnected() || isSyncInProgress() || isHighPriorityProcessing) { + if (!isReadyToProcessGeneralEvents()) { log.warn("[{}][{}] Session is not ready (connected={}, syncInProgress={}, highPriority={}), skip starting edge event consumer", tenantId, edge != null ? edge.getId() : null, isConnected(), isSyncInProgress(), isHighPriorityProcessing); return Futures.immediateFuture(Boolean.FALSE); @@ -126,6 +132,7 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { .consumerCreator(() -> tbCoreQueueFactory.createEdgeEventMsgConsumer(tenantId, edge.getId())) .consumerExecutor(consumerExecutor) .threadPrefix("edge-events-" + edge.getId()) + .readinessCheck(this::isReadyToProcessGeneralEvents) .build(); consumer.subscribe(); consumer.launch(); diff --git a/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java b/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java index 4adc0354c4..9ee6bc0b50 100644 --- a/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java +++ b/common/queue/src/main/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManager.java @@ -30,6 +30,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.BooleanSupplier; import java.util.function.Supplier; @Slf4j @@ -40,6 +41,8 @@ public class QueueConsumerManager { private final long pollInterval; private final ExecutorService consumerExecutor; private final String threadPrefix; + /** Optional poll gate: while {@code false} the loop skips polling so the position doesn't advance; {@code null} = always ready (default). */ + private final BooleanSupplier readinessCheck; @Getter private final TbQueueConsumer consumer; @@ -49,12 +52,13 @@ public class QueueConsumerManager { @Builder public QueueConsumerManager(String name, MsgPackProcessor msgPackProcessor, long pollInterval, Supplier> consumerCreator, - ExecutorService consumerExecutor, String threadPrefix) { + ExecutorService consumerExecutor, String threadPrefix, BooleanSupplier readinessCheck) { this.name = name; this.pollInterval = pollInterval; this.msgPackProcessor = msgPackProcessor; this.consumerExecutor = consumerExecutor; this.threadPrefix = threadPrefix; + this.readinessCheck = readinessCheck; this.consumer = consumerCreator.get(); } @@ -84,6 +88,12 @@ public class QueueConsumerManager { private void consumerLoop(TbQueueConsumer consumer) { while (!stopped && !consumer.isStopped()) { try { + if (!isReadyToProcess()) { + if (!awaitNextReadinessCheck()) { + return; + } + continue; + } List msgs = consumer.poll(pollInterval); if (msgs.isEmpty()) { continue; @@ -102,6 +112,25 @@ public class QueueConsumerManager { } } + private boolean isReadyToProcess() { + return readinessCheck == null || readinessCheck.getAsBoolean(); + } + + /** + * Waits one poll interval before readiness is re-checked. Returns {@code false} if interrupted, which is treated as + * a stop signal so the consumer loop exits. + */ + private boolean awaitNextReadinessCheck() { + log.trace("[{}] Consumer is not ready to process messages yet, skipping poll iteration", name); + try { + Thread.sleep(pollInterval); + return true; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + } + public void stop() { log.debug("[{}] Stopping consumer", name); stopped = true; diff --git a/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java b/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java new file mode 100644 index 0000000000..c987fb502e --- /dev/null +++ b/common/queue/src/test/java/org/thingsboard/server/queue/common/consumer/QueueConsumerManagerTest.java @@ -0,0 +1,251 @@ +/** + * Copyright © 2016-2026 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.queue.common.consumer; + +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; +import org.thingsboard.server.queue.TbQueueConsumer; +import org.thingsboard.server.queue.TbQueueMsg; + +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.Mockito.mock; + +@Slf4j +class QueueConsumerManagerTest { + + private static final long POLL_INTERVAL_MS = 20L; + // Before asserting the consumer never polled we wait until the loop has evaluated the readiness gate at least + // this many times. That proves the consumer thread is actually running and deciding not to poll, rather than the + // assertion passing vacuously because the thread simply has not started yet. + private static final int MIN_READINESS_CHECKS = 3; + + private final AtomicBoolean readyToProcess = new AtomicBoolean(false); + private final AtomicInteger readinessChecks = new AtomicInteger(); + private final TestQueueConsumer consumer = new TestQueueConsumer(); + private ExecutorService consumerExecutor; + private QueueConsumerManager manager; + + @AfterEach + void tearDown() { + if (manager != null) { + manager.stop(); + } + if (consumerExecutor != null) { + consumerExecutor.shutdownNow(); + } + } + + @Test + void eventQueuedWhileNotReadyIsDeliveredAfterReadinessGateOpensInsteadOfBeingDropped() { + List delivered = new CopyOnWriteArrayList<>(); + + consumer.enqueue(List.of(mock(TbQueueMsg.class))); + + // The processor is unconditional: only the readiness gate may hold the event back, so delivery proves the + // gate (not the processor) is what kept the event queued while not ready. + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> { + delivered.addAll(msgs); + c.commit(); + }); + + // The loop is running and repeatedly evaluating the gate during the not-ready (sync) window... + awaitReadinessGateEvaluated(readinessChecks); + // ...yet the queued event is neither polled nor delivered - it stays in the queue rather than being dropped. + assertThat(consumer.getPollCount()) + .as("consumer must not poll while not ready") + .isZero(); + assertThat(delivered) + .as("event must not be delivered while not ready") + .isEmpty(); + + // Sync completes - the processor becomes ready. + readyToProcess.set(true); + + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(delivered) + .as("event queued during the not-ready window must be delivered, not dropped") + .hasSize(1)); + } + + @Test + void consumerIsNotPolledWhileNotReadyToProcess() { + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> c.commit()); + + awaitReadinessGateEvaluated(readinessChecks); + assertThat(consumer.getPollCount()) + .as("consumer must not be polled while not ready to process") + .isZero(); + + readyToProcess.set(true); + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(consumer.getPollCount()) + .as("consumer resumes polling once ready") + .isPositive()); + } + + @Test + void consumerWithoutReadinessCheckPollsAndDeliversImmediately() { + List delivered = new CopyOnWriteArrayList<>(); + + consumer.enqueue(List.of(mock(TbQueueMsg.class))); + + // No readiness gate configured - the consumer must default to "always ready", preserving the behaviour every + // consumer that does not opt in relies on. + manager = launchManager(consumer, null, (msgs, c) -> { + delivered.addAll(msgs); + c.commit(); + }); + + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(delivered) + .as("consumer without a readiness gate must poll and deliver immediately") + .hasSize(1)); + } + + @Test + void consumerLoopExitsWhenInterruptedWhileNotReady() throws Exception { + manager = launchManager(consumer, countingReadiness(readyToProcess, readinessChecks), (msgs, c) -> c.commit()); + + // The loop is parked in the not-ready wait... + awaitReadinessGateEvaluated(readinessChecks); + + // ...interrupting the worker (as shutdownNow does on stop) must end the loop, not spin or hang. + consumerExecutor.shutdownNow(); + assertThat(consumerExecutor.awaitTermination(5, TimeUnit.SECONDS)) + .as("consumer loop must exit when interrupted while waiting to become ready") + .isTrue(); + } + + private static void awaitReadinessGateEvaluated(AtomicInteger readinessChecks) { + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(readinessChecks.get()) + .as("consumer loop must be running and repeatedly evaluating the readiness gate") + .isGreaterThanOrEqualTo(MIN_READINESS_CHECKS)); + } + + private static BooleanSupplier countingReadiness(AtomicBoolean ready, AtomicInteger readinessChecks) { + return () -> { + readinessChecks.incrementAndGet(); + return ready.get(); + }; + } + + private QueueConsumerManager launchManager(TestQueueConsumer consumer, BooleanSupplier readinessCheck, + QueueConsumerManager.MsgPackProcessor processor) { + consumerExecutor = Executors.newSingleThreadExecutor(); + QueueConsumerManager queueConsumerManager = QueueConsumerManager.builder() + .name("test-consumer") + .pollInterval(POLL_INTERVAL_MS) + .consumerCreator(() -> consumer) + .consumerExecutor(consumerExecutor) + .readinessCheck(readinessCheck) + .msgPackProcessor(processor) + .build(); + queueConsumerManager.subscribe(); + queueConsumerManager.launch(); + return queueConsumerManager; + } + + private static class TestQueueConsumer implements TbQueueConsumer { + + private final Queue> batches = new ConcurrentLinkedQueue<>(); + private final AtomicInteger pollCount = new AtomicInteger(); + private volatile boolean stopped; + + void enqueue(List batch) { + batches.add(batch); + } + + int getPollCount() { + return pollCount.get(); + } + + @Override + public List poll(long durationInMillis) { + pollCount.incrementAndGet(); + List batch = batches.poll(); + if (batch != null) { + return batch; + } + try { + Thread.sleep(durationInMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return Collections.emptyList(); + } + + @Override + public String getTopic() { + return "test-topic"; + } + + @Override + public void subscribe() { + } + + @Override + public void subscribe(Set partitions) { + } + + @Override + public void stop() { + stopped = true; + } + + @Override + public void unsubscribe() { + stopped = true; + } + + @Override + public void commit() { + } + + @Override + public boolean isStopped() { + return stopped; + } + + @Override + public Set getPartitions() { + return Collections.emptySet(); + } + + @Override + public List getFullTopicNames() { + return Collections.emptyList(); + } + + } + +} From 48ff41b57b1d6861c67ee7e43a2af1ce298d028b Mon Sep 17 00:00:00 2001 From: Andrii Landiak Date: Thu, 11 Jun 2026 16:01:28 +0300 Subject: [PATCH 4/4] KafkaEdgeGrpcSessionTest: added --- .../edge/rpc/KafkaEdgeGrpcSession.java | 7 +- .../edge/rpc/KafkaEdgeGrpcSessionTest.java | 299 ++++++++++++++++++ 2 files changed, 304 insertions(+), 2 deletions(-) create mode 100644 application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java diff --git a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java index 082452e944..1f2d378d72 100644 --- a/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java +++ b/application/src/main/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSession.java @@ -147,8 +147,11 @@ public class KafkaEdgeGrpcSession extends EdgeGrpcSession { @Override public void processHighPriorityEvents() { isHighPriorityProcessing = true; - super.processHighPriorityEvents(); - isHighPriorityProcessing = false; + try { + super.processHighPriorityEvents(); + } finally { + isHighPriorityProcessing = false; + } } @Override diff --git a/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java b/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java new file mode 100644 index 0000000000..39ba71b927 --- /dev/null +++ b/application/src/test/java/org/thingsboard/server/service/edge/rpc/KafkaEdgeGrpcSessionTest.java @@ -0,0 +1,299 @@ +/** + * Copyright © 2016-2026 The Thingsboard Authors + * + * Licensed 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.thingsboard.server.service.edge.rpc; + +import io.grpc.stub.StreamObserver; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.test.util.ReflectionTestUtils; +import org.thingsboard.server.common.data.edge.Edge; +import org.thingsboard.server.common.data.id.EdgeId; +import org.thingsboard.server.common.data.id.TenantId; +import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; +import org.thingsboard.server.gen.edge.v1.ResponseMsg; +import org.thingsboard.server.gen.transport.TransportProtos.ToEdgeEventNotificationMsg; +import org.thingsboard.server.queue.TbQueueConsumer; +import org.thingsboard.server.queue.common.TbProtoQueueMsg; +import org.thingsboard.server.queue.common.consumer.QueueConsumerManager; +import org.thingsboard.server.queue.discovery.TopicService; +import org.thingsboard.server.queue.kafka.KafkaAdmin; +import org.thingsboard.server.queue.provider.TbCoreQueueFactory; +import org.thingsboard.server.service.edge.EdgeContextComponent; + +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class KafkaEdgeGrpcSessionTest { + + private static final long POLL_INTERVAL_MS = 20L; + + private EdgeContextComponent ctx; + private TbCoreQueueFactory tbCoreQueueFactory; + private KafkaEdgeGrpcSession session; + + @BeforeEach + void setUp() { + ctx = mock(EdgeContextComponent.class); + TopicService topicService = mock(TopicService.class); + tbCoreQueueFactory = mock(TbCoreQueueFactory.class); + KafkaAdmin kafkaAdmin = mock(KafkaAdmin.class); + @SuppressWarnings("unchecked") + StreamObserver outputStream = mock(StreamObserver.class); + + session = new KafkaEdgeGrpcSession(ctx, topicService, tbCoreQueueFactory, kafkaAdmin, outputStream, + (edgeId, s) -> {}, (edge, uuid) -> {}, null, 0, 0); + + ReflectionTestUtils.setField(session, "edge", new Edge(new EdgeId(UUID.randomUUID()))); + ReflectionTestUtils.setField(session, "tenantId", TenantId.fromUUID(UUID.randomUUID())); + } + + @AfterEach + void tearDown() { + if (session != null) { + session.destroy(); + } + } + + @Test + void readyOnlyWhenConnectedNotSyncingNotHighPriority() { + setReadinessFlags(true, false, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("connected, not syncing, no high-priority work -> ready") + .isTrue(); + } + + @Test + void notReadyWhenDisconnected() { + setReadinessFlags(false, false, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("disconnected -> not ready") + .isFalse(); + } + + @Test + void notReadyWhileSyncInProgress() { + setReadinessFlags(true, true, false); + assertThat(isReadyToProcessGeneralEvents()) + .as("sync in progress -> not ready (this is the window where events were being dropped)") + .isFalse(); + } + + @Test + void notReadyWhileHighPriorityProcessing() { + setReadinessFlags(true, false, true); + assertThat(isReadyToProcessGeneralEvents()) + .as("high-priority processing -> not ready") + .isFalse(); + } + + @Test + void processEdgeEventsWiresReadinessPredicateIntoConsumerGate() { + // processEdgeEvents() builds the consumer lazily; stub just enough of that path. + EdgeEventStorageSettings storageSettings = new EdgeEventStorageSettings(); + storageSettings.setNoRecordsSleepInterval(1000L); + when(ctx.getEdgeEventStorageSettings()).thenReturn(storageSettings); + + @SuppressWarnings("unchecked") + TbQueueConsumer> queueConsumer = mock(TbQueueConsumer.class); + // Report stopped so the launched consumer loop exits immediately - this test asserts on wiring, not polling. + when(queueConsumer.isStopped()).thenReturn(true); + when(tbCoreQueueFactory.createEdgeEventMsgConsumer(any(), any())).thenReturn(queueConsumer); + + // The consumer is only started when the session is ready. + setReadinessFlags(true, false, false); + session.processEdgeEvents(); + + QueueConsumerManager> manager = session.getConsumer(); + assertThat(manager).as("processEdgeEvents must build the consumer when ready").isNotNull(); + + BooleanSupplier readinessCheck = (BooleanSupplier) ReflectionTestUtils.getField(manager, "readinessCheck"); + assertThat(readinessCheck) + .as("the edge consumer must be wired with a readinessCheck (the .readinessCheck(...) builder line)") + .isNotNull(); + + // It must be the live predicate, not a snapshot: flipping the session's state must flip the gate. + assertThat(readinessCheck.getAsBoolean()).as("ready session -> gate open").isTrue(); + setReadinessFlags(true, true, false); + assertThat(readinessCheck.getAsBoolean()).as("sync starts -> gate closes, consumer pauses polling").isFalse(); + } + + @Test + void eventArrivingDuringSyncIsHeldByTheEdgeConsumerUntilSyncCompletes() { + EdgeEventStorageSettings storageSettings = new EdgeEventStorageSettings(); + storageSettings.setNoRecordsSleepInterval(POLL_INTERVAL_MS); + when(ctx.getEdgeEventStorageSettings()).thenReturn(storageSettings); + + RecordingEdgeEventConsumer queueConsumer = new RecordingEdgeEventConsumer(); + when(tbCoreQueueFactory.createEdgeEventMsgConsumer(any(), any())).thenReturn(queueConsumer); + + // The consumer is launched only while the session is ready - that is how it starts in production. + setReadinessFlags(true, false, false); + session.processEdgeEvents(); + + // Sync starts: the gate closes. Wait until the loop has actually parked on it (poll count stops advancing) + // before enqueuing - otherwise we would race an in-flight poll() and the test would be non-deterministic. + setReadinessFlags(true, true, false); + awaitParkedOnClosedGate(queueConsumer); + + // An event lands in the edge-event topic during the sync window - exactly the case that used to be dropped. + @SuppressWarnings("unchecked") + TbProtoQueueMsg event = mock(TbProtoQueueMsg.class); + int pollsBeforeEvent = queueConsumer.getPollCount(); + queueConsumer.enqueue(List.of(event)); + + // While sync is in progress the consumer stays parked: it neither polls nor consumes the event, + // so the event survives in the queue instead of being read-and-skipped. + sleepQuietly(POLL_INTERVAL_MS * 5); + assertThat(queueConsumer.getPolledEvents()) + .as("event must not be polled while sync is in progress (it must stay queued, not be dropped)") + .isEmpty(); + assertThat(queueConsumer.getPollCount()) + .as("consumer must not poll at all while the gate is closed") + .isEqualTo(pollsBeforeEvent); + + // Sync completes: the gate opens and the held event is finally picked up by the consumer. + // (We assert at the poll boundary - the actual drop point - since the downlink-send path that + // processMsgs drives afterwards is not reachable from a unit test.) + setReadinessFlags(true, false, false); + await().atMost(5, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(queueConsumer.getPolledEvents()) + .as("event held during sync must be picked up once sync completes, not lost") + .hasSize(1)); + } + + private boolean isReadyToProcessGeneralEvents() { + return Boolean.TRUE.equals(ReflectionTestUtils.invokeMethod(session, "isReadyToProcessGeneralEvents")); + } + + private void setReadinessFlags(boolean connected, boolean syncInProgress, boolean highPriorityProcessing) { + ReflectionTestUtils.setField(session, "connected", connected); + ReflectionTestUtils.setField(session, "syncInProgress", syncInProgress); + ReflectionTestUtils.setField(session, "isHighPriorityProcessing", highPriorityProcessing); + } + + private static void awaitParkedOnClosedGate(RecordingEdgeEventConsumer consumer) { + await().atMost(5, TimeUnit.SECONDS).until(() -> { + int before = consumer.getPollCount(); + // Several poll intervals with no new poll means the loop is parked on the readiness gate rather than + // blocked inside poll(), so it is safe to enqueue without racing an in-flight read. + sleepQuietly(POLL_INTERVAL_MS * 5); + return consumer.getPollCount() == before; + }); + } + + private static void sleepQuietly(long millis) { + try { + Thread.sleep(millis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + /** + * Fake edge-event consumer that records what it was actually polled, so a test can prove an event stayed in the + * queue while the readiness gate was closed and was only read once it opened. + */ + private static class RecordingEdgeEventConsumer implements TbQueueConsumer> { + + private final Queue>> pending = new ConcurrentLinkedQueue<>(); + private final List> polledEvents = new CopyOnWriteArrayList<>(); + private final AtomicInteger pollCount = new AtomicInteger(); + private volatile boolean stopped; + + void enqueue(List> batch) { + pending.add(batch); + } + + int getPollCount() { + return pollCount.get(); + } + + List> getPolledEvents() { + return polledEvents; + } + + @Override + public List> poll(long durationInMillis) { + pollCount.incrementAndGet(); + List> batch = pending.poll(); + if (batch != null) { + polledEvents.addAll(batch); + return batch; + } + sleepQuietly(durationInMillis); + return Collections.emptyList(); + } + + @Override + public String getTopic() { + return "test-edge-event-topic"; + } + + @Override + public void subscribe() { + } + + @Override + public void subscribe(Set partitions) { + } + + @Override + public void stop() { + stopped = true; + } + + @Override + public void unsubscribe() { + stopped = true; + } + + @Override + public void commit() { + } + + @Override + public boolean isStopped() { + return stopped; + } + + @Override + public List getFullTopicNames() { + return Collections.emptyList(); + } + + @Override + public Set getPartitions() { + return Collections.emptySet(); + } + + } + +}