Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
* Used to track source records that have been (or are about to be) dispatched to a producer and their accompanying
Expand All @@ -42,9 +46,12 @@ class SubmittedRecords {

// Visible for testing
final Map<Map<String, Object>, Deque<SubmittedRecord>> records;
private AtomicInteger numUnackedMessages;
Comment thread
C0urante marked this conversation as resolved.
Outdated
private CountDownLatch messageDrainLatch;

public SubmittedRecords() {
this.records = new HashMap<>();
this.numUnackedMessages = new AtomicInteger(0);
}

/**
Expand All @@ -68,6 +75,9 @@ SubmittedRecord submit(Map<String, Object> partition, Map<String, Object> offset
SubmittedRecord result = new SubmittedRecord(partition, offset);
records.computeIfAbsent(result.partition(), p -> new LinkedList<>())
.add(result);
synchronized (this) {
numUnackedMessages.incrementAndGet();
}
Comment thread
C0urante marked this conversation as resolved.
return result;
}

Expand All @@ -89,7 +99,9 @@ public boolean removeLastOccurrence(SubmittedRecord record) {
if (deque.isEmpty()) {
records.remove(record.partition());
}
if (!result) {
if (result) {
messageAcked();
} else {
log.warn("Attempted to remove record from submitted queue for partition {}, but the record has not been submitted or has already been removed", record.partition());
}
return result;
Expand Down Expand Up @@ -132,6 +144,27 @@ public CommittableOffsets committableOffsets() {
return new CommittableOffsets(offsets, totalCommittableMessages, totalUncommittableMessages, records.size(), largestDequeSize, largestDequePartition);
}

/**
* Wait for all currently in-flight messages to be acknowledged, up to the requested timeout.
Comment thread
C0urante marked this conversation as resolved.
* @param timeout the maximum time to wait
* @param timeUnit the time unit of the timeout argument
* @return whether all in-flight messages were acknowledged before the timeout elapsed
*/
public boolean awaitAllMessages(long timeout, TimeUnit timeUnit) {
Comment thread
rhauch marked this conversation as resolved.
// Create a new message drain latch as a local variable to avoid SpotBugs warnings about inconsistent synchronization
// on an instance variable when invoking CountDownLatch::await outside a synchronized block
CountDownLatch messageDrainLatch;
synchronized (this) {
messageDrainLatch = new CountDownLatch(numUnackedMessages.get());
this.messageDrainLatch = messageDrainLatch;
}
Comment on lines +154 to +158

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this code be simplified and the synchronized block removed altogether by changing the messageDrainLatch field to:

    private final AtomicReference<CountDownLatch> messageDrainLatch = new AtomicReference<>();

and then changing these lines above to:

Suggested change
// Create a new message drain latch as a local variable to avoid SpotBugs warnings about inconsistent synchronization
// on an instance variable when invoking CountDownLatch::await outside a synchronized block
CountDownLatch messageDrainLatch;
synchronized (this) {
messageDrainLatch = new CountDownLatch(numUnackedMessages.get());
this.messageDrainLatch = messageDrainLatch;
}
// Create a new message drain latch as a local variable to avoid SpotBugs warnings about inconsistent synchronization
// on an instance variable when invoking CountDownLatch::await outside a synchronized block
CountDownLatch messageDrainLatch = this.messageDrainLatch.updateAndGet(existing -> new CountDownLatch(numUnackedMessages.get()));

This synchronized block ensures that the latch is initialized sized with the number of unacked messages at the time this method is called but does not prevent new messages from being added. Using concurrent types for the two fields solves the first issue, while the second is prevented by having the WorkerSourceTask#execute() method call submit(...), which increments numUnackedMessages, and then only call awaitAllMessages() after the task has been told to stop (at which point submit(...) will not be called again and numAckedMessages will not be incremented).

And calling out that last assumption would be good. Might be as simple as adding the following to thesubmitMessages(...) JavaDoc:

     *
     * <p>This method should never be called after {@link #awaitAllMessages(long, TimeUnit)} has been called.

You'd also have to change the messageAcked() method to use the atomic reference:

    private void messageAcked() {
        numUnackedMessages.decrementAndGet();
        CountDownLatch messageDrainLatch = this.messageDrainLatch.get();
        if (messageDrainLatch != null) {
            messageDrainLatch.countDown();
        }
    }

and remove the synchronized keyword. Again, I don't think we need to atomically update both the number of acked messages and counts down the latch atomically; we really just need them each to be updated consistently.

There are a few reasons why this might be better:

  1. We avoid synchronized keyword, which might be unclear to our future selves ("Is the documentation about not being thread-safe wrong?").
  2. We can make the fields be final, which IMO makes the logic a bit easier to follow.
  3. We don't need stricter synchronization than individual atomic updates of each field.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Randall. I agree that the synchronization here is, even if necessary, inelegant, and I hope that we can improve things. But I'm worried that the proposal here may be prone to a race condition.

Imagine we restructure the code with your suggestions and the result is this:

class SubmittedRecords {

    private final AtomicReference messageDrainLatch = new AtomicReference<>();

    private boolean awaitAllMessages(long timeout, TimeUnit timeUnit) {
        // (2)
        CountDownLatch messageDrainLatch = this.messageDrainLatch.updateAndGet(existing -> new CountDownLatch(numUnackedMessages.get()));
        try {
            return messageDrainLatch.await(timeout, timeUnit);
        } catch (InterruptedException e) {
            return false;
        }
    }

    private void messageAcked() {
        // (1)
        numUnackedMessages.decrementAndGet();
        // (3)
        CountDownLatch messageDrainLatch = this.messageDrainLatch.get();
        if (messageDrainLatch != null) {
            messageDrainLatch.countDown();
        }
    }
}

Isn't it still possible that the lines marked (1), (2), and (3) could execute in that order? And in that case, wouldn't it cause awaitAllMessages to return early since the CountDownLatch created in part (2) would use the already-decremented value of numUnackedMessages after part (1) was executed, but then also be counted down for the same message in part (3)?

FWIW, I originally used a volatile int for the numUnackedMessages field, but got a SpotBugs warning about incrementing a volatile field being a non-atomic operation for lines like numUnackedMessages++; in SubmittedRecords::submit. If we synchronize every access to that field, it shouldn't matter that increments/decrements are non-atomic, and we can consider adding an exemption to spotbugs-exclude.xml.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've just pushed an update that simplifies things a little bit but still utilizes synchronized blocks. I realized that, since we're synchronizing around every read of the numUnackedMessages field, it doesn't even need to be declared volatile, and so we can just replace it with a regular, primitive int.

try {
return messageDrainLatch.await(timeout, timeUnit);
} catch (InterruptedException e) {
return false;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this clear the interrupted flag before returning, since we're handling the interruption here?

@C0urante C0urante Nov 23, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this is already accomplished. According to the Javadocs for CountDownLatch::await:

If the current thread:
     • has its interrupted status set on entry to this method; or
     • is interrupted while waiting,
then InterruptedException is thrown and the current thread's interrupted status is cleared.

}
}

// Note that this will return null if either there are no committable offsets for the given deque, or the latest
// committable offset is itself null. The caller is responsible for distinguishing between the two cases.
private Map<String, Object> committableOffset(Deque<SubmittedRecord> queuedRecords) {
Expand All @@ -146,27 +179,36 @@ private boolean canCommitHead(Deque<SubmittedRecord> queuedRecords) {
return queuedRecords.peek() != null && queuedRecords.peek().acked();
}

static class SubmittedRecord {
private synchronized void messageAcked() {
numUnackedMessages.decrementAndGet();
if (messageDrainLatch != null) {
messageDrainLatch.countDown();
}
}

class SubmittedRecord {
private final Map<String, Object> partition;
private final Map<String, Object> offset;
private volatile boolean acked;
private final AtomicBoolean acked;

public SubmittedRecord(Map<String, Object> partition, Map<String, Object> offset) {
this.partition = partition;
this.offset = offset;
this.acked = false;
this.acked = new AtomicBoolean(false);
}

/**
* Acknowledge this record; signals that its offset may be safely committed.
* This is safe to be called from a different thread than what called {@link SubmittedRecords#submit(SourceRecord)}.
*/
public void ack() {
this.acked = true;
if (this.acked.compareAndSet(false, true)) {
messageAcked();
}
Comment on lines +206 to +208

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice simplification.

}

private boolean acked() {
return acked;
return acked.get();
}

private Map<String, Object> partition() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,6 @@
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.kafka.connect.runtime.SubmittedRecords.SubmittedRecord;
import static org.apache.kafka.connect.runtime.SubmittedRecords.CommittableOffsets;
import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;

Expand Down Expand Up @@ -260,6 +259,10 @@ public void execute() {
} catch (InterruptedException e) {
// Ignore and allow to exit.
} finally {
submittedRecords.awaitAllMessages(
workerConfig.getLong(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_CONFIG),

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Waiting for up to offset.flush.timeout.ms milliseconds here may cause shutdown to block for double that time (since the subsequent call to WorkerSourceTask::commitOffsets may also block for offset.flush.timeout.ms milliseconds while flushing offset information to the persistent backing store).

I initially considered an approach where we would require the entire combination of (awaiting in-flight messages, computing committable offsets, committing offsets to backing store) to complete in offset.flush.timeout.ms milliseconds, but this came with the unfortunate drawback that, if any in-flight record were undeliverable within the flush timeout, it would prevent any offsets from being committed, which seems counter to the original motivation of the changes we made for KAFKA-12226 in #11323.

We could also consider blocking for at most half of offset.flush.timeout.ms here and then blocking for the remaining time in commitOffsets if we want to honor the docstring for offset.flush.timeout.ms as strictly as possible:

Maximum number of milliseconds to wait for records to flush and partition offset data to be committed to offset storage before cancelling the process and restoring the offset data to be committed in a future attempt.

But I don't know if that's really necessary since in this case, these offsets have no chance of being committed in a future attempt (at least, not by the same task instance), and the offset commit is taking place on the task's work thread instead of the source task offset commit thread, so there's no worry about squatting on that thread and blocking other tasks from being able to commit offsets while it's in use.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I think I agree. The previous behavior called commitOffsets() on shutdown, and that method blocked for up to offset.flush.timeout.ms anyway. So IMO it makes sense to block up to that amount of time before calling updateCommittableOffsets() and commitOffsets() on the subsequent lines.

TimeUnit.MILLISECONDS
);
// It should still be safe to commit offsets since any exception would have
// simply resulted in not getting more records but all the existing records should be ok to flush
// and commit offsets. Worst case, task.flush() will also throw an exception causing the offset commit
Expand Down Expand Up @@ -356,7 +359,7 @@ private boolean sendRecords() {
}

log.trace("{} Appending record to the topic {} with key {}, value {}", this, record.topic(), record.key(), record.value());
SubmittedRecord submittedRecord = submittedRecords.submit(record);
SubmittedRecords.SubmittedRecord submittedRecord = submittedRecords.submit(record);
try {
maybeCreateTopic(record.topic());
final String topic = producerRecord.topic();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import static org.apache.kafka.connect.runtime.SubmittedRecords.SubmittedRecord;
import static org.apache.kafka.connect.runtime.SubmittedRecords.CommittableOffsets;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
Expand Down Expand Up @@ -86,7 +88,7 @@ public void testNoCommittedRecords() {
public void testSingleAck() {
Map<String, Object> offset = newOffset();

SubmittedRecord submittedRecord = submittedRecords.submit(PARTITION1, offset);
SubmittedRecords.SubmittedRecord submittedRecord = submittedRecords.submit(PARTITION1, offset);
CommittableOffsets committableOffsets = submittedRecords.committableOffsets();
// Record has been submitted but not yet acked; cannot commit offsets for it yet
assertFalse(committableOffsets.isEmpty());
Expand Down Expand Up @@ -117,10 +119,10 @@ public void testMultipleAcksAcrossMultiplePartitions() {
Map<String, Object> partition2Offset1 = newOffset();
Map<String, Object> partition2Offset2 = newOffset();

SubmittedRecord partition1Record1 = submittedRecords.submit(PARTITION1, partition1Offset1);
SubmittedRecord partition1Record2 = submittedRecords.submit(PARTITION1, partition1Offset2);
SubmittedRecord partition2Record1 = submittedRecords.submit(PARTITION2, partition2Offset1);
SubmittedRecord partition2Record2 = submittedRecords.submit(PARTITION2, partition2Offset2);
SubmittedRecords.SubmittedRecord partition1Record1 = submittedRecords.submit(PARTITION1, partition1Offset1);
SubmittedRecords.SubmittedRecord partition1Record2 = submittedRecords.submit(PARTITION1, partition1Offset2);
SubmittedRecords.SubmittedRecord partition2Record1 = submittedRecords.submit(PARTITION2, partition2Offset1);
SubmittedRecords.SubmittedRecord partition2Record2 = submittedRecords.submit(PARTITION2, partition2Offset2);

CommittableOffsets committableOffsets = submittedRecords.committableOffsets();
// No records ack'd yet; can't commit any offsets
Expand Down Expand Up @@ -169,7 +171,7 @@ public void testMultipleAcksAcrossMultiplePartitions() {

@Test
public void testRemoveLastSubmittedRecord() {
SubmittedRecord submittedRecord = submittedRecords.submit(PARTITION1, newOffset());
SubmittedRecords.SubmittedRecord submittedRecord = submittedRecords.submit(PARTITION1, newOffset());

CommittableOffsets committableOffsets = submittedRecords.committableOffsets();
assertEquals(Collections.emptyMap(), committableOffsets.offsets());
Expand All @@ -193,8 +195,8 @@ public void testRemoveNotLastSubmittedRecord() {
Map<String, Object> partition1Offset = newOffset();
Map<String, Object> partition2Offset = newOffset();

SubmittedRecord recordToRemove = submittedRecords.submit(PARTITION1, partition1Offset);
SubmittedRecord lastSubmittedRecord = submittedRecords.submit(PARTITION2, partition2Offset);
SubmittedRecords.SubmittedRecord recordToRemove = submittedRecords.submit(PARTITION1, partition1Offset);
SubmittedRecords.SubmittedRecord lastSubmittedRecord = submittedRecords.submit(PARTITION2, partition2Offset);

CommittableOffsets committableOffsets = submittedRecords.committableOffsets();
assertMetadata(committableOffsets, 0, 2, 2, 1, PARTITION1, PARTITION2);
Expand Down Expand Up @@ -232,7 +234,7 @@ public void testRemoveNotLastSubmittedRecord() {

@Test
public void testNullPartitionAndOffset() {
SubmittedRecord submittedRecord = submittedRecords.submit(null, null);
SubmittedRecords.SubmittedRecord submittedRecord = submittedRecords.submit(null, null);
CommittableOffsets committableOffsets = submittedRecords.committableOffsets();
assertMetadata(committableOffsets, 0, 1, 1, 1, (Map<String, Object>) null);

Expand All @@ -244,6 +246,95 @@ public void testNullPartitionAndOffset() {
assertNoEmptyDeques();
}

@Test
public void testAwaitMessagesNoneSubmitted() {
assertTrue(submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS));
}

@Test
public void testAwaitMessagesAfterAllAcknowledged() {
SubmittedRecords.SubmittedRecord recordToAck = submittedRecords.submit(PARTITION1, newOffset());
assertFalse(submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS));
recordToAck.ack();
assertTrue(submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS));
}

@Test
public void testAwaitMessagesAfterAllRemoved() {
SubmittedRecords.SubmittedRecord recordToRemove1 = submittedRecords.submit(PARTITION1, newOffset());
SubmittedRecords.SubmittedRecord recordToRemove2 = submittedRecords.submit(PARTITION1, newOffset());
assertFalse(
"Await should fail since neither of the in-flight records has been removed so far",
submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS)
);

submittedRecords.removeLastOccurrence(recordToRemove1);
assertFalse(
"Await should fail since only one of the two submitted records has been removed so far",
submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS)
);

submittedRecords.removeLastOccurrence(recordToRemove1);
assertFalse(
"Await should fail since only one of the two submitted records has been removed so far, "
+ "even though that record has been removed twice",
submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS)
);

submittedRecords.removeLastOccurrence(recordToRemove2);
assertTrue(
"Await should succeed since both submitted records have now been removed",
submittedRecords.awaitAllMessages(0, TimeUnit.MILLISECONDS)
);
}

@Test
public void testAwaitMessagesTimesOut() {
submittedRecords.submit(PARTITION1, newOffset());
assertFalse(submittedRecords.awaitAllMessages(10, TimeUnit.MILLISECONDS));
}

@Test
public void testAwaitMessagesReturnsAfterAsynchronousAck() throws Exception {
SubmittedRecords.SubmittedRecord inFlightRecord1 = submittedRecords.submit(PARTITION1, newOffset());
SubmittedRecords.SubmittedRecord inFlightRecord2 = submittedRecords.submit(PARTITION2, newOffset());

AtomicBoolean awaitResult = new AtomicBoolean();
CountDownLatch awaitComplete = new CountDownLatch(1);
new Thread(() -> {
awaitResult.set(submittedRecords.awaitAllMessages(5, TimeUnit.SECONDS));
awaitComplete.countDown();
}).start();

assertTrue(
"Should not have finished awaiting message delivery before either in-flight record was acknowledged",
awaitComplete.getCount() > 0
);

inFlightRecord1.ack();
assertTrue(
"Should not have finished awaiting message delivery before one in-flight record was acknowledged",
awaitComplete.getCount() > 0
);

inFlightRecord1.ack();
assertTrue(
"Should not have finished awaiting message delivery before one in-flight record was acknowledged, "
+ "even though the other record has been acknowledged twice",
awaitComplete.getCount() > 0
);

inFlightRecord2.ack();
assertTrue(
"Should have finished awaiting message delivery after both in-flight records were acknowledged",
awaitComplete.await(1, TimeUnit.SECONDS)
);
assertTrue(
"Await of in-flight messages should have succeeded",
awaitResult.get()
);
}

private void assertNoRemainingDeques() {
assertEquals("Internal records map should be completely empty", Collections.emptyMap(), submittedRecords.records);
}
Expand Down