[backport camel-4.22.x] CAMEL-25781: camel-azure-eventhubs - an event that completes when the checkpoint timer of its window has run is checkpointed by the next timer

components/camel-azure/camel-azure-eventhubs/src/main/java/org/apache/camel/component/azure/eventhubs/EventHubsCheckpointUpdaterTask.java

@@ -17,6 +17,8 @@
 package org.apache.camel.component.azure.eventhubs;
 
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 import com.azure.messaging.eventhubs.models.EventContext;
 import org.slf4j.Logger;
@@ -32,23 +34,60 @@ public class EventHubsCheckpointUpdaterTask implements Runnable {
     private volatile EventContext eventContext;
     private final AtomicInteger processedEvents;
     private volatile long scheduledTime;
+    // the lock of the consumer, which counts the events under it
+    private final Lock lock;
+    // set when the task has taken the count of its window, or was cancelled; read and written under the lock only
+    private boolean done;
 
+    /**
+     * Creates a task that synchronizes on a private lock of its own, unrelated to any consumer: for standalone or test
+     * use. {@link EventHubsConsumer} uses {@link #EventHubsCheckpointUpdaterTask(EventContext, AtomicInteger, Lock)}
+     * with its own lock, which it holds while it counts the events; with a private lock an event could be counted in a
+     * window whose timer has already taken the count.
+     *
+     * @param eventContext    the event context to checkpoint
+     * @param processedEvents the count of the events processed since the last checkpoint
+     */
     public EventHubsCheckpointUpdaterTask(EventContext eventContext, AtomicInteger processedEvents) {
+        this(eventContext, processedEvents, new ReentrantLock());
+    }
+
+    public EventHubsCheckpointUpdaterTask(EventContext eventContext, AtomicInteger processedEvents, Lock lock) {
         this.eventContext = eventContext;
         this.processedEvents = processedEvents;
+        this.lock = lock;
     }
 
     @Override
     public void run() {
-        if (processedEvents.get() > 0) {
-            LOG.debug("checkpointing offset after reaching timeout, with a batch of {}", processedEvents.get());
-            eventContext.updateCheckpointAsync()
+        // take the count when the checkpoint is started: an event that completes while it is being stored is counted
+        // for the next window from 0, and a checkpoint by batch size in the meantime cannot be confused with this one.
+        // This is done under the lock of the consumer, so an event is either counted (and its context set) before, and
+        // checkpointed here, or after, and then it sees that this task is done and starts a new window.
+        final int processed;
+        final EventContext context;
+        lock.lock();
+        try {
+            if (done) {
+                LOG.debug("skip checkpointing offset: the window was already checkpointed by batch size");
+                return;
+            }
+            done = true;
+            processed = processedEvents.getAndSet(0);
+            context = eventContext;
+        } finally {
+            lock.unlock();
+        }
+        if (processed > 0) {
+            LOG.debug("checkpointing offset after reaching timeout, with a batch of {}", processed);
+            context.updateCheckpointAsync()
                     .subscribe(unused -> LOG.debug("Processed one event..."),
-                            error -> LOG.warn("Error when updating Checkpoint: {}", error.getMessage(), error),
-                            () -> {
-                                LOG.debug("Checkpoint updated.");
-                                processedEvents.set(0);
-                            });
+                            error -> {
+                                LOG.warn("Error when updating Checkpoint: {}", error.getMessage(), error);
+                                // these events are not checkpointed: count them again for the next checkpoint
+                                processedEvents.addAndGet(processed);
+                            },
+                            () -> LOG.debug("Checkpoint updated."));
         } else {
             LOG.debug("skip checkpointing offset even if timeout is reached. No events processed");
         }
@@ -66,6 +105,26 @@ public long getScheduledTime() {
         return scheduledTime;
     }
 
+    /**
+     * Whether this task has taken the count of its window (or was cancelled), so a new event must start a new window.
+     * Called under the lock of the consumer.
+     */
+    public boolean isDone() {
+        return done;
+    }
+
+    /**
+     * Cancels the task under the lock of the consumer: if it is already running, it takes nothing.
+     */
+    public void cancel() {
+        lock.lock();
+        try {
+            done = true;
+        } finally {
+            lock.unlock();
+        }
+    }
+
     public boolean isExpired() {
         return System.currentTimeMillis() > scheduledTime;
     }

components/camel-azure/camel-azure-eventhubs/src/main/java/org/apache/camel/component/azure/eventhubs/EventHubsConsumer.java

@@ -220,9 +220,10 @@ private void onErrorListener(final ErrorContext errorContext) {
     /**
      * Strategy to commit the offset after message being processed successfully.
      *
-     * @param exchange the exchange
+     * @param exchange     the exchange
+     * @param eventContext the event context of the event, used to update the checkpoint
      */
-    private void processCommit(final Exchange exchange, final EventContext eventContext) {
+    void processCommit(final Exchange exchange, final EventContext eventContext) {
         lock.lock();
         try {
             final String partitionId = eventContext.getPartitionContext().getPartitionId();
@@ -231,8 +232,11 @@ private void processCommit(final Exchange exchange, final EventContext eventCont
             EventHubsCheckpointUpdaterTask checkpointTask = checkpointTasksByPartition.get(partitionId);
             ScheduledFuture<?> scheduledTask = scheduledTasksByPartition.get(partitionId);
 
-            if (checkpointTask == null || checkpointTask.isExpired()) {
-                checkpointTask = new EventHubsCheckpointUpdaterTask(eventContext, processedEvents);
+            // a new window starts when the timer of the current one has run, not when the wall clock passed its time: the
+            // timer can run in the millisecond of its scheduled time, or later, and the event must be counted in the
+            // window whose timer has not taken the count yet
+            if (checkpointTask == null || checkpointTask.isDone()) {
+                checkpointTask = new EventHubsCheckpointUpdaterTask(eventContext, processedEvents, lock);
                 // delegate the checkpoint update to a dedicated Thread
                 long timeout = getConfiguration().getCheckpointBatchTimeout();
                 checkpointTask.setScheduledTime(System.currentTimeMillis() + timeout);
@@ -246,13 +250,20 @@ private void processCommit(final Exchange exchange, final EventContext eventCont
 
             try {
                 var cnt = processedEvents.incrementAndGet();
-                if (cnt == getConfiguration().getCheckpointBatchSize()) {
+                // >= and not ==: a failed timeout checkpoint adds its count back, so the count can pass the batch size
+                if (cnt >= getConfiguration().getCheckpointBatchSize()) {
                     processedEvents.set(0);
                     exchange.getIn().setHeader(EventHubsConstants.CHECKPOINT_UPDATED_BY, COMPLETED_BY_SIZE);
                     LOG.debug("eventhub consumer batch size of reached for partition {}", partitionId);
+                    // a timer that has already started takes nothing (it waits for the lock)
+                    checkpointTask.cancel();
                     if (scheduledTask != null) {
                         scheduledTask.cancel(false);
                     }
+                    // the next event starts a new timeout window: the cancelled task must not be reused until its
+                    // scheduled time, else the events of the rest of that window are not checkpointed by the timeout
+                    checkpointTasksByPartition.remove(partitionId);
+                    scheduledTasksByPartition.remove(partitionId);
                     eventContext.updateCheckpointAsync()
                             .subscribe(unused -> LOG.debug("Processed one event..."),
                                     error -> LOG.warn("Error when updating Checkpoint: {}", error.getMessage(), error),

components/camel-azure/camel-azure-eventhubs/src/test/java/org/apache/camel/component/azure/eventhubs/EventHubsConsumerCheckpointAfterBatchSizeTest.java

@@ -0,0 +1,229 @@
+/*
+ * 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.camel.component.azure.eventhubs;
+
+import java.lang.reflect.Method;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+
+import com.azure.messaging.eventhubs.models.EventContext;
+import com.azure.messaging.eventhubs.models.PartitionContext;
+import org.apache.camel.AsyncProcessor;
+import org.apache.camel.CamelContext;
+import org.apache.camel.Endpoint;
+import org.apache.camel.Exchange;
+import org.apache.camel.ExtendedCamelContext;
+import org.apache.camel.spi.ExchangeFactory;
+import org.apache.camel.support.DefaultExchange;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+
+import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * The events completed after a checkpoint by batch size must be checkpointed by the batch timeout, also when no further
+ * event arrives (end of a burst). Uses a real scheduler with a short batch timeout.
+ */
+class EventHubsConsumerCheckpointAfterBatchSizeTest {
+
+    private final EventHubsEndpoint endpoint = mock();
+    private final EventHubsConfiguration configuration = new EventHubsConfiguration();
+    private final AsyncProcessor processor = mock();
+    private final CamelContext context = mock();
+    private final ExtendedCamelContext ecc = mock();
+    private final ExchangeFactory exchangeFactory = mock();
+    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
+
+    private EventHubsConsumer consumer;
+    private Method processCommit;
+
+    @BeforeEach
+    void setUp() throws Exception {
+        configuration.setCheckpointBatchSize(2);
+        configuration.setCheckpointBatchTimeout(1000);
+
+        when(endpoint.getCamelContext()).thenReturn(context);
+        when(context.getCamelContextExtension()).thenReturn(ecc);
+        when(ecc.getExchangeFactory()).thenReturn(exchangeFactory);
+        when(exchangeFactory.newExchangeFactory(any())).thenReturn(exchangeFactory);
+        when(exchangeFactory.create(any(Endpoint.class), anyBoolean()))
+                .thenAnswer(invocation -> DefaultExchange.newFromEndpoint(invocation.getArgument(0)));
+        when(endpoint.getConfiguration()).thenReturn(configuration);
+        consumer = new EventHubsConsumer(endpoint, processor);
+
+        var field = EventHubsConsumer.class.getDeclaredField("scheduledExecutorService");
+        field.setAccessible(true);
+        field.set(consumer, scheduler);
+
+        processCommit = EventHubsConsumer.class.getDeclaredMethod("processCommit", Exchange.class, EventContext.class);
+        processCommit.setAccessible(true);
+    }
+
+    @AfterEach
+    void tearDown() {
+        scheduler.shutdownNow();
+    }
+
+    @Test
+    void eventAfterBatchSizeCheckpointIsCheckpointedByTimeout() throws Exception {
+        EventContext first = eventContext("0");
+        EventContext second = eventContext("0");
+        EventContext third = eventContext("0");
+
+        complete(first);
+        complete(second);
+        // the batch size is reached: checkpoint at the second event
+        verify(second, times(1)).updateCheckpointAsync();
+
+        // the last event of the burst: no other event follows
+        complete(third);
+
+        await().atMost(10, TimeUnit.SECONDS)
+                .untilAsserted(() -> verify(third, times(1)).updateCheckpointAsync());
+    }
+
+    @Test
+    void eventsAfterSecondBatchSizeCheckpointAreCheckpointedByTimeout() throws Exception {
+        EventContext[] events = new EventContext[5];
+        for (int i = 0; i < events.length; i++) {
+            events[i] = eventContext("0");
+            complete(events[i]);
+        }
+        // checkpoints by batch size at the second and the fourth event
+        verify(events[1], times(1)).updateCheckpointAsync();
+        verify(events[3], times(1)).updateCheckpointAsync();
+
+        await().atMost(10, TimeUnit.SECONDS)
+                .untilAsserted(() -> verify(events[4], times(1)).updateCheckpointAsync());
+    }
+
+    @Test
+    void eventCompletedWhileTheTimeoutCheckpointIsStoredIsCheckpointedByTheNextTimeout() throws Exception {
+        configuration.setCheckpointBatchSize(10);
+        // the checkpoint of the first window is stored only when the test says so
+        Sinks.Empty<Void> stored = Sinks.empty();
+        AtomicLong checkpointedAt = new AtomicLong();
+        EventContext first = eventContext("0");
+        doAnswer(invocation -> {
+            checkpointedAt.set(System.currentTimeMillis());
+            return stored.asMono();
+        }).when(first).updateCheckpointAsync();
+        EventContext second = eventContext("0");
+
+        complete(first);
+        await().atMost(10, TimeUnit.SECONDS).untilAsserted(() -> verify(first, times(1)).updateCheckpointAsync());
+        // the window of the first task is over, so the next event starts a new task
+        await().atMost(10, TimeUnit.SECONDS).until(() -> System.currentTimeMillis() > checkpointedAt.get());
+
+        // an event completes while the checkpoint of the first window is being stored
+        complete(second);
+        stored.tryEmitEmpty();
+
+        await().atMost(10, TimeUnit.SECONDS)
+                .untilAsserted(() -> verify(second, times(1)).updateCheckpointAsync());
+    }
+
+    @Test
+    void eventAfterBatchSizeCheckpointDuringTimeoutCheckpointIsCheckpointed() throws Exception {
+        configuration.setCheckpointBatchSize(3);
+        // the checkpoint of the first window is stored only when the test says so
+        Sinks.Empty<Void> stored = Sinks.empty();
+        AtomicLong checkpointedAt = new AtomicLong();
+        EventContext first = eventContext("0");
+        doAnswer(invocation -> {
+            checkpointedAt.set(System.currentTimeMillis());
+            return stored.asMono();
+        }).when(first).updateCheckpointAsync();
+        EventContext second = eventContext("0");
+        EventContext third = eventContext("0");
+        EventContext fourth = eventContext("0");
+
+        complete(first);
+        await().atMost(10, TimeUnit.SECONDS).untilAsserted(() -> verify(first, times(1)).updateCheckpointAsync());
+        await().atMost(10, TimeUnit.SECONDS).until(() -> System.currentTimeMillis() > checkpointedAt.get());
+
+        // while the timeout checkpoint of the first event is stored, more events complete, and the count climbs back
+        // to the value that checkpoint started with
+        complete(second);
+        complete(third);
+        complete(fourth);
+        stored.tryEmitEmpty();
+
+        // the last event of the burst must still be checkpointed, by batch size or by the next timeout
+        await().atMost(10, TimeUnit.SECONDS)
+                .untilAsserted(() -> verify(fourth, times(1)).updateCheckpointAsync());
+    }
+
+    @Test
+    void batchSizeCheckpointStillWorksAfterFailedTimeoutCheckpoint() throws Exception {
+        configuration.setCheckpointBatchSize(3);
+        // the checkpoint of the first window fails only when the test says so
+        Sinks.Empty<Void> stored = Sinks.empty();
+        AtomicLong checkpointedAt = new AtomicLong();
+        EventContext first = eventContext("0");
+        EventContext second = eventContext("0");
+        doAnswer(invocation -> {
+            checkpointedAt.set(System.currentTimeMillis());
+            return stored.asMono();
+        }).when(second).updateCheckpointAsync();
+        EventContext third = eventContext("0");
+        EventContext fourth = eventContext("0");
+
+        complete(first);
+        complete(second);
+        await().atMost(10, TimeUnit.SECONDS).untilAsserted(() -> verify(second, times(1)).updateCheckpointAsync());
+        await().atMost(10, TimeUnit.SECONDS).until(() -> System.currentTimeMillis() > checkpointedAt.get());
+
+        // one event completes while the timeout checkpoint of two events is stored, then storing fails: the two
+        // events are counted again, so the count is 3, and the next event takes it past the batch size
+        complete(third);
+        stored.tryEmitError(new IllegalStateException("store failed"));
+
+        Exchange exchange = DefaultExchange.newFromEndpoint(endpoint);
+        processCommit.invoke(consumer, exchange, fourth);
+
+        verify(fourth, times(1)).updateCheckpointAsync();
+        assertEquals(EventHubsConstants.COMPLETED_BY_SIZE,
+                exchange.getIn().getHeader(EventHubsConstants.CHECKPOINT_UPDATED_BY));
+    }
+
+    private void complete(EventContext eventContext) throws Exception {
+        processCommit.invoke(consumer, DefaultExchange.newFromEndpoint(endpoint), eventContext);
+    }
+
+    private static EventContext eventContext(String partitionId) {
+        EventContext eventContext = mock(EventContext.class);
+        PartitionContext partitionContext = mock(PartitionContext.class);
+        when(partitionContext.getPartitionId()).thenReturn(partitionId);
+        when(eventContext.getPartitionContext()).thenReturn(partitionContext);
+        when(eventContext.updateCheckpointAsync()).thenReturn(Mono.empty());
+        return eventContext;
+    }
+}