From f68a4612c8654d70a8e205fd535e75c203f4b833 Mon Sep 17 00:00:00 2001
From: Jon Meredith <jmeredithco@gmail.com>
Date: Thu, 18 Oct 2018 09:56:46 -0600
Subject: [PATCH] In BufferPool, make allocating thread receive a Chunk.

Prevents a condition where the thread allocating the macro chunk
could have all of the chunks taken from the queue before it
is able to use one.

For CASSANDRA-14832
---
 .../cassandra/utils/memory/BufferPool.java    | 31 ++++++++++---------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/java/org/apache/cassandra/utils/memory/BufferPool.java b/src/java/org/apache/cassandra/utils/memory/BufferPool.java
index f9720594e908..3c876a49077a 100644
--- a/src/java/org/apache/cassandra/utils/memory/BufferPool.java
+++ b/src/java/org/apache/cassandra/utils/memory/BufferPool.java
@@ -237,23 +237,23 @@ static final class GlobalPool
         /** Return a chunk, the caller will take owership of the parent chunk. */
         public Chunk get()
         {
-            while (true)
-            {
-                Chunk chunk = chunks.poll();
-                if (chunk != null)
-                    return chunk;
+            Chunk chunk = chunks.poll();
+            if (chunk != null)
+                return chunk;
 
-                if (!allocateMoreChunks())
-                    // give it one last attempt, in case someone else allocated before us
-                    return chunks.poll();
-            }
+            chunk = allocateMoreChunks();
+            if (chunk != null)
+                return chunk;
+
+            // another thread may have just allocated last macro chunk, so make one final attempt before returning null
+            return chunks.poll();
         }
 
         /**
          * This method might be called by multiple threads and that's fine if we add more
          * than one chunk at the same time as long as we don't exceed the MEMORY_USAGE_THRESHOLD.
          */
-        private boolean allocateMoreChunks()
+        private Chunk allocateMoreChunks()
         {
             while (true)
             {
@@ -262,7 +262,7 @@ private boolean allocateMoreChunks()
                 {
                     noSpamLogger.info("Maximum memory usage reached ({} bytes), cannot allocate chunk of {} bytes",
                                       MEMORY_USAGE_THRESHOLD, MACRO_CHUNK_SIZE);
-                    return false;
+                    return null;
                 }
                 if (memoryUsage.compareAndSet(cur, cur + MACRO_CHUNK_SIZE))
                     break;
@@ -272,15 +272,18 @@ private boolean allocateMoreChunks()
             Chunk chunk = new Chunk(allocateDirectAligned(MACRO_CHUNK_SIZE));
             chunk.acquire(null);
             macroChunks.add(chunk);
-            for (int i = 0 ; i < MACRO_CHUNK_SIZE ; i += CHUNK_SIZE)
+
+            final Chunk callerChunk = new Chunk(chunk.get(CHUNK_SIZE));
+            if (DEBUG)
+                debug.register(callerChunk);
+            for (int i = CHUNK_SIZE ; i < MACRO_CHUNK_SIZE; i += CHUNK_SIZE)
             {
                 Chunk add = new Chunk(chunk.get(CHUNK_SIZE));
                 chunks.add(add);
                 if (DEBUG)
                     debug.register(add);
             }
-
-            return true;
+            return callerChunk;
         }
 
         public void recycle(Chunk chunk)