com.gitee.yanfanvip.single.SingleDatabase Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rocksdb Show documentation
RocksDB Cluster and Replication cluster
The newest version!
package com.gitee.yanfanvip.single;

import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
import org.rocksdb.FlushOptions;
import org.rocksdb.InfoLogLevel;
import org.rocksdb.Logger;
import org.rocksdb.Options;
import org.rocksdb.ReadOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
import org.rocksdb.Snapshot;
import org.rocksdb.WriteBatch;
import org.rocksdb.WriteOptions;

import com.gitee.yanfanvip.interfaces.Database;
import com.gitee.yanfanvip.model.Row;

public class SingleDatabase implements Database{
	
	static {
		RocksDB.loadLibrary();
	}
	
	static String DATABASE_PATH = "DATABASE";
	final Options OPTIONS = new Options();
	
	final DBOptions DB_OPTIONS = new DBOptions()
			.setCreateIfMissing(true)
			.setCreateMissingColumnFamilies(true)
			.setMaxOpenFiles(500)
			.setLogger(new Logger(OPTIONS) {
				@Override
				protected void log(InfoLogLevel infoLogLevel, String logMsg) {}
			});
	
	final ColumnFamilyOptions COLUMN_FAMILY_OPTIONS = new ColumnFamilyOptions();
	
	final ReadOptions READ_OPTIONS = new ReadOptions().setFillCache(false);
	final WriteOptions WRITE_OPTIONS = new WriteOptions().setDisableWAL(false).setSync(false);
	final HashMap columnFamilys = new HashMap();
	RocksDB db;
	
	static SingleDatabase SYSTEM = null;
	static SingleDatabase database = null;
	
	public static void setDataPath(String path){
		if(path == null){ throw new IllegalArgumentException("path can not null"); }
		DATABASE_PATH = path;
	}
	
	public static synchronized SingleDatabase getSystem() throws RocksDBException {
		if(SYSTEM != null) {
			return SYSTEM;
		}else {
			SYSTEM = new SingleDatabase(DATABASE_PATH + "/SYSTEM");
			Runtime.getRuntime().addShutdownHook(new Thread() {
				@Override
				public void run() { SYSTEM.close(); }
			});
			return SYSTEM;
		}
	}
	
	public static synchronized SingleDatabase get() throws RocksDBException {
		if(database != null) {
			return database;
		}else {
			database = new SingleDatabase(DATABASE_PATH + "/DATA");
			Runtime.getRuntime().addShutdownHook(new Thread() {
				@Override
				public void run() { database.close(); }
			});
			return database;
		}
	}
	
	private SingleDatabase(String path) throws RocksDBException {
		if(!new File(path).exists()){ new File(path).mkdirs(); }
		List cfs = columnFamilys(path);
		List cfds = new ArrayList();
		cfs.forEach(cf->{
			cfds.add(new ColumnFamilyDescriptor(cf, COLUMN_FAMILY_OPTIONS));
		});
		if(cfds.size() == 0) { cfds.add(new ColumnFamilyDescriptor("default".getBytes())); }
		List cfhs = new ArrayList();
		db = RocksDB.open(DB_OPTIONS, path, cfds, cfhs);
		for (int j = 0; j < cfs.size(); j++) {
			byte[] bs = cfs.get(j);
			columnFamilys.put(new String(bs), cfhs.get(j));
		}
	}
	
	private static List columnFamilys(String path) throws RocksDBException {
		try(Options options = new Options()){
			return RocksDB.listColumnFamilies(options, path);
		}
	}
	
	ColumnFamilyHandle getColumnFamily(String columnFamilyName) throws RocksDBException {
		if(columnFamilyName.equals("default")){ throw new RocksDBException("can not use default"); }
		if (columnFamilys.containsKey(columnFamilyName)) {
			return columnFamilys.get(columnFamilyName);
		} else {
			ColumnFamilyHandle handle = db.createColumnFamily(new ColumnFamilyDescriptor(columnFamilyName.getBytes(), COLUMN_FAMILY_OPTIONS));
			columnFamilys.put(columnFamilyName, handle);
			return handle;
		}
	}

	@Override
	public boolean containsKey(String keyspaces, String key) throws RocksDBException {
		byte[] bs = db.get(getColumnFamily(keyspaces), READ_OPTIONS, key.getBytes());
		if(bs != null) { return true; }
		return false;
	}

	@Override
	public byte[] get(String keyspaces, String key) throws RocksDBException {
		return db.get(getColumnFamily(keyspaces), READ_OPTIONS, key.getBytes());
	}

	@SuppressWarnings("resource")
	@Override
	public void forEach(String keyspaces, BiConsumer consumer) throws RocksDBException {
		Snapshot snapshot = db.getSnapshot();
		try {
			try(ReadOptions readOptions = new ReadOptions().setFillCache(true).setPinData(true).setTotalOrderSeek(true).setSnapshot(snapshot)){
				try(RocksIterator iterator = db.newIterator(getColumnFamily(keyspaces), readOptions)){
					for(iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
						consumer.accept(new String(iterator.key()), iterator.value());
					}
				}
			}
		}finally {
			db.releaseSnapshot(snapshot);
			snapshot.close();
		}
	}

	@SuppressWarnings("resource")
	@Override
	public void forEach(Consumer consumer) throws RocksDBException {
		Snapshot snapshot = db.getSnapshot();
		try {
			try(ReadOptions readOptions = new ReadOptions().setFillCache(true).setPinData(true).setTotalOrderSeek(true).setSnapshot(snapshot)){
				columnFamilys.forEach((keyspaces, handler)->{
					try(RocksIterator iterator = db.newIterator(handler, readOptions)){
						for(iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
							consumer.accept(new Row(keyspaces, new String(iterator.key()), iterator.value()));
						}
					}
				});
			}
		}finally {
			db.releaseSnapshot(snapshot);
			snapshot.close();
		}
	}
	
	@SuppressWarnings("resource")
	public Row getLastData(String keyspaces) throws RocksDBException {
		Snapshot snapshot = db.getSnapshot();
		try {
			try(ReadOptions readOptions = new ReadOptions().setFillCache(true).setPinData(true).setTotalOrderSeek(true).setSnapshot(snapshot)){
				try(RocksIterator iterator = db.newIterator(getColumnFamily(keyspaces), readOptions)){
					iterator.seekToLast();
					if(iterator.isValid()){
						return new Row(keyspaces, new String(iterator.key()), iterator.value());
					}else{
						return null;
					}
				}
			}
		}finally {
			db.releaseSnapshot(snapshot);
			snapshot.close();
		}
	}
	
	@Override
	public void put(String keyspaces, String key, byte[] value) throws RocksDBException {
		db.put(getColumnFamily(keyspaces), WRITE_OPTIONS, key.getBytes(), value);
	}
	
	@Override
	public boolean putIfChanged(String keyspaces, String key, byte[] value) throws RocksDBException {
		byte[] old_value = get(keyspaces, key);
		if(!Arrays.equals(old_value, value)){
			put(keyspaces, key, value);
			return true;
		}else{
			return false;
		}
	}
	
	@Override
	public void putAll(String keyspaces, Map datas) throws RocksDBException {
		ColumnFamilyHandle handle = getColumnFamily(keyspaces);
		try(WriteBatch batch = new WriteBatch()){
			for (Entry entry : datas.entrySet()) {
				batch.put(handle, entry.getKey().getBytes(), entry.getValue());
			}
			db.write(WRITE_OPTIONS, batch);
		}
	}
	
	@Override
	public Map putAllIfChanged(String keyspaces, Map datas) throws RocksDBException {
		ColumnFamilyHandle handle = getColumnFamily(keyspaces);
		Map changed = new HashMap<>();
		for (Entry entry : datas.entrySet()) {
			byte[] old_value = db.get(handle, READ_OPTIONS, entry.getKey().getBytes());
			if(!Arrays.equals(old_value, entry.getValue())){
				changed.put(entry.getKey(), entry.getValue());
			}
		}
		if(changed.size() == 0){ return changed; }
		try(WriteBatch batch = new WriteBatch()){
			for (Entry entry : changed.entrySet()) {
				batch.put(handle, entry.getKey().getBytes(), entry.getValue());
			}
			db.write(WRITE_OPTIONS, batch);
		}
		return changed;
	}

	@Override
	public void remove(String keyspaces, String key) throws RocksDBException {
		db.delete(getColumnFamily(keyspaces), WRITE_OPTIONS, key.getBytes());
	}
	
	@Override
	public boolean removeIfExist(String keyspaces, String key) throws RocksDBException {
		if(containsKey(keyspaces, key)){
			remove(keyspaces, key);
			return true;
		}else{
			return false;
		}
	}

	@Override
	public void clear(String keyspaces) throws RocksDBException {
		org.rocksdb.ColumnFamilyHandle columnFamilyHandle = getColumnFamily(keyspaces);
		columnFamilys.remove(keyspaces);
		db.dropColumnFamily(columnFamilyHandle);
	}
	
	public void clearAll() throws RocksDBException{
		for (Entry entry : columnFamilys.entrySet()) {
			if(!entry.getKey().equals("default")){
				db.dropColumnFamily(entry.getValue());
			}
		}
		columnFamilys.clear();
	}
	
	@SuppressWarnings("resource")
	private void flush() {
		try {
			db.compactRange();
			db.flush(new FlushOptions().setWaitForFlush(true));
		} catch (RocksDBException e) {
			e.printStackTrace();
		}
	}
	
	@Override
	public void close() {
		columnFamilys.forEach((keyspaces, cf)->{ cf.close(); });
		columnFamilys.clear();
		flush();
		db.close();
		READ_OPTIONS.close();
		WRITE_OPTIONS.close();
		COLUMN_FAMILY_OPTIONS.close();
		DB_OPTIONS.close();
		OPTIONS.close();
	}
}