src/java/org/apache/cassandra/net/BufferPoolAllocator.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.cassandra.net;

import java.nio.ByteBuffer;

import io.netty.buffer.AbstractByteBufAllocator;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.buffer.UnpooledUnsafeDirectByteBuf;
import org.apache.cassandra.io.compress.BufferType;
import org.apache.cassandra.utils.memory.BufferPool;
import org.apache.cassandra.utils.memory.BufferPools;
import org.assertj.core.util.VisibleForTesting;

import static java.lang.Integer.max;

/**
 * A trivial wrapper around BufferPool for integrating with Netty, but retaining ownership of pooling behaviour
 * that is integrated into Cassandra's other pooling.
 */
public abstract class BufferPoolAllocator extends AbstractByteBufAllocator
{
    private static final BufferPool bufferPool = BufferPools.forNetworking();

    BufferPoolAllocator()
    {
        super(true);
    }

    @Override
    public boolean isDirectBufferPooled()
    {
        return true;
    }

    /** shouldn't be invoked */
    @Override
    protected ByteBuf newHeapBuffer(int minCapacity, int maxCapacity)
    {
        return Unpooled.buffer(minCapacity, maxCapacity);
    }

    @Override
    protected ByteBuf newDirectBuffer(int minCapacity, int maxCapacity)
    {
        ByteBuf result = new Wrapped(this, getAtLeast(minCapacity), maxCapacity);
        result.clear();
        return result;
    }

    ByteBuffer get(int size)
    {
        return bufferPool.get(size, BufferType.OFF_HEAP);
    }

    ByteBuffer getAtLeast(int size)
    {
        return bufferPool.getAtLeast(size, BufferType.OFF_HEAP);
    }

    void put(ByteBuffer buffer)
    {
        bufferPool.put(buffer);
    }

    void putUnusedPortion(ByteBuffer buffer)
    {
        bufferPool.putUnusedPortion(buffer);
    }

    @VisibleForTesting
    public long usedSizeInBytes()
    {
        return bufferPool.usedSizeInBytes();
    }

    void release()
    {
    }

    /**
     * A simple extension to UnpooledUnsafeDirectByteBuf that returns buffers to BufferPool on deallocate,
     * and permits extracting the buffer from it to take ownership and use directly.
     */
    public static class Wrapped extends UnpooledUnsafeDirectByteBuf
    {
        private ByteBuffer wrapped;

        Wrapped(BufferPoolAllocator allocator, ByteBuffer wrap, int maxCapacity)
        {
            super(allocator, wrap, max(wrap.capacity(), maxCapacity));
            wrapped = wrap;
        }

        @Override
        public ByteBuf capacity(int newCapacity)
        {
            if (newCapacity == capacity())
                return this;

            ByteBuf newBuffer = super.capacity(newCapacity);
            ByteBuffer nioBuffer = newBuffer.nioBuffer(0, newBuffer.capacity());

            bufferPool.put(wrapped);
            wrapped = nioBuffer;
            return newBuffer;
        }

        @Override
        protected ByteBuffer allocateDirect(int initialCapacity)
        {
            return bufferPool.getAtLeast(initialCapacity, BufferType.OFF_HEAP);
        }

        @Override
        protected void freeDirect(ByteBuffer buffer)
        {
            // noop
            // buffer is put back into the pool by deallocate()
        }

        @Override
        public void deallocate()
        {
            super.deallocate();
            if (wrapped != null)
                bufferPool.put(wrapped);
        }

        public ByteBuffer adopt()
        {
            if (refCnt() > 1)
                throw new IllegalStateException();
            ByteBuffer adopt = wrapped;
            adopt.position(readerIndex()).limit(writerIndex());
            wrapped = null;
            return adopt;
        }
    }
}