@Override
    public ByteBuffer tryAllocate(int sizeBytes) {
        if (sizeBytes > batchSize) {
            throw new IllegalArgumentException("Cannot allocate buffers larger than max " +
                "batch size of " + batchSize);
        }

        lock.lock();
        try {
            ByteBuffer buffer = free.poll();
            // Always allocation a new buffer if there are no free buffers
            if (buffer == null) {
                buffer = ByteBuffer.allocate(batchSize);
                numAllocatedBatches += 1;
            }

            return buffer;
        } finally {
            lock.unlock();
        }
    }
@Test
    public void testOversizeAllocation() {
        int batchSize = 1024;
        int maxRetainedBatches = 3;

        BatchMemoryPool pool = new BatchMemoryPool(maxRetainedBatches, batchSize);
        assertThrows(IllegalArgumentException.class, () -> pool.tryAllocate(batchSize + 1));
    }