strangedb-core/src/main/java/it/cavallium/strangedb/database/blocks/DatabaseBlocksMetadata.java

package it.cavallium.strangedb.database.blocks;


import it.cavallium.strangedb.database.IBlocksMetadata;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousFileChannel;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicLong;

import static it.cavallium.strangedb.database.IDatabase.DISK_BLOCK_SIZE;

public class DatabaseBlocksMetadata implements IBlocksMetadata {
	public static final BlockInfo ERROR_BLOCK_INFO = new BlockInfo(-2, 0);
	private static final int BLOCK_META_BYTES_COUNT = Long.BYTES + Integer.BYTES;
	public static final int BLOCK_META_READS_AT_EVERY_READ = (DISK_BLOCK_SIZE - DISK_BLOCK_SIZE % BLOCK_META_BYTES_COUNT) / BLOCK_META_BYTES_COUNT;

	private final AsynchronousFileChannel metaFileChannel;
	private final DatabaseBlocksMetadataCache cache;
	private AtomicLong firstFreeBlock;

	public DatabaseBlocksMetadata(Path metaFile) throws IOException {
		metaFileChannel = AsynchronousFileChannel.open(metaFile, StandardOpenOption.READ, StandardOpenOption.WRITE);
		firstFreeBlock = new AtomicLong(metaFileChannel.size() / BLOCK_META_BYTES_COUNT);
		this.cache = new DatabaseBlocksMetadataCache(this::writeBlockToDisk);
	}

	@Override
	public BlockInfo getBlockInfo(long blockId) throws IOException {
		if (blockId == EMPTY_BLOCK_ID) {
			return EMPTY_BLOCK_INFO;
		}
		if (blockId == ERROR_BLOCK_ID) {
			throw new IOException("Errored block id");
		}
		BlockInfo blockInfo;
		if ((blockInfo = cache.get(blockId)) != ERROR_BLOCK_INFO) {
			return blockInfo;
		}
		long position = blockId * BLOCK_META_BYTES_COUNT;
		int size = BLOCK_META_READS_AT_EVERY_READ * BLOCK_META_BYTES_COUNT;
		long currentFirstFreeBlock = this.firstFreeBlock.get();
		if (blockId > currentFirstFreeBlock) {
			return EMPTY_BLOCK_INFO;
		}
		if (blockId + (size - 1) / BLOCK_META_BYTES_COUNT >= currentFirstFreeBlock) {
			size = (int) ((currentFirstFreeBlock - blockId) * BLOCK_META_BYTES_COUNT);
		}
		int blocksCount = size / BLOCK_META_BYTES_COUNT;

		ByteBuffer buffer = ByteBuffer.allocate(size);
		try {
			metaFileChannel.read(buffer, position).get();
		} catch (InterruptedException e) {
			throw new IOException(e);
		} catch (ExecutionException e) {
			throw new IOException(e.getCause());
		}
		buffer.flip();

		if (blocksCount < 1) {
			throw new IOException("Trying to read <1 blocks");
		}
		if (buffer.limit() % BLOCK_META_BYTES_COUNT != 0 || buffer.limit() < BLOCK_META_BYTES_COUNT) {
			throw new IOException("The buffer is smaller than the data requested.");
		} else if (buffer.limit() != size) {
			size = buffer.limit();
			blocksCount = size / BLOCK_META_BYTES_COUNT;
		}

		long[] allBlockIds = new long[blocksCount];
		BlockInfo[] allBlockInfo = new BlockInfo[blocksCount];

		blockInfo = EMPTY_BLOCK_INFO;
		for (int delta = 0; delta < blocksCount; delta++) {
			long blockToLoad = blockId + delta;
			long blockIndex = buffer.getLong();
			int blockSize = buffer.getInt();
			BlockInfo currentBlockInfo = new BlockInfo(blockIndex, blockSize);
			allBlockIds[delta] = blockToLoad;
			allBlockInfo[delta] = currentBlockInfo;
			if (blockToLoad == blockId) {
				blockInfo = currentBlockInfo;
			}
		}
		cache.putAll(allBlockIds, allBlockInfo);
		return blockInfo;
	}

	@Override
	public long newBlock(long index, int size) throws IOException {
		if (size == 0) {
			return EMPTY_BLOCK_ID;
		}
		long newBlockId = firstFreeBlock.getAndIncrement();
		BlockInfo blockInfo = new BlockInfo(index, size);
		cache.put(newBlockId, blockInfo);
		return newBlockId;
	}

	@Override
	public void close() throws IOException {
		cache.close();
		metaFileChannel.close();
	}

	private void writeBlockToDisk(long blockId, long index, int size) throws IOException {
		ByteBuffer data = ByteBuffer.allocate(BLOCK_META_BYTES_COUNT);
		data.putLong(index);
		data.putInt(size);
		data.flip();
		try {
			metaFileChannel.write(data, blockId * BLOCK_META_BYTES_COUNT).get();
		} catch (InterruptedException e) {
			throw new IOException(e);
		} catch (ExecutionException e) {
			throw new IOException(e.getCause());
		}
	}

	@Override
	public long getTotalBlocksCount() {
		return firstFreeBlock.get();
	}
}