DirectBufferPool { public ByteBuffer allocate(int size) 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.allocateDirect(size); } else if (!blockOnExhaustion) { throw new BufferOverflowException(); } else { int accumulated = 0; ByteBuffer buffer = null; Condition moreMemory = this.lock.newCondition(); this.waiters.addLast(moreMemory); while (accumulated < size) { moreMemory.await(); 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.allocateDirect(size); else return buffer; } } finally { if (lock.isHeldByCurrentThread()) lock.unlock(); } } DirectBufferPool(long memory, int poolableSize, boolean blockOnExhaustion); ByteBuffer allocate(int size); void deallocate(ByteBuffer buffer, int size); void deallocate(ByteBuffer buffer); long availableMemory(); long unallocatedMemory(); int queued(); int poolableSize(); long totalMemory();  }
@Test(expected = IllegalArgumentException.class) public void shouldNotAllocateMoreMemoryThanTotalMemoryWeHave() throws InterruptedException { DirectBufferPool pool = new DirectBufferPool(1024, 512, true); pool.allocate(1025); }
@Test(expected = BufferOverflowException.class) public void shouldThrowExceptionOnBufferExhaustion() throws InterruptedException { DirectBufferPool pool = new DirectBufferPool(1024, 512, false); pool.allocate(512); pool.allocate(768); }
@Test public void shouldBlockOnAllocationUntilMemoryIsAvailable() throws Exception { int total = 5 * 1024; DirectBufferPool pool = new DirectBufferPool(total, 1024, true); ByteBuffer buffer = pool.allocate(1024); CountDownLatch doDeallocation = asyncDeallocate(pool, buffer); CountDownLatch allocation = asyncAllocate(pool, total); assertEquals("Allocation shouldn't have happened yet, waiting on memory.", 1L, allocation.getCount()); doDeallocation.countDown(); allocation.await(); }