BufferPool { public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException { if (size > this.totalMemory) { throw new IllegalArgumentException("Attempt to allocate " + size + " bytes, but there is a hard limit of " + this.totalMemory + " on memory allocations."); } this.lock.lock(); try { if (size == poolableSize && !this.free.isEmpty()) { return this.free.pollFirst(); } int freeListSize = this.free.size() * this.poolableSize; if (this.availableMemory + freeListSize >= size) { freeUp(size); this.availableMemory -= size; lock.unlock(); return ByteBuffer.allocate(size); } else { int accumulated = 0; ByteBuffer buffer = null; Condition moreMemory = this.lock.newCondition(); long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); this.waiters.addLast(moreMemory); while (accumulated < size) { long startWaitNs = time.nanoseconds(); long timeNs; boolean waitingTimeElapsed; try { waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); } catch (InterruptedException e) { this.waiters.remove(moreMemory); throw e; } finally { long endWaitNs = time.nanoseconds(); timeNs = Math.max(0L, endWaitNs - startWaitNs); this.waitTime.record(timeNs, time.milliseconds()); } if (waitingTimeElapsed) { this.waiters.remove(moreMemory); throw new TimeoutException("Failed to allocate memory within the configured max blocking time " + maxTimeToBlockMs + " ms."); } remainingTimeToBlockNs -= timeNs; if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) { buffer = this.free.pollFirst(); accumulated = size; } else { freeUp(size - accumulated); int got = (int) Math.min(size - accumulated, this.availableMemory); this.availableMemory -= got; accumulated += got; } } Condition removed = this.waiters.removeFirst(); if (removed != moreMemory) { throw new IllegalStateException("Wrong condition: this shouldn't happen."); } if (this.availableMemory > 0 || !this.free.isEmpty()) { if (!this.waiters.isEmpty()) { this.waiters.peekFirst() .signal(); } } lock.unlock(); if (buffer == null) { return ByteBuffer.allocate(size); } else { return buffer; } } } finally { if (lock.isHeldByCurrentThread()) { lock.unlock(); } } } BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName); ByteBuffer allocate(int size, long maxTimeToBlockMs); void deallocate(ByteBuffer buffer, int size); void deallocate(ByteBuffer buffer); long availableMemory(); long unallocatedMemory(); int queued(); int poolableSize(); long totalMemory();  }
@Test public void testDelayedAllocation() throws Exception { BufferPool pool = new BufferPool(5 * 1024, 1024, metrics, time, metricGroup); ByteBuffer buffer = pool.allocate(1024, maxBlockTimeMs); CountDownLatch doDealloc = asyncDeallocate(pool, buffer); CountDownLatch allocation = asyncAllocate(pool, 5 * 1024); assertEquals("Allocation shouldn't have happened yet, waiting on memory.", 1L, allocation.getCount()); doDealloc.countDown(); assertTrue("Allocation should succeed soon after de-allocation", allocation.await(1, TimeUnit.SECONDS)); }
@Test public void testBlockTimeout() throws Exception { BufferPool pool = new BufferPool(10, 1, metrics, systemTime, metricGroup); ByteBuffer buffer1 = pool.allocate(1, maxBlockTimeMs); ByteBuffer buffer2 = pool.allocate(1, maxBlockTimeMs); ByteBuffer buffer3 = pool.allocate(1, maxBlockTimeMs); delayedDeallocate(pool, buffer1, maxBlockTimeMs / 2); delayedDeallocate(pool, buffer2, maxBlockTimeMs); delayedDeallocate(pool, buffer3, maxBlockTimeMs / 2 * 5); long beginTimeMs = systemTime.milliseconds(); try { pool.allocate(10, maxBlockTimeMs); fail("The buffer allocated more memory than its maximum value 10"); } catch (TimeoutException e) { } long endTimeMs = systemTime.milliseconds(); assertTrue("Allocation should finish not much later than maxBlockTimeMs", endTimeMs - beginTimeMs < maxBlockTimeMs + 1000); }