fix(backPressure): fix a race condition that may block enqueue.

PhantomThief · Oct 12, 2020 · dd54b63 · dd54b63
1 parent a4d4dee
commit dd54b63
Showing 1 changed file with 22 additions and 17 deletions.
diff --git a/src/main/java/com/github/phantomthief/collection/impl/SimpleBufferTrigger.java b/src/main/java/com/github/phantomthief/collection/impl/SimpleBufferTrigger.java
@@ -139,7 +139,23 @@ public void enqueue(E element) {
         long currentCount = counter.get();
         long thisMaxBufferCount = maxBufferCount.getAsLong();
         if (thisMaxBufferCount > 0 && currentCount >= thisMaxBufferCount) {
-            boolean pass = fireRejectHandler(element);
+            boolean pass = false;
+            if (rejectHandler != null) {
+                if (writeLock != null && writeCondition != null) {
+                    writeLock.lock();
+                }
+                try {
+                    currentCount = counter.get();
+                    thisMaxBufferCount = maxBufferCount.getAsLong();
+                    if (thisMaxBufferCount > 0 && currentCount >= thisMaxBufferCount) {
+                        pass = fireRejectHandler(element);
+                    }
+                } finally {
+                    if (writeLock != null && writeCondition != null) {
+                        writeLock.unlock();
+                    }
+                }
+            }
             if (!pass) {
                 return;
             }
@@ -167,23 +183,12 @@ public void enqueue(E element) {
     }
 
     private boolean fireRejectHandler(E element) {
-        boolean pass = false;
-        if (rejectHandler != null) {
-            if (writeLock != null && writeCondition != null) {
-                writeLock.lock();
-            }
-            try {
-                pass = rejectHandler.onReject(element, writeCondition);
-            } catch (Throwable e) {
-                throwIfUnchecked(e);
-                throw new RuntimeException(e);
-            } finally {
-                if (writeLock != null && writeCondition != null) {
-                    writeLock.unlock();
-                }
-            }
+        try {
+            return rejectHandler.onReject(element, writeCondition);
+        } catch (Throwable e) {
+            throwIfUnchecked(e);
+            throw new RuntimeException(e);
         }
-        return pass;
     }
 
     @Override