strangedb/src/main/java/org/warp/jcwdb/FileAllocator.java

package org.warp.jcwdb;

import it.unimi.dsi.fastutil.longs.Long2IntMap;
import it.unimi.dsi.fastutil.longs.Long2IntRBTreeMap;
import it.unimi.dsi.fastutil.objects.ObjectBidirectionalIterator;

import java.io.IOException;
import java.nio.channels.SeekableByteChannel;

public class FileAllocator implements AutoCloseable {
	private static final int MAXIMUM_UNALLOCATED_ENTRIES = 500000;

	private final SeekableByteChannel dataFileChannel;
	private volatile long fileSize;
	private volatile boolean closed;
	private final Object closeLock = new Object();
	private final Object allocateLock = new Object();
	/**
	 * index -> free space size
	 */
	private final Long2IntRBTreeMap freeBytes = new Long2IntRBTreeMap((a, b) -> (int) (a - b));

	public FileAllocator(SeekableByteChannel dataFileChannel) throws IOException {
		this.dataFileChannel = dataFileChannel;
		this.fileSize = this.dataFileChannel.size();
	}

	/**
	 * TODO: not implemented
	 *
	 * @param size
	 * @return offset
	 */
	public long allocate(int size) {
		checkClosed();
		synchronized (allocateLock) {
			long offset = allocateIntoUnusedParts(size);
			if (offset == -1) {
				return allocateToEnd(size);
			} else {
				return offset;
			}
		}
	}

	private long allocateIntoUnusedParts(int size) {
		ObjectBidirectionalIterator<Long2IntMap.Entry> it = freeBytes.long2IntEntrySet().iterator();
		long holeOffset = -1;
		int holeSize = 0;
		while (it.hasNext()) {
			Long2IntMap.Entry entry = it.next();
			int currentHoleSize = entry.getIntValue();
			if (currentHoleSize < size) {
				freeBytes.remove(holeOffset);
				if (holeSize > size) {
					freeBytes.put(holeOffset + size, holeSize - size);
				}
				break;
			}
			holeOffset = entry.getLongKey();
			holeSize = currentHoleSize;
		}
		return holeOffset;
	}

	private long allocateToEnd(int size) {
		long allocatedOffset = fileSize;
		fileSize += size;
		return allocatedOffset;
	}


	public void close() throws IOException {
		if (closed) {
			return;
		}
		synchronized (closeLock) {
			if (closed) {
				return;
			}
			closed = true;
		}
	}

	/**
	 * Frees the unused bytes
	 *
	 * @param startPosition
	 * @param length
	 */
	public void markFree(long startPosition, int length) {
		checkClosed();

		if (freeBytes.containsKey(startPosition + length)) {
			int secondLength = freeBytes.remove(startPosition + length);
			freeBytes.put(startPosition, length + secondLength);
		} else {
			boolean addedToList = false;
			for (Long2IntMap.Entry entry : freeBytes.long2IntEntrySet()) {
				if (entry.getLongKey() + entry.getIntValue() == startPosition) {
					freeBytes.put(entry.getLongKey(), entry.getIntValue() + length);
					addedToList = true;
					break;
				}
			}
			if (!addedToList) {
				freeBytes.put(startPosition, length);
			}
		}

		if (freeBytes.size() > MAXIMUM_UNALLOCATED_ENTRIES) {
			freeBytes.remove(freeBytes.lastLongKey());
		}
	}


	private void checkClosed() {
		if (closed) {
			throw new RuntimeException("Index Manager is closed.");
		}
	}
}