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

package org.warp.jcwdb;

import it.unimi.dsi.fastutil.longs.Long2IntLinkedOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2IntMap;

import java.io.IOException;
import java.nio.channels.SeekableByteChannel;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Stream;

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

	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 Long2IntMap freeBytes = new Long2IntLinkedOpenHashMap();

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

	public FileAllocator(SeekableByteChannel dataFileChannel, long fileSize, Long2IntMap freeBytes) {
		this.dataFileChannel = dataFileChannel;
		this.fileSize = fileSize;
		this.freeBytes.putAll(freeBytes);
	}

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

	private long allocateIntoUnusedParts(int size) {
		if (FileIndexManager.ALWAYS_ALLOCATE_NEW) return -1;
		Stream<Map.Entry<Long,Integer>> sorted =
						freeBytes.entrySet().stream()
										.sorted(Collections.reverseOrder(Map.Entry.comparingByValue()));
		final VariableWrapper<Long> holeOffset = new VariableWrapper<>(-1L);
		final VariableWrapper<Integer> holeSize = new VariableWrapper<>(0);
		sorted.anyMatch((entry) -> {
			int currentHoleSize = entry.getValue();
			if (currentHoleSize < size) {
				return true;
			}
			holeOffset.var = entry.getKey();
			holeSize.var = currentHoleSize;
			return false;
		});
		if (holeOffset.var != -1L) {
			freeBytes.remove(holeOffset.var);
			if (holeSize.var > size) {
				freeBytes.put(holeOffset.var + size, holeSize.var - size);
			}
		}
		return holeOffset.var;
	}

	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 (FileIndexManager.ALWAYS_ALLOCATE_NEW) return;

		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 && length > 0) {
				freeBytes.put(startPosition, length);
			}
		}

		if (startPosition + length >= fileSize) {
			fileSize = startPosition;
		}

		// Remove the smallest hole in the file
		if (freeBytes.size() > MAXIMUM_UNALLOCATED_ENTRIES) {
			Stream<Map.Entry<Long,Integer>> sorted =
							freeBytes.entrySet().stream()
											.sorted(Map.Entry.comparingByValue());
			Optional<Map.Entry<Long, Integer>> first = sorted.findFirst();
			if (first.isPresent()) {
				freeBytes.remove(first.get().getKey());
			}
		}
	}


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