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

package org.warp.jcwdb;

import java.io.IOException;
import java.lang.ref.WeakReference;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.function.Consumer;

import it.unimi.dsi.fastutil.longs.Long2ObjectLinkedOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap.Entry;
import it.unimi.dsi.fastutil.longs.LongArrayList;
import it.unimi.dsi.fastutil.objects.ObjectIterator;

public class JCWDatabase implements AutoCloseable, Cleanable {
	public final static long MAX_LOADED_REFERENCES = 1000;
	public final static long MAX_LOADED_INDICES = 10000;
	
	private final TypesManager typesManager;
	private final MixedIndexDatabase indices;
	private final Cleaner databaseCleaner;
	private final EntryReferenceTools entryReferenceTools = new EntryReferenceTools();
	private final Long2ObjectMap<WeakReference<EntryReference<?>>> references;
	private volatile boolean closed;
	private final Object closeLock = new Object();
	private final Object indicesAccessLock = new Object();
	private final Object referencesAccessLock = new Object();

	public JCWDatabase(Path dataFile, Path metadataFile) throws IOException {
		this.typesManager = new TypesManager(this);
		this.indices = new MixedIndexDatabase(typesManager, dataFile, metadataFile);
		this.references = new Long2ObjectLinkedOpenHashMap<>();
		Runtime.getRuntime().addShutdownHook(new Thread(() -> {
			try {
				JCWDatabase.this.close();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}));
		this.databaseCleaner = new Cleaner(this);
		
		this.databaseCleaner.start();
	}

	public <T> EntryReference<LightList<T>> getRoot() throws IOException {
		return getRoot(Object.class).cast();
	}
	
	public <T> EntryReference<LightList<T>> getRoot(Class<T> clazz) throws IOException {
		checkClosed();
		if (exists(0)) {
			return get(0);
		} else {
			LightList<T> newRoot = new LightList<T>(this, new LongArrayList());
			return set(0, newRoot);
		}
	}

	public <T> EntryReference<T> get(long index) throws IOException {
		checkClosed();
		synchronized (referencesAccessLock) {
			WeakReference<EntryReference<?>> refRef = this.references.getOrDefault(index, null);
			EntryReference<T> ref;
			if (refRef == null || (ref = (EntryReference<T>) refRef.get()) == null) {
				int type;
				long hash;
				synchronized (indicesAccessLock) {
					type = this.indices.getType(index);
					hash = this.indices.getHash(index);
				}
				DBTypeParser<T> typeParser = this.typesManager.get(type);
				ref = new EntryReference<>(entryReferenceTools, index, hash, typeParser);
				refRef = new WeakReference<>(ref);
				this.references.put(index, refRef);	
			}
			return ref;
		}
	}

	protected <T> EntryReference<T> add(T value) throws IOException {
		checkClosed();
		synchronized (referencesAccessLock) {
			EntryReference<T> ref;
			DBTypeParser<T> typeParser = this.typesManager.get((Class<T>) value.getClass());
			long index;
			long hash;
			synchronized (indicesAccessLock) {
				index = indices.add(typeParser.getWriter(value));
				hash = indices.getHash(index);
			}
			ref = new EntryReference<>(entryReferenceTools, index, hash, typeParser, value);
			this.references.put(index, new WeakReference<>(ref));
			return ref;
		}
	}

	protected boolean exists(long index) {
		checkClosed();
		synchronized (referencesAccessLock) {
			synchronized (indicesAccessLock) {
				return this.references.containsKey(index) || this.indices.has(index);
			}
		}
	}

	protected <T> EntryReference<T> set(long index, T value) throws IOException {
		checkClosed();
		synchronized (referencesAccessLock) {
			EntryReference<T> ref;
			if (exists(index)) {
				ref = get(index);
				ref.setValue(value);
				return ref;
			} else {
				@SuppressWarnings("unchecked")
				DBTypeParser<T> typeParser = this.typesManager.get((Class<T>) value.getClass());
				long hash;
				synchronized (indicesAccessLock) {
					IndexDetails returnedDetails = indices.set(index, typeParser.getWriter(value));
					hash = returnedDetails.getHash();
				}
				ref = new EntryReference<>(entryReferenceTools, index, hash, typeParser);
				this.references.put(index, new WeakReference<EntryReference<?>>(ref));
				return ref;
			}
		}
	}

	public boolean isOpen() {
		return !closed;
	}

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

		this.databaseCleaner.stop();

		synchronized (referencesAccessLock) {
			ObjectIterator<WeakReference<EntryReference<?>>> iterator = references.values().iterator();
			while (iterator.hasNext()) {
				WeakReference<EntryReference<?>> referenceRef = iterator.next();
				EntryReference<?> reference = referenceRef.get();
				if (reference != null) {
					reference.close();
					iterator.remove();
				}

			}
		}
		synchronized (indicesAccessLock) {
			this.indices.close();
		}
		System.out.println("Database closed.");
	}

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

	@Override
	public long clean() {
		long removedItems = cleanEmptyReferences()
				+ cleanExtraReferences()
				+ indices.clean();
		return removedItems;
	}
	

	private long cleanEmptyReferences() {
		long removed = 0;
		synchronized(referencesAccessLock) {
			ObjectIterator<Entry<WeakReference<EntryReference<?>>>> iterator = references.long2ObjectEntrySet().iterator();
			while (iterator.hasNext()) {
				Entry<WeakReference<EntryReference<?>>> entry = iterator.next();
				if (entry.getValue().get() == null) {
					iterator.remove();
					removed++;
				}
			}
		}
		return removed;
	}

	private long cleanExtraReferences() {
		long removedReferences = 0;
		synchronized(referencesAccessLock) {
			if (references.size() > MAX_LOADED_REFERENCES) {
				long count = 0;
				ObjectIterator<Entry<WeakReference<EntryReference<?>>>> iterator = references.long2ObjectEntrySet().iterator();
				while (iterator.hasNext()) {
					Entry<WeakReference<EntryReference<?>>> entry = iterator.next();
					if (count > MAX_LOADED_REFERENCES * 3l / 2l) {
						WeakReference<EntryReference<?>> weakRef = entry.getValue();
						EntryReference<?> ref = weakRef.get();
						if (ref != null) {
							try {
								ref.close();
							} catch (IOException e) {
								e.printStackTrace();
							}
						}
						iterator.remove();
						removedReferences++;
					} else {
						count++;
					}
				}
			}
		}
		return removedReferences;
	}

	public class EntryReferenceTools {
		private EntryReferenceTools() {

		}

		public <T> T read(long index, DBReader<T> reader) throws IOException {
			return indices.get(index, reader);
		}

		public <T> IndexDetails write(long index, DBDataOutput<T> writer) throws IOException {
			return indices.set(index, writer);
		}
	}

	@SuppressWarnings("unchecked")
	protected <T> long calculateHash(T o) {
		return ((DBTypeParser<T>) typesManager.get(o.getClass())).calculateHash(o);
	}

}