diff --git a/pulsar-io/kafka-connect-adaptor/pom.xml b/pulsar-io/kafka-connect-adaptor/pom.xml index b70996b77b11a..1a50edafb044b 100644 --- a/pulsar-io/kafka-connect-adaptor/pom.xml +++ b/pulsar-io/kafka-connect-adaptor/pom.xml @@ -185,6 +185,12 @@ test + + org.awaitility + awaitility + test + + ${project.groupId} pulsar-broker diff --git a/pulsar-io/kafka-connect-adaptor/src/main/java/org/apache/pulsar/io/kafka/connect/AbstractKafkaConnectSource.java b/pulsar-io/kafka-connect-adaptor/src/main/java/org/apache/pulsar/io/kafka/connect/AbstractKafkaConnectSource.java index 270decd8ab533..e048ddf724414 100644 --- a/pulsar-io/kafka-connect-adaptor/src/main/java/org/apache/pulsar/io/kafka/connect/AbstractKafkaConnectSource.java +++ b/pulsar-io/kafka-connect-adaptor/src/main/java/org/apache/pulsar/io/kafka/connect/AbstractKafkaConnectSource.java @@ -31,11 +31,13 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; -import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import lombok.Getter; import lombok.extern.slf4j.Slf4j; import org.apache.kafka.connect.connector.Task; +import org.apache.kafka.connect.errors.ConnectException; import org.apache.kafka.connect.runtime.TaskConfig; import org.apache.kafka.connect.source.SourceConnector; import org.apache.kafka.connect.source.SourceRecord; @@ -68,7 +70,6 @@ public abstract class AbstractKafkaConnectSource implements Source { // pulsar io related variables private Iterator currentBatch = null; - private CompletableFuture flushFuture; private OffsetBackingStore offsetStore; private OffsetStorageReader offsetReader; private String topicNamespace; @@ -76,6 +77,8 @@ public abstract class AbstractKafkaConnectSource implements Source { public OffsetStorageWriter offsetWriter; // number of outstandingRecords that have been polled but not been acked private final AtomicInteger outstandingRecords = new AtomicInteger(0); + private final AtomicBoolean flushing = new AtomicBoolean(false); + private final AtomicReference> flushFutureRef = new AtomicReference<>(); public static final String CONNECTOR_CLASS = "kafkaConnectorSourceClass"; @@ -162,38 +165,103 @@ public void open(Map config, SourceContext sourceContext) throws sourceTask.initialize(sourceTaskContext); sourceTask.start(taskConfig); } + private void onOffsetsFlushed(Throwable error, CompletableFuture snapshotFlushFuture) { + if (error != null) { + log.error("Failed to flush offsets to storage: ", error); + offsetWriter.cancelFlush(); + snapshotFlushFuture.completeExceptionally(new Exception("No Offsets Added Error", error)); + return; + } + try { + sourceTask.commit(); + if (log.isDebugEnabled()) { + log.debug("Finished flushing offsets to storage"); + } + snapshotFlushFuture.complete(null); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + log.warn("Flush interrupted, cancelling", ie); + offsetWriter.cancelFlush(); + snapshotFlushFuture.completeExceptionally(new Exception("Failed to commit offsets", ie)); + } catch (Throwable t) { + log.warn("Flush failed, cancelling", t); + offsetWriter.cancelFlush(); + snapshotFlushFuture.completeExceptionally(new Exception("Failed to commit offsets", t)); + } + } + + private void triggerOffsetsFlushIfNeeded() { + final CompletableFuture snapshotFlushFuture = flushFutureRef.get(); + // Only flush when we have a batch in flight, nothing outstanding, and a pending future + if (snapshotFlushFuture == null || snapshotFlushFuture.isDone() || outstandingRecords.get() != 0) { + return; + } + if (!flushing.compareAndSet(false, true)) { + return; // someone else is flushing + } + try { + if (offsetWriter.beginFlush()) { + offsetWriter.doFlush((error, ignored) -> { + try { + onOffsetsFlushed(error, snapshotFlushFuture); + } finally { + flushing.set(false); + } + }); + } else { + try { + onOffsetsFlushed(null, snapshotFlushFuture); + } finally { + flushing.set(false); + } + } + } catch (ConnectException alreadyFlushing) { + // Another thread initiated the flush; let their callback complete the future. + // Keep 'flushing' = true until read() finalizes the batch. + } catch (Exception t) { + try { + onOffsetsFlushed(t, snapshotFlushFuture); + } finally { + flushing.set(false); + } + } + } @Override public synchronized Record read() throws Exception { while (true) { if (currentBatch == null) { - flushFuture = new CompletableFuture<>(); List recordList = sourceTask.poll(); if (recordList == null || recordList.isEmpty()) { continue; } outstandingRecords.addAndGet(recordList.size()); currentBatch = recordList.iterator(); + + final CompletableFuture newFuture = new CompletableFuture<>(); + flushFutureRef.set(newFuture); } if (currentBatch.hasNext()) { AbstractKafkaSourceRecord processRecord = processSourceRecord(currentBatch.next()); if (processRecord == null || processRecord.isEmpty()) { outstandingRecords.decrementAndGet(); - continue; + triggerOffsetsFlushIfNeeded(); } else { return processRecord; } } else { - // there is no records any more, then waiting for the batch to complete writing - // to sink and the offsets are committed as well, then do next round read. + final CompletableFuture snapshotFlushFuture = flushFutureRef.get(); try { - flushFuture.get(); + if (snapshotFlushFuture != null) { + snapshotFlushFuture.get(); + } } catch (ExecutionException ex) { - // log the error, continue execution log.error("execution exception while get flushFuture", ex); throw new Exception("Flush failed", ex.getCause()); } finally { - flushFuture = null; + flushing.set(false); + // Clear only if this is still the current batch future + flushFutureRef.compareAndSet(snapshotFlushFuture, null); currentBatch = null; } } @@ -272,62 +340,19 @@ public boolean isEmpty() { return this.value == null; } - private void completedFlushOffset(Throwable error, Void result) { - if (error != null) { - log.error("Failed to flush offsets to storage: ", error); - currentBatch = null; - offsetWriter.cancelFlush(); - flushFuture.completeExceptionally(new Exception("No Offsets Added Error")); - } else { - try { - sourceTask.commit(); - - log.info("Finished flushing offsets to storage"); - currentBatch = null; - flushFuture.complete(null); - } catch (InterruptedException exception) { - log.warn("Flush of {} offsets interrupted, cancelling", this); - Thread.currentThread().interrupt(); - offsetWriter.cancelFlush(); - flushFuture.completeExceptionally(new Exception("Failed to commit offsets", exception)); - } catch (Throwable t) { - // SourceTask can throw unchecked ConnectException/KafkaException. - // Make sure the future is cancelled in that case - log.warn("Flush of {} offsets failed, cancelling", this); - offsetWriter.cancelFlush(); - flushFuture.completeExceptionally(new Exception("Failed to commit offsets", t)); - } - } - } - @Override public void ack() { - // TODO: should flush for each batch. not wait for a time for acked all. - // How to handle order between each batch. QueueList>. check if head is all acked. - boolean canFlush = (outstandingRecords.decrementAndGet() == 0); - - // consumed all the records, flush the offsets - if (canFlush && flushFuture != null) { - if (!offsetWriter.beginFlush()) { - log.error("When beginFlush, No offsets to commit!"); - flushFuture.completeExceptionally(new Exception("No Offsets Added Error when beginFlush")); - return; - } - - Future doFlush = offsetWriter.doFlush(this::completedFlushOffset); - if (doFlush == null) { - // Offsets added in processSourceRecord, But here no offsets to commit - log.error("No offsets to commit!"); - flushFuture.completeExceptionally(new Exception("No Offsets Added Error")); - return; - } + // Decrement and let the centralized flusher decide if we should flush now + if (outstandingRecords.decrementAndGet() == 0) { + triggerOffsetsFlushIfNeeded(); } } @Override public void fail() { - if (flushFuture != null) { - flushFuture.completeExceptionally(new Exception("Sink Error")); + final CompletableFuture snapshotFlushFuture = flushFutureRef.get(); + if (snapshotFlushFuture != null) { + snapshotFlushFuture.completeExceptionally(new Exception("Sink Error")); } } } diff --git a/pulsar-io/kafka-connect-adaptor/src/test/java/org/apache/pulsar/io/kafka/connect/KafkaConnectSourceTest.java b/pulsar-io/kafka-connect-adaptor/src/test/java/org/apache/pulsar/io/kafka/connect/KafkaConnectSourceTest.java index dfa3031abb2c1..5d414b6373ca7 100644 --- a/pulsar-io/kafka-connect-adaptor/src/test/java/org/apache/pulsar/io/kafka/connect/KafkaConnectSourceTest.java +++ b/pulsar-io/kafka-connect-adaptor/src/test/java/org/apache/pulsar/io/kafka/connect/KafkaConnectSourceTest.java @@ -19,6 +19,8 @@ package org.apache.pulsar.io.kafka.connect; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; @@ -28,19 +30,25 @@ import java.io.File; import java.io.OutputStream; import java.nio.file.Files; +import java.time.Duration; import java.util.HashMap; import java.util.Map; import java.util.Optional; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.kafka.connect.file.FileStreamSourceConnector; import org.apache.kafka.connect.runtime.TaskConfig; import org.apache.kafka.connect.source.SourceRecord; +import org.apache.kafka.connect.storage.OffsetStorageWriter; import org.apache.pulsar.client.api.ProducerConsumerBase; import org.apache.pulsar.client.api.PulsarClient; import org.apache.pulsar.common.schema.KeyValue; import org.apache.pulsar.functions.api.Record; import org.apache.pulsar.io.core.SourceContext; import org.apache.pulsar.io.kafka.connect.KafkaConnectSource.KafkaSourceRecord; +import org.awaitility.Awaitility; +import org.mockito.ArgumentMatchers; +import org.mockito.Mockito; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -49,7 +57,7 @@ * Test the implementation of {@link KafkaConnectSource}. */ @Slf4j -public class KafkaConnectSourceTest extends ProducerConsumerBase { +public class KafkaConnectSourceTest extends ProducerConsumerBase { private String offsetTopicName; // The topic to publish data to, for kafkaSource @@ -96,6 +104,174 @@ public void testOpenAndReadConnectorConfig() throws Exception { testOpenAndReadTask(config); } + private OffsetStorageWriter wrapOffsetWriterWithSpy(KafkaConnectSource source) throws IllegalAccessException { + OffsetStorageWriter original = source.getOffsetWriter(); + OffsetStorageWriter spy = org.mockito.Mockito.spy(original); + FieldUtils.writeField(source, "offsetWriter", spy, true); + return spy; + } + + + @Test(timeOut = 30000) + public void testFlushWhenAllMessagesFilteredWithoutBlocking() throws Exception { + + Map config = getConfig(); + config.put(TaskConfig.TASK_CLASS_CONFIG, "org.apache.kafka.connect.file.FileStreamSourceTask"); + + config.put("transforms", "Filter"); + config.put("transforms.Filter.type", "org.apache.kafka.connect.transforms.Filter"); + config.put("transforms.Filter.predicate", "DropMeTopic"); + + config.put("predicates", "DropMeTopic"); + config.put("predicates.DropMeTopic.type", "org.apache.kafka.connect.transforms.predicates.TopicNameMatches"); + config.put("predicates.DropMeTopic.pattern", ".*my-property/my-ns/kafka-connect-source.*"); + + kafkaConnectSource = new KafkaConnectSource(); + kafkaConnectSource.open(config, context); + + OffsetStorageWriter spyWriter = wrapOffsetWriterWithSpy(kafkaConnectSource); + + try (OutputStream os = Files.newOutputStream(tempFile.toPath())) { + os.write("first\n".getBytes()); + os.flush(); + os.write("second\n".getBytes()); + os.flush(); + } + + Thread t = new Thread(() -> { + try { + kafkaConnectSource.read(); + } catch (Exception e) { + // Ignore, it just needs the loop running through a batch where every message is filtered out + log.error("Exception in read thread", e); + } + }); + t.setDaemon(true); + t.start(); + + try { + // First ensure offsets are being written for filtered records. + Awaitility.await() + .atMost(Duration.ofSeconds(10)) + .untilAsserted(() -> { + Mockito.verify(spyWriter, Mockito.atLeastOnce()) + .offset(ArgumentMatchers.any(), ArgumentMatchers.any()); + }); + + // Then ensure a flush cycle is triggered without waiting for any acks (since all are filtered). + Awaitility.await() + .atMost(java.time.Duration.ofSeconds(10)) + .untilAsserted(() -> { + Mockito.verify(spyWriter, Mockito.atLeastOnce()) + .beginFlush(); + Mockito.verify(spyWriter, Mockito.atLeastOnce()) + .doFlush(ArgumentMatchers.any()); + }); + } finally { + kafkaConnectSource.close(); + t.interrupt(); + } + } + + @Test(timeOut = 40000) + public void testNoFlushUntilAck() throws Exception { + Map config = getConfig(); + config.put(TaskConfig.TASK_CLASS_CONFIG, "org.apache.kafka.connect.file.FileStreamSourceTask"); + + kafkaConnectSource = new KafkaConnectSource(); + kafkaConnectSource.open(config, context); + + OffsetStorageWriter spyWriter = wrapOffsetWriterWithSpy(kafkaConnectSource); + + try (OutputStream os = Files.newOutputStream(tempFile.toPath())) { + os.write("one\n".getBytes()); + os.flush(); + } + + Record> readRecord = kafkaConnectSource.read(); + + // Verify no flush happens while there is one outstanding record (no ack yet) + verify(spyWriter, timeout(10000).times(0)).beginFlush(); + verify(spyWriter, timeout(10000).times(0)) + .doFlush(ArgumentMatchers.any()); + + readRecord.ack(); + + // Verify flush happens after ack + verify(spyWriter, timeout(10000).atLeastOnce()).beginFlush(); + verify(spyWriter, timeout(10000).atLeastOnce()) + .doFlush(ArgumentMatchers.any()); + + kafkaConnectSource.close(); + } + + @Test(timeOut = 40000) + public void testPartialAckNoFlush_ThenFlushOnAllAck() throws Exception { + Map config = getConfig(); + config.put(TaskConfig.TASK_CLASS_CONFIG, "org.apache.kafka.connect.file.FileStreamSourceTask"); + + kafkaConnectSource = new KafkaConnectSource(); + kafkaConnectSource.open(config, context); + + OffsetStorageWriter spyWriter = wrapOffsetWriterWithSpy(kafkaConnectSource); + + try (OutputStream os = Files.newOutputStream(tempFile.toPath())) { + os.write("first\n".getBytes()); + os.flush(); + os.write("second\n".getBytes()); + os.flush(); + } + + Record> r1 = kafkaConnectSource.read(); + Record> r2 = kafkaConnectSource.read(); + + // Ack only the first; one outstanding remains -> no flush should happen + r1.ack(); + verify(spyWriter, timeout(10000).times(0)).beginFlush(); + verify(spyWriter, timeout(10000).times(0)) + .doFlush(ArgumentMatchers.any()); + + // Ack the second; outstanding reaches zero -> flush should be triggered + r2.ack(); + verify(spyWriter, timeout(10000).atLeastOnce()).beginFlush(); + verify(spyWriter, timeout(10000).atLeastOnce()) + .doFlush(ArgumentMatchers.any()); + + kafkaConnectSource.close(); + } + + @Test(timeOut = 20000) + public void testAckFlush() throws Exception { + Map config = getConfig(); + config.put(TaskConfig.TASK_CLASS_CONFIG, "org.apache.kafka.connect.file.FileStreamSourceTask"); + + kafkaConnectSource = new KafkaConnectSource(); + kafkaConnectSource.open(config, context); + + OffsetStorageWriter spyWriter = wrapOffsetWriterWithSpy(kafkaConnectSource); + + try (OutputStream os = Files.newOutputStream(tempFile.toPath())) { + os.write("first\n".getBytes()); + os.flush(); + os.write("second\n".getBytes()); + os.flush(); + } + + Record> r1 = kafkaConnectSource.read(); + Record> r2 = kafkaConnectSource.read(); + + // Ack both, no outstanding remains -> flush should happen + r1.ack(); + r2.ack(); + + verify(spyWriter, timeout(10000).atLeastOnce()).beginFlush(); + verify(spyWriter, timeout(10000).atLeastOnce()) + .doFlush(ArgumentMatchers.any()); + + kafkaConnectSource.close(); + } + + @Test public void testOpenAndReadTaskDirect() throws Exception { Map config = getConfig(); @@ -157,8 +333,8 @@ private void runTopicNameTest(Map config, String topicName, Stri sourceOffset.put("test", 0); value.put("myField", "42"); SourceRecord srcRecord = new SourceRecord( - sourcePartition, sourceOffset, topicName, null, - null, null, null, value + sourcePartition, sourceOffset, topicName, null, + null, null, null, value ); KafkaSourceRecord record = kafkaConnectSource.processSourceRecord(srcRecord); @@ -197,8 +373,8 @@ private void runTransformTest(Map config, boolean expectTransfor Map value = new HashMap<>(); value.put("myField", "42"); SourceRecord record = new SourceRecord( - null, null, "test-topic", null, - null, null, null, value + null, null, "test-topic", null, + null, null, null, value ); SourceRecord transformed = kafkaConnectSource.applyTransforms(record);