Repository: apache/kafka
Pull request #19759: KAFKA-19312: Avoiding concurrent execution of onComplete and tryComplete
Author: login: apoorvmittal10
Labels: name: core
description: Kafka Broker
name: KIP-932
description: Queues for Kafka
name: ci-approved
name: small
description: Small PRs
Created: 2025-05-19T15:35:10Z | Merged: 2025-05-25T06:36:43Z

The `onComplete` method in DelayedOperation is guaranteed to run only
once, through `forceComplete`, invoked either by `tryComplete` or when
operation is expired (`run` method). The invocation of  `tryComplete` is
done by attaining `lock` so no concurrent execution of  `tryComplete`
happens for same delayed operation. However, there can be  concurrent
execution of `tryComplete` and `onComplete` as the  `expiration` thread
can trigger a separte run of `onComplete` while  `tryComplete` is still
executing. This behaviour is not ideal as there  are parallel runs where
1 threads method execution is wasteful i.e. if  `onComplete` is already
invoked by another thread then execution of  `tryComplete` is not
required.

Reviewers: Jun Rao <junrao@gmail.com>, Andrew Schofield
 <aschofield@confluent.io>, Chia-Ping Tsai <chia7712@gmail.com>



Diff:
diff --git a/core/src/main/java/kafka/server/share/DelayedShareFetch.java b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
index 0a46e834eb7f5..8446b7acad6c9 100644
--- a/core/src/main/java/kafka/server/share/DelayedShareFetch.java
+++ b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
@@ -202,32 +202,29 @@ public void onExpiration() {
      * Complete the share fetch operation by fetching records for all partitions in the share fetch request irrespective
      * of whether they have any acquired records. This is called when the fetch operation is forced to complete either
      * because records can be acquired for some partitions or due to MaxWaitMs timeout.
+     * <p>
+     * On operation timeout, onComplete is invoked, last try occurs to acquire partitions and read
+     * from log, if acquired. The fetch will only happen from local log and not remote storage, on
+     * operation expiration.
      */
     @Override
     public void onComplete() {
-        // We are utilizing lock so that onComplete doesn't do a dirty read for instance variables -
-        // partitionsAcquired and localPartitionsAlreadyFetched, since these variables can get updated in a different tryComplete thread.
-        lock.lock();
         log.trace("Completing the delayed share fetch request for group {}, member {}, "
             + "topic partitions {}", shareFetch.groupId(), shareFetch.memberId(),
             partitionsAcquired.keySet());
 
-        try {
-            if (remoteStorageFetchException.isPresent()) {
-                completeErroneousRemoteShareFetchRequest();
-            } else if (pendingRemoteFetchesOpt.isPresent()) {
-                if (maybeRegisterCallbackPendingRemoteFetch()) {
-                    log.trace("Registered remote storage fetch callback for group {}, member {}, "
-                            + "topic partitions {}", shareFetch.groupId(), shareFetch.memberId(),
-                        partitionsAcquired.keySet());
-                    return;
-                }
-                completeRemoteStorageShareFetchRequest();
-            } else {
-                completeLocalLogShareFetchRequest();
+        if (remoteStorageFetchException.isPresent()) {
+            completeErroneousRemoteShareFetchRequest();
+        } else if (pendingRemoteFetchesOpt.isPresent()) {
+            if (maybeRegisterCallbackPendingRemoteFetch()) {
+                log.trace("Registered remote storage fetch callback for group {}, member {}, "
+                        + "topic partitions {}", shareFetch.groupId(), shareFetch.memberId(),
+                    partitionsAcquired.keySet());
+                return;
             }
-        } finally {
-            lock.unlock();
+            completeRemoteStorageShareFetchRequest();
+        } else {
+            completeLocalLogShareFetchRequest();
         }
     }
 
@@ -358,7 +355,7 @@ public boolean tryComplete() {
                 if (anyPartitionHasLogReadError(replicaManagerReadResponse) || isMinBytesSatisfied(topicPartitionData, partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes, topicPartitionData.keySet(), topicPartitionData.size()))) {
                     partitionsAcquired = topicPartitionData;
                     localPartitionsAlreadyFetched = replicaManagerReadResponse;
-                    return forceCompleteRequest();
+                    return forceComplete();
                 } else {
                     log.debug("minBytes is not satisfied for the share fetch request for group {}, member {}, " +
                             "topic partitions {}", shareFetch.groupId(), shareFetch.memberId(),
@@ -366,7 +363,7 @@ public boolean tryComplete() {
                     releasePartitionLocks(topicPartitionData.keySet());
                 }
             } else {
-                log.trace("Can't acquire records for any partition in the share fetch request for group {}, member {}, " +
+                log.trace("Can't acquire any partitions in the share fetch request for group {}, member {}, " +
                         "topic partitions {}", shareFetch.groupId(), shareFetch.memberId(),
                     sharePartitions.keySet());
             }
@@ -381,10 +378,8 @@ public boolean tryComplete() {
                 releasePartitionLocks(topicPartitionData.keySet());
                 partitionsAcquired.clear();
                 localPartitionsAlreadyFetched.clear();
-                return forceCompleteRequest();
-            } else {
-                return forceCompleteRequest();
             }
+            return forceComplete();
         }
     }
 
@@ -785,7 +780,7 @@ private boolean maybeCompletePendingRemoteFetch() {
         }
 
         if (canComplete || pendingRemoteFetchesOpt.get().isDone()) { // Case d
-            return forceCompleteRequest();
+            return forceComplete();
         } else
             return false;
     }
@@ -944,16 +939,6 @@ private void cancelRemoteFetchTask(RemoteFetch remoteFetch) {
         }
     }
 
-    private boolean forceCompleteRequest() {
-        boolean completedByMe = forceComplete();
-        // If the delayed operation is completed by me, the acquired locks are already released in onComplete().
-        // Otherwise, we need to release the acquired locks.
-        if (!completedByMe) {
-            releasePartitionLocksAndAddToActionQueue(partitionsAcquired.keySet());
-        }
-        return completedByMe;
-    }
-
     private void completeRemoteShareFetchRequestOutsidePurgatory() {
         if (outsidePurgatoryCallbackLock.compareAndSet(false, true)) {
             completeRemoteStorageShareFetchRequest();
diff --git a/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java b/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
index f3c818cb9c6c6..16bfdabdb67a3 100644
--- a/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
+++ b/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.server.util.timer.TimerTask;
 
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -29,10 +28,9 @@
  * a delayed fetch operation could be waiting for a given number of bytes to accumulate.
  * <br/>
  * The logic upon completing a delayed operation is defined in onComplete() and will be called exactly once.
- * Once an operation is completed, isCompleted() will return true. onComplete() can be triggered by either
- * forceComplete(), which forces calling onComplete() after delayMs if the operation is not yet completed,
- * or tryComplete(), which first checks if the operation can be completed or not now, and if yes calls
- * forceComplete().
+ * Once an operation is completed, isCompleted() will return true. onComplete() is called from forceComplete(),
+ * which is triggered by either expiration, if the operation is not completed after delayMs; or tryComplete(),
+ * if the operation can be completed now.
  * <br/>
  * A subclass of DelayedOperation needs to provide an implementation of both onComplete() and tryComplete().
  * <br/>
@@ -41,7 +39,7 @@
  */
 public abstract class DelayedOperation extends TimerTask {
 
-    private final AtomicBoolean completed = new AtomicBoolean(false);
+    private volatile boolean completed = false;
 
     protected final Lock lock;
 
@@ -68,24 +66,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
+        try {
+            // Re-check, if the operation is already completed by some other thread.
+            if (!completed) {
+                completed = true;
+                // cancel the timeout timer
+                cancel();
+                onComplete();
+                return true;
+            } else {
+                return false;
+            }
+        } finally {
+            lock.unlock();
+        }
     }
 
     /**
      * Check if the delayed operation is already completed
      */
     public boolean isCompleted() {
-        return completed.get();
+        return completed;
     }
 
     /**


Reviews:
user: login: junrao
body: @apoorvmittal10 : Thanks for the PR. The code looks reasonable to me. Could you run TestPurgatoryPerformance to see if there is any perf degradation?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T23:09:27Z
user: login: AndrewJSchofield
body: Thanks for the PR. This was a tricky bug to find, so nice catch.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T08:56:52Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T09:58:48Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T09:58:58Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T10:03:03Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T13:20:28Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-21T17:51:28Z
user: login: junrao
body: @apoorvmittal10 : Thanks for the explanation. It's clear to me why this is a bug now. A couple of comments.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T00:17:29Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T06:33:16Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T06:34:30Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T06:38:59Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T07:09:10Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T08:40:33Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T09:24:45Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T12:37:52Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T12:47:50Z
user: login: junrao
body: @apoorvmittal10 : Thanks for the updated PR. A couple of more comments.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T18:02:40Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T08:24:10Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T08:33:29Z
user: login: junrao
body: @apoorvmittal10 : Thanks for the updated PR. Just a minor comment.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T17:45:23Z
user: login: apoorvmittal10
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T21:57:18Z
user: login: junrao
body: @apoorvmittal10 : Thanks for the updated PR. LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-24T23:43:27Z
user: login: chia7712
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-25T06:34:23Z

Review comments:
diff_hunk: @@ -191,26 +191,23 @@ public void onExpiration() {
      * Complete the share fetch operation by fetching records for all partitions in the share fetch request irrespective
      * of whether they have any acquired records. This is called when the fetch operation is forced to complete either
      * because records can be acquired for some partitions or due to MaxWaitMs timeout.
+     * <p>
+     * On operation timeout, onComplete is invoked, last try occurs to acquire partitions and read
+     * from log, if acquired. The fetch will only happen from local log and not tier log, on operation
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: AndrewJSchofield
body: nit: I'd say "remote storage" rather than "tier log".
created_at: 2025-05-20T08:53:22Z
author_association: MEMBER
diff_hunk: @@ -349,7 +346,7 @@ public boolean tryComplete() {
                     releasePartitionLocks(topicPartitionData.keySet());
                 }
             } else {
-                log.trace("Can't acquire records for any partition in the share fetch request for group {}, member {}, " +
+                log.trace("Can't acquire any partition in the share fetch request for group {}, member {}, " +
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: AndrewJSchofield
body: nit: "any partitions"
created_at: 2025-05-20T08:56:26Z
author_association: MEMBER
diff_hunk: @@ -349,7 +346,7 @@ public boolean tryComplete() {
                     releasePartitionLocks(topicPartitionData.keySet());
                 }
             } else {
-                log.trace("Can't acquire records for any partition in the share fetch request for group {}, member {}, " +
+                log.trace("Can't acquire any partition in the share fetch request for group {}, member {}, " +
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: apoorvmittal10
body: Done.
created_at: 2025-05-20T09:58:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,26 +191,23 @@ public void onExpiration() {
      * Complete the share fetch operation by fetching records for all partitions in the share fetch request irrespective
      * of whether they have any acquired records. This is called when the fetch operation is forced to complete either
      * because records can be acquired for some partitions or due to MaxWaitMs timeout.
+     * <p>
+     * On operation timeout, onComplete is invoked, last try occurs to acquire partitions and read
+     * from log, if acquired. The fetch will only happen from local log and not tier log, on operation
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: apoorvmittal10
body: Done.
created_at: 2025-05-20T09:58:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: chia7712
body: After this change, the lock must be reentrant, correct? If so, we should consider replacing the `Lock` type with `ReentrantLock` to prevent potential deadlocks.

Additionally, `DelayedOperation(long delayMs, Optional<Lock> lockOpt)` seems to lose its meaning once the old group coordinator is removed. We might want to clean that up.

These points can be addressed in a follow-up.
created_at: 2025-05-20T10:03:03Z
author_association: MEMBER
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: The default implemenation of Lock is over ReentrantLock itself in DelayedOperation and elsewhere, abstracted over Lock API in DelayedOperation. Are you suggesting to make it more tighter and use ReentrantLock explicitly?
created_at: 2025-05-21T13:20:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: chia7712
body: > Are you suggesting to make it more tighter and use ReentrantLock explicitly?

yes, but this change needs to cleanup the code. It can be addressed later if we want to merge this bug fix first.
created_at: 2025-05-21T17:51:28Z
author_association: MEMBER
diff_hunk: @@ -191,26 +191,23 @@ public void onExpiration() {
      * Complete the share fetch operation by fetching records for all partitions in the share fetch request irrespective
      * of whether they have any acquired records. This is called when the fetch operation is forced to complete either
      * because records can be acquired for some partitions or due to MaxWaitMs timeout.
+     * <p>
+     * On operation timeout, onComplete is invoked, last try occurs to acquire partitions and read
+     * from log, if acquired. The fetch will only happen from local log and not remote storage, on
+     * operation expiration.
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: junrao
body: For better consistency, should we remove released partitions from `partitionsAcquired` in `releasePartitionLocks()`? It seems that if we had done that, we wouldn't get into the situation of releasing the lock more than once.
created_at: 2025-05-22T00:13:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: junrao
body: An implication of this change is that when forceComplete() is called inside `tryComplete`, it always returns true. Currently, we have multiple places that check the return value of forceComplete() inside `tryComplete`. Should we adjust those?

created_at: 2025-05-22T00:13:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,26 +191,23 @@ public void onExpiration() {
      * Complete the share fetch operation by fetching records for all partitions in the share fetch request irrespective
      * of whether they have any acquired records. This is called when the fetch operation is forced to complete either
      * because records can be acquired for some partitions or due to MaxWaitMs timeout.
+     * <p>
+     * On operation timeout, onComplete is invoked, last try occurs to acquire partitions and read
+     * from log, if acquired. The fetch will only happen from local log and not remote storage, on
+     * operation expiration.
path: core/src/main/java/kafka/server/share/DelayedShareFetch.java
user: login: apoorvmittal10
body: I thought about it but it complicates the code execution in DelayedShareFetch. Clearing partitionsAcquired alone is not sufficient, as there could be pending remote fetch as well which might need additional time to finish. Hence I didn't go with that approach.
created_at: 2025-05-22T06:33:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: Sure, I ll create a folow up.
created_at: 2025-05-22T06:34:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: I can revert it to AtomicBoolean#compareAndSet but I didn't understand the problem. There would only be first execution of forceComplete which should be successful and return true. Subsequent executions will return false.
created_at: 2025-05-22T06:38:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: chia7712
body: @junrao that is a good point. After avoiding concurrent execution, `forceComplete` MUST return true if it is executed by `tryComplete`. Hence, checking the returned value of `forceComplete` gets meaningless.
created_at: 2025-05-22T07:09:10Z
author_association: MEMBER
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: Ok, so are we saying that first execution of forceComplete from tryComplete will return true? And as the operation is marked completed hence no further execution of tryComplete can trigger, which makes all executions of forceComplete within tryComplete always returns true?

I have removed the check in DelayedShareFetch.
created_at: 2025-05-22T08:40:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: chia7712
body: > Ok, so are we saying that first execution of forceComplete from tryComplete will return true? And as the operation is marked completed hence no further execution of tryComplete can trigger, which makes all executions of forceComplete within tryComplete always returns true?

yes, we are on the same page. I hope I'm not misunderstanding @junrao comment :)
created_at: 2025-05-22T09:24:45Z
author_association: MEMBER
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: chia7712
body: I was a little bored, so I created Jira KAFKA-19322 (https://issues.apache.org/jira/browse/KAFKA-19322). Feel free to take it over, or I can assign it to a mentee :smile: 
created_at: 2025-05-22T12:37:52Z
author_association: MEMBER
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
+        if (completed) {
             return false;
         }
+        // Attain lock prior completing the request.
+        lock.lock();
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: You are awesome, feel free to assign to a mentee and I can review or else I will pick it up.
created_at: 2025-05-22T12:47:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: junrao
body: Only the expiration logic checks the return value of forceComplete(). We could do the following in the expiration logic and change `forceComplete()` to return void and avoid using the lock there (the caller will get the lock instead). 

```
        lock.lock();
        try {
            if (!isCompleted()) {
              forceComplete();
              onExpiration();
             }
        } finally {
            lock.unlock();
        }

```
 
created_at: 2025-05-22T17:35:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,7 +40,7 @@
  */
 public abstract class DelayedOperation extends TimerTask {
 
-    private final AtomicBoolean completed = new AtomicBoolean(false);
+    private volatile boolean completed = false;
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: junrao
body: > onComplete() can be triggered by either forceComplete(), which forces calling onComplete() after delayMs if the operation is not yet completed, or tryComplete(), which first checks if the operation can be completed or not now, and if yes calls forceComplete().

Could we adjust the comment to sth like the following?

onComplete() is called from forceComplete(), which is triggered once by either expiration if the operation is not completed after delayMs, or tryComplete() if the operation can be completed now.
created_at: 2025-05-22T18:01:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: > and change forceComplete() to return void

Though I agree that now if forceComplete is invoked by tryComplete then it will return true but tryComplete still need to return boolean hence all places where we return the value from tryComplete based on forceComplete has to be changed and also requires explicit `return true` from all the places. I am of the opinion to not do this. Or if you think we should then shall we take it as a separate PR for just refactor?
created_at: 2025-05-23T08:24:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,7 +40,7 @@
  */
 public abstract class DelayedOperation extends TimerTask {
 
-    private final AtomicBoolean completed = new AtomicBoolean(false);
+    private volatile boolean completed = false;
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: Done.
created_at: 2025-05-23T08:33:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,24 +67,36 @@ public DelayedOperation(long delayMs, Lock lock) {
      * Return true iff the operation is completed by the caller: note that
      * concurrent threads can try to complete the same operation, but only
      * the first thread will succeed in completing the operation and return
-     * true, others will still return false
+     * true, others will still return false.
      */
     public boolean forceComplete() {
-        if (completed.compareAndSet(false, true)) {
-            // cancel the timeout timer
-            cancel();
-            onComplete();
-            return true;
-        } else {
+        // Do not proceed if the operation is already completed.
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: junrao
body: Well, the issue is that there are quite a few places where tryComplete() calls forceComplete() and checks the return value of forceComplete(). You fixed the issue in DelayedShareFetch, but not in other places. As we evolve the code, it would be useful to leave the code base in a better place for future developers, even though it takes a bit of more work. This can be done in a followup jira. Filed https://issues.apache.org/jira/browse/KAFKA-19325 to track it.
created_at: 2025-05-23T17:43:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -29,10 +28,9 @@
  * a delayed fetch operation could be waiting for a given number of bytes to accumulate.
  * <br/>
  * The logic upon completing a delayed operation is defined in onComplete() and will be called exactly once.
- * Once an operation is completed, isCompleted() will return true. onComplete() can be triggered by either
- * forceComplete(), which forces calling onComplete() after delayMs if the operation is not yet completed,
- * or tryComplete(), which first checks if the operation can be completed or not now, and if yes calls
- * forceComplete().
+ * Once an operation is completed, isCompleted() will return true. onComplete() is called from forceComplete(),
+ * which is triggered once by either expiration if the operation is not completed after delayMs, or tryComplete()
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: junrao
body: Actually, it's possible for forceComplete() to be called once from tryComplete() and another from expiration. So, we need to remove `once`. 
created_at: 2025-05-23T17:44:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -29,10 +28,9 @@
  * a delayed fetch operation could be waiting for a given number of bytes to accumulate.
  * <br/>
  * The logic upon completing a delayed operation is defined in onComplete() and will be called exactly once.
- * Once an operation is completed, isCompleted() will return true. onComplete() can be triggered by either
- * forceComplete(), which forces calling onComplete() after delayMs if the operation is not yet completed,
- * or tryComplete(), which first checks if the operation can be completed or not now, and if yes calls
- * forceComplete().
+ * Once an operation is completed, isCompleted() will return true. onComplete() is called from forceComplete(),
+ * which is triggered once by either expiration if the operation is not completed after delayMs, or tryComplete()
path: server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java
user: login: apoorvmittal10
body: Done.
created_at: 2025-05-23T21:57:18Z
author_association: CONTRIBUTOR

PR comments:
user: login: chia7712
created_at: 2025-05-19T19:22:45Z
body: > The onComplete method in DelayedOperation is guaranteed to run only
once, either by tryComplete or onExpiration. 

`onExpiration`? or `forceComplete`?

> if
onComplete is already invoked by another thread then execution of
tryComplete is not required.

Excuse me, how to address it? do you plan to check `isCompleted` before executing `tryComplete`?
author_association: MEMBER
user: login: chia7712
created_at: 2025-05-20T09:09:37Z
body: @apoorvmittal10 It is too cool to add response in comment, so I don't notice your response :)

> For 2) tryComplete invocation is in lock itself and there is also a checck for isCompleted prior invocation: 

Should we update `safeTryCompleteOrElse` to ensure it check for `isCompleted` prior invocation.
author_association: MEMBER
user: login: apoorvmittal10
created_at: 2025-05-20T09:53:47Z
body: > @apoorvmittal10 It is too cool to add response in comment, so I don't notice your response :)
> 
> > For 2) tryComplete invocation is in lock itself and there is also a checck for isCompleted prior invocation:
> 
> Should we update `safeTryCompleteOrElse` to ensure it check for `isCompleted` prior invocation.

@chia7712 Ahhhh, my bad. I clicked edit instead of reply 🫣🤦‍♂️ My bad.

safeTryCompleteOrElse is invoked when the operation is being added to purgatory, if tryComplete returns false then it gets added in purgatory. An operation should never have completed true when invoking safeTryCompleteOrElse.
author_association: CONTRIBUTOR
user: login: apoorvmittal10
created_at: 2025-05-20T09:56:19Z
body: > > The onComplete method in DelayedOperation is guaranteed to run only
> > once, either by tryComplete or onExpiration.
> 
> `onExpiration`? or `forceComplete`?
> 
> > if
> > onComplete is already invoked by another thread then execution of
> > tryComplete is not required.
> 
> Excuse me, how to address it? do you plan to check `isCompleted` before executing `tryComplete`?

For 1) I meant the invocation by run method, I have updated the description. Thanks.

For 2) tryComplete invocation is in `lock` itself and there is also a checck for `isCompleted` prior invocation: https://github.com/apache/kafka/blob/d3f897948630d6f387969cdc9542fcabb6f1d2a9/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java#L136

@chia7712 just rewriting, so we have corrected flow of comments.
author_association: CONTRIBUTOR
user: login: junrao
created_at: 2025-05-20T16:42:50Z
body: > This was a tricky bug to find, so nice catch.

@apoorvmittal10 : My understanding is that this PR is not fixing a bug. Without the PR, we still guarantee that onComplete() will be called exactly once. So, this PR just saves some unnecessary work after the operation completes. Is that correct?
author_association: CONTRIBUTOR
user: login: apoorvmittal10
created_at: 2025-05-21T09:28:27Z
body: > > This was a tricky bug to find, so nice catch.
> 
> @apoorvmittal10 : My understanding is that this PR is not fixing a bug. Without the PR, we still guarantee that onComplete() will be called exactly once. So, this PR just saves some unnecessary work after the operation completes. Is that correct?

It does both, without the PR, though we guarantee that onComplete will be called exactly once but there can be a thread which is inside onComplete (when timeout occurs) and a separate thread invoking tryComplete. For DelayedShareFetch, prior this PR, a thread can be waiting in onComplete for lock which has already marked `completed` true in DelayedOperation but parallel thread in `tryComplete` is running, which acquires some partitions. Now the thread in tryComplete cannot invoke `onComplete` as `completed` is marked `true` in DelayedOperation, hence acquired partitions will be released. Then other thread which was waiting in onComplete for the lock, starts executing and releases the same acquired topic partitions again. This PR fixes this issue as well.


author_association: CONTRIBUTOR
user: login: apoorvmittal10
created_at: 2025-05-21T09:31:23Z
body: > @apoorvmittal10 : Thanks for the PR. The code looks reasonable to me. Could you run TestPurgatoryPerformance to see if there is any perf degradation?

I ran some tests and performance is same.

### After the chages:

```
--num 10000 --rate 100 --timeout 1000 --pct50 0.5 --pct75 0.75

# latency samples: pct75 = 0, pct50 = 0, min = 0, max = 7
# interval samples: rate = 100.068948, min = 0, max = 129
# enqueue rate (10000 requests):
# <elapsed time ms>	<target rate>	<actual rate>	<process cpu time ms>	<G1 Old Generation count> <G1 Young Generation count>	<G1 Old Generation time ms> <G1 Young Generation time ms>
101196	99.809364	99.806376	3240	0 2	0 8
```

```
--num 10000 --rate 1000 --timeout 1000 --pct50 0.5 --pct75 0.75

# latency samples: pct75 = 0, pct50 = 0, min = 0, max = 9
# interval samples: rate = 999.371395, min = 0, max = 14
# enqueue rate (10000 requests):
# <elapsed time ms>	<target rate>	<actual rate>	<process cpu time ms>	<G1 Old Generation count> <G1 Young Generation count>	<G1 Old Generation time ms> <G1 Young Generation time ms>
11104	989.902990	989.805008	1349	0 2	0 7
```

### Before changes:

```
--num 10000 --rate 100 --timeout 1000 --pct50 0.5 --pct75 0.75

# latency samples: pct75 = 0, pct50 = 0, min = 0, max = 9
# interval samples: rate = 100.020304, min = 0, max = 130
# enqueue rate (10000 requests):
# <elapsed time ms>	<target rate>	<actual rate>	<process cpu time ms>	<G1 Old Generation count> <G1 Young Generation count>	<G1 Old Generation time ms> <G1 Young Generation time ms>
102366	98.657274	98.652408	3444	0 2	0 8

--num 10000 --rate 1000 --timeout 1000 --pct50 0.5 --pct75 0.75

# latency samples: pct75 = 0, pct50 = 0, min = 0, max = 8
# interval samples: rate = 997.134236, min = 0, max = 14
# enqueue rate (10000 requests):
# <elapsed time ms>	<target rate>	<actual rate>	<process cpu time ms>	<G1 Old Generation count> <G1 Young Generation count>	<G1 Old Generation time ms> <G1 Young Generation time ms>
11218	978.665101	978.665101	1624	0 2	0 7
author_association: CONTRIBUTOR
user: login: junrao
created_at: 2025-05-21T16:25:53Z
body: > Now the thread in tryComplete cannot invoke onComplete as completed is marked true in DelayedOperation, hence acquired partitions will be released. 

Hmm, is this true? tryComplete will only release the acquired locks if forceComplete() returns true, right?
author_association: CONTRIBUTOR
user: login: apoorvmittal10
created_at: 2025-05-21T16:36:03Z
body: > > Now the thread in tryComplete cannot invoke onComplete as completed is marked true in DelayedOperation, hence acquired partitions will be released.
> 
> Hmm, is this true? tryComplete will only release the acquired locks if forceComplete() returns true, right?

Yes it's true, it can happen in current implementation. [This](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/core/src/main/java/kafka/server/share/DelayedShareFetch.java#L902) line when invoked from tryComplete can release the locks, which is correct but conflicts with concurrent execution on request timeout. To make sure that our acquired share partitions release is correct, we added a check [here](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/core/src/main/java/kafka/server/share/SharePartition.java#L1384). In our performance runs we saw the issue. However it will not cause issue per se for share fetches, as SharePartition will not acquire record offsets which are already acquired but it will have some wasted fetches.
author_association: CONTRIBUTOR
user: login: junrao
created_at: 2025-05-21T18:00:40Z
body: @apoorvmittal10 : Thanks for the explanation. It seems that we found a case that an acquired sharePartition lock was released more than once. I am still not sure how it happened. forceComplete() is supposed to only return true once and the lock is only released when forceComplete() returns true. How is the acquired lock released a second time?
author_association: CONTRIBUTOR
user: login: apoorvmittal10
created_at: 2025-05-21T22:52:19Z
body: > @apoorvmittal10 : Thanks for the explanation. It seems that we found a case that an acquired sharePartition lock was released more than once. I am still not sure how it happened. forceComplete() is supposed to only return true once and the lock is only released when forceComplete() returns true. How is the acquired lock released a second time?

Prior to this PR there is no lock in [forceComplete](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java#L73) but a way to make sure that onComplete is not executed twice. However, consider that thread 1 has started executing [safeTryComplete](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java#L133)  which attains lock and checks that if the operation has marked completed, if not then tryComplete is executed. Now thread 1 is executing tryComplete, then thread 2 on timeout executes [forceComplete](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/server-common/src/main/java/org/apache/kafka/server/purgatory/DelayedOperation.java#L148) which marks completed as true and waits [here](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/core/src/main/java/kafka/server/share/DelayedShareFetch.java#L199) for the operation lock. Now thread 1, in tryComplete, acquires some share partitions but as [completed](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/core/src/main/java/kafka/server/share/DelayedShareFetch.java#L899) is already marked true then the share partitions are released and so the lock for delayed operation. Thread 2 now starts executing `onComplete` but as [partitionsAcquired](https://github.com/apache/kafka/blob/e88c10d5952c464a0f597b295deaa51e4c9fa978/core/src/main/java/kafka/server/share/DelayedShareFetch.java#L106) is at instance level hence `onComplete` tries to execute further with those share partitions which are already released. Once completed then onCompete releases those share partitions again.
author_association: CONTRIBUTOR