All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.thinkaurelius.thrift.util.mem.Memory Maven / Gradle / Ivy

/**
 * 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 com.thinkaurelius.thrift.util.mem;

import java.io.UnsupportedEncodingException;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;

import com.sun.jna.Pointer;
import sun.misc.Unsafe;

/**
 * An off-heap region of memory that must be manually free'd when no longer needed.
 */
class Memory
{
    private static final Unsafe unsafe;

    static
    {
        try
        {
            Field field = sun.misc.Unsafe.class.getDeclaredField("theUnsafe");
            field.setAccessible(true);
            unsafe = (sun.misc.Unsafe) field.get(null);
        }
        catch (Exception e)
        {
            throw new AssertionError(e);
        }
    }

    protected long peer, size; // size could be changed by 'reallocate'

    protected Memory(long size)
    {
        this(unsafe.allocateMemory(size), size);
    }

    protected Memory(long peer, long size)
    {
        this.peer = peer;
        this.size = size;
    }

    public static Memory allocate(long bytes)
    {
        if (bytes < 0)
            throw new IllegalArgumentException();

        return new Memory(bytes);
    }

    public static Memory wrap(long peer, int limit)
    {
        return new Memory(peer, limit);
    }

    public void setByte(long offset, byte b)
    {
        checkPosition(offset);
        setByteUnsafe(offset, b);
    }

    void setByteUnsafe(long offset, byte b)
    {
        unsafe.putByte(peer + offset, b);
    }

    public void setMemory(long offset, long bytes, byte b)
    {
        // check if the last element will fit into the memory
        checkPosition(offset + bytes - 1);
        unsafe.setMemory(peer + offset, bytes, b);
    }

    void copyMemory(long offset, long dstAddr, long bytes)
    {
        if (bytes > size - offset)
            throw new IllegalArgumentException();

        unsafe.copyMemory(peer + offset, dstAddr, bytes);
    }

    public void setLong(long offset, long l)
    {
        checkPosition(offset);
        unsafe.putLong(peer + offset, l);
    }

    /**
     * Transfers count bytes from buffer to Memory
     *
     * @param memoryOffset start offset in the memory
     * @param buffer the data buffer
     * @param bufferOffset start offset of the buffer
     * @param count number of bytes to transfer
     */
    public void setBytes(long memoryOffset, byte[] buffer, int bufferOffset, int count)
    {
        if (buffer == null)
            throw new NullPointerException();
        else if (bufferOffset < 0
                 || count < 0
                 || bufferOffset + count > buffer.length)
            throw new IndexOutOfBoundsException();
        else if (count == 0)
            return;

        checkPosition(memoryOffset);
        long end = memoryOffset + count;
        checkPosition(end - 1);
        while (memoryOffset < end)
        {
            unsafe.putByte(peer + memoryOffset, buffer[bufferOffset]);
            memoryOffset++;
            bufferOffset++;
        }
    }

    /**
     * Transfers count bytes from buffer to Memory
     *
     * @param memoryOffset start offset in the memory
     * @param buffer the data buffer
     */
    public void setBytes(long memoryOffset, ByteBuffer buffer)
    {
        if (buffer == null)
            throw new NullPointerException();
        else if (buffer.remaining() == 0)
            return;

        int bufferOffset = buffer.position();

        checkPosition(memoryOffset);
        long end = memoryOffset + buffer.remaining();
        checkPosition(end - 1);
        while (memoryOffset < end)
        {
            unsafe.putByte(peer + memoryOffset, buffer.get(bufferOffset));
            memoryOffset++;
            bufferOffset++;
        }
    }

    public byte getByte(long offset)
    {
        checkPosition(offset);
        return getByteUnsafe(offset);
    }

    byte getByteUnsafe(long offset)
    {
        return unsafe.getByte(peer + offset);
    }

    public long getLong(long offset)
    {
        checkPosition(offset);
        return unsafe.getLong(peer + offset);
    }

    /**
     * Transfers count bytes from Memory starting at memoryOffset to buffer starting at bufferOffset
     *
     * @param memoryOffset start offset in the memory
     * @param buffer the data buffer
     * @param bufferOffset start offset of the buffer
     * @param count number of bytes to transfer
     */
    public void getBytes(long memoryOffset, byte[] buffer, int bufferOffset, int count)
    {
        if (buffer == null)
            throw new NullPointerException();
        else if (bufferOffset < 0 || count < 0 || count > buffer.length - bufferOffset)
            throw new IndexOutOfBoundsException();
        else if (count == 0)
            return;

        checkPosition(memoryOffset);
        long end = memoryOffset + count;
        checkPosition(end - 1);
        while (memoryOffset < end)
        {
            buffer[bufferOffset++] = unsafe.getByte(peer + memoryOffset);
            memoryOffset++;
        }
    }

    private void checkPosition(long offset)
    {
        assert peer != 0 : "Memory was freed";
        assert offset >= 0 && offset < size : "Illegal offset: " + offset + ", size: " + size;
    }

    public long getPeer()
    {
        return peer;
    }

    public Memory reallocate(int newSize)
    {
        peer = unsafe.reallocateMemory(peer, newSize);
        size = newSize;

        return this;
    }

    public void free()
    {
        assert peer != 0;
        unsafe.freeMemory(peer);
        peer = 0;
    }

    public long size()
    {
        return size;
    }

    @Override
    public boolean equals(Object o)
    {
        if (this == o)
            return true;
        if (!(o instanceof Memory))
            return false;
        Memory b = (Memory) o;
        if (peer == b.peer && size == b.size)
            return true;
        return false;
    }

    public String toString()
    {
        return new String(toBytes());
    }

    public String toString(String enc) throws UnsupportedEncodingException
    {
        return new String(toBytes(), enc);
    }

    public byte[] toBytes()
    {
        return toBytes((int) size);
    }

    public byte[] toBytes(int size)
    {
        assert size <= (int) this.size;

        byte[] buf = new byte[size];

        for (int i = 0; i < size; i++)
            buf[i] = getByte(i);

        return buf;
    }

    public ByteBuffer toByteBuffer()
    {
        ByteBuffer buffer = new Pointer(peer).getByteBuffer(0, size);

        if (buffer != null)
            return buffer;

        buffer = ByteBuffer.allocate((int) size);

        for (int i = 0; i < size; i++)
            buffer.put(getByteUnsafe(i));

        buffer.flip();

        return buffer;
    }
}





© 2015 - 2024 Weber Informatics LLC | Privacy Policy