private void initPoolWithAllBuffers(ByteBufferPool pool, int maxBuffersInPool) {
ByteBuffer[] buffers = new ByteBuffer[maxBuffersInPool];
// Just call getBuffer() on pool 'maxBuffersInPool' so as to init all buffers and then put back
// all. Makes pool with max #buffers.
for (int i = 0; i < maxBuffersInPool; i++) {
buffers[i] = pool.getBuffer();
}
for (ByteBuffer buf : buffers) {
pool.putbackBuffer(buf);
}
}
@Test
public void testAllocateByteBuffToReadInto() throws Exception {
int maxBuffersInPool = 10;
ByteBufferPool pool = new ByteBufferPool(6 * 1024, maxBuffersInPool);
initPoolWithAllBuffers(pool, maxBuffersInPool);
ByteBuff buff = null;
Pair<ByteBuff, CallCleanup> pair;
// When the request size is less than 1/6th of the pool buffer size. We should use on demand
// created on heap Buffer
pair =
RpcServer.allocateByteBuffToReadInto(pool, RpcServer.getMinSizeForReservoirUse(pool), 200);
buff = pair.getFirst();
assertTrue(buff.hasArray());
assertEquals(maxBuffersInPool, pool.getQueueSize());
assertNull(pair.getSecond());
// When the request size is > 1/6th of the pool buffer size.
pair =
RpcServer.allocateByteBuffToReadInto(pool, RpcServer.getMinSizeForReservoirUse(pool), 1024);
buff = pair.getFirst();
assertFalse(buff.hasArray());
assertEquals(maxBuffersInPool - 1, pool.getQueueSize());
assertNotNull(pair.getSecond());
pair.getSecond().run(); // CallCleanup#run should put back the BB to pool.
assertEquals(maxBuffersInPool, pool.getQueueSize());
// Request size> pool buffer size
pair =
RpcServer.allocateByteBuffToReadInto(
pool, RpcServer.getMinSizeForReservoirUse(pool), 7 * 1024);
buff = pair.getFirst();
assertFalse(buff.hasArray());
assertTrue(buff instanceof MultiByteBuff);
ByteBuffer[] bbs = ((MultiByteBuff) buff).getEnclosingByteBuffers();
assertEquals(2, bbs.length);
assertTrue(bbs[0].isDirect());
assertTrue(bbs[1].isDirect());
assertEquals(6 * 1024, bbs[0].limit());
assertEquals(1024, bbs[1].limit());
assertEquals(maxBuffersInPool - 2, pool.getQueueSize());
assertNotNull(pair.getSecond());
pair.getSecond().run(); // CallCleanup#run should put back the BB to pool.
assertEquals(maxBuffersInPool, pool.getQueueSize());
pair =
RpcServer.allocateByteBuffToReadInto(
pool, RpcServer.getMinSizeForReservoirUse(pool), 6 * 1024 + 200);
buff = pair.getFirst();
assertFalse(buff.hasArray());
assertTrue(buff instanceof MultiByteBuff);
bbs = ((MultiByteBuff) buff).getEnclosingByteBuffers();
assertEquals(2, bbs.length);
assertTrue(bbs[0].isDirect());
assertFalse(bbs[1].isDirect());
assertEquals(6 * 1024, bbs[0].limit());
assertEquals(200, bbs[1].limit());
assertEquals(maxBuffersInPool - 1, pool.getQueueSize());
assertNotNull(pair.getSecond());
pair.getSecond().run(); // CallCleanup#run should put back the BB to pool.
assertEquals(maxBuffersInPool, pool.getQueueSize());
ByteBuffer[] buffers = new ByteBuffer[maxBuffersInPool - 1];
for (int i = 0; i < maxBuffersInPool - 1; i++) {
buffers[i] = pool.getBuffer();
}
pair =
RpcServer.allocateByteBuffToReadInto(
pool, RpcServer.getMinSizeForReservoirUse(pool), 20 * 1024);
buff = pair.getFirst();
assertFalse(buff.hasArray());
assertTrue(buff instanceof MultiByteBuff);
bbs = ((MultiByteBuff) buff).getEnclosingByteBuffers();
assertEquals(2, bbs.length);
assertTrue(bbs[0].isDirect());
assertFalse(bbs[1].isDirect());
assertEquals(6 * 1024, bbs[0].limit());
assertEquals(14 * 1024, bbs[1].limit());
assertEquals(0, pool.getQueueSize());
assertNotNull(pair.getSecond());
pair.getSecond().run(); // CallCleanup#run should put back the BB to pool.
assertEquals(1, pool.getQueueSize());
pool.getBuffer();
pair =
RpcServer.allocateByteBuffToReadInto(
pool, RpcServer.getMinSizeForReservoirUse(pool), 7 * 1024);
buff = pair.getFirst();
assertTrue(buff.hasArray());
assertTrue(buff instanceof SingleByteBuff);
assertEquals(7 * 1024, ((SingleByteBuff) buff).getEnclosingByteBuffer().limit());
assertNull(pair.getSecond());
}