Protect the producer index in case of OutOfMemoryError

JCTools · Apr 3, 2019 · d64b9b4 · d64b9b4
1 parent 1e638a6
commit d64b9b4
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/jctools-core/src/main/java/org/jctools/queues/BaseMpscLinkedArrayQueue.java b/jctools-core/src/main/java/org/jctools/queues/BaseMpscLinkedArrayQueue.java
@@ -260,7 +260,7 @@ public boolean offer(final E e)
                     case QUEUE_FULL:
                         return false;
                     case QUEUE_RESIZE:
-                        resize(mask, buffer, pIndex, e);
+                        resize(mask, buffer, pIndex, e, null);
                         return true;
                 }
             }
@@ -563,7 +563,7 @@ public int fill(Supplier<E> s, int batchSize)
                     case QUEUE_FULL:
                         return 0;
                     case QUEUE_RESIZE:
-                        resize(mask, buffer, pIndex, s.get());
+                        resize(mask, buffer, pIndex, null, s);
                         return 1;
                 }
             }
@@ -641,10 +641,21 @@ public void drain(Consumer<E> c, WaitStrategy w, ExitCondition exit)
         }
     }
 
-    private void resize(long oldMask, E[] oldBuffer, long pIndex, E e)
+    private void resize(long oldMask, E[] oldBuffer, long pIndex, E e, Supplier<E> s)
     {
+        assert (e != null && s == null) || (e == null || s != null);
         int newBufferLength = getNextBufferSize(oldBuffer);
-        final E[] newBuffer = allocate(newBufferLength);
+        final E[] newBuffer;
+        try
+        {
+            newBuffer = allocate(newBufferLength);
+        }
+        catch (OutOfMemoryError oom)
+        {
+            assert lvProducerIndex() == pIndex + 1;
+            soProducerIndex(pIndex);
+            throw oom;
+        }
 
         producerBuffer = newBuffer;
         final int newMask = (newBufferLength - 2) << 1;
@@ -653,7 +664,7 @@ private void resize(long oldMask, E[] oldBuffer, long pIndex, E e)
         final long offsetInOld = modifiedCalcElementOffset(pIndex, oldMask);
         final long offsetInNew = modifiedCalcElementOffset(pIndex, newMask);
 
-        soElement(newBuffer, offsetInNew, e);// element in new array
+        soElement(newBuffer, offsetInNew, e == null ? s.get() : e);// element in new array
         soElement(oldBuffer, nextArrayOffset(oldMask), newBuffer);// buffer linked
 
         // ASSERT code