com.gitee.yanfanvip.cluster.ClusterDatabase 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.cluster;

import org.jgroups.Address;
import org.jgroups.ChannelListener;
import org.jgroups.Event;
import org.jgroups.Global;
import org.jgroups.Header;
import org.jgroups.JChannel;
import org.jgroups.Message;
import org.jgroups.View;
import org.jgroups.annotations.MBean;
import org.jgroups.conf.ClassConfigurator;
import org.jgroups.stack.Protocol;
import org.jgroups.util.MessageBatch;
import org.rocksdb.RocksDBException;

import com.gitee.yanfanvip.interfaces.Database;
import com.gitee.yanfanvip.model.OptionsEnum;
import com.gitee.yanfanvip.model.Row;
import com.gitee.yanfanvip.single.SingleDatabase;

import java.io.BufferedOutputStream;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Supplier;

@MBean(description="集群化数据库")
public class ClusterDatabase extends Protocol implements Database, ClusterReceiver{
	protected static final short ClusterDatabase_ID = 11500;
	static{
		ClassConfigurator.add((short) 25001, ClusterDatabaseHeader.class);
		ClassConfigurator.addProtocol(ClusterDatabase_ID, ClusterDatabase.class);
	}
    public interface Notification{
        void put(String keyspaces, String key, byte[] value);
        void putAll(String keyspaces, Map datas);
        void remove(String keyspaces, String key);
        void clear(String keyspaces);
    }
	SingleDatabase database;
	String path = null;
	protected Address localAddress;
	protected List members = new ArrayList<>();
	protected Notification notification;
	protected boolean enable = false;
	
	public ClusterDatabase(String cluster, int port) {
		path = cluster + "-" + port;
		SingleDatabase.setDataPath("DATABASE/" + path);
	}
	
	public void start() throws Exception {
		try {
			database = SingleDatabase.get();
			JChannel channel = getProtocolStack().getChannel();
			channel.setReceiver(this);
			channel.addChannelListener(new ChannelListener() {
				@Override
				public void channelDisconnected(JChannel channel) {
					
				}
				
				@Override
				public void channelConnected(JChannel channel) {
					try {
						channel.getState(null, 60 * 60 * 1000, false);
					} catch (Exception e) {
						e.printStackTrace();
						System.exit(-1);
					}
				}
				
				@Override
				public void channelClosed(JChannel channel) {
					
				}
			});
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}
	
	public void setListener(Notification notification){
		this.notification = notification;
	}
	
	@Override
    public void getState(OutputStream out) throws Exception {
        try(ObjectOutputStream output = new ObjectOutputStream(new BufferedOutputStream(out, 1024))) {
        	database.forEach(row->{
        		try { output.writeObject(row); } catch (IOException e) { throw new RuntimeException(e); }
        	});
        	output.writeObject(null);
        }
    }

    @Override
	public void setState(InputStream in) throws Exception {
    	database.clearAll();
        try(ObjectInputStream input = new ObjectInputStream(in)) {
        	Row row = null;
        	do {
        		row = (Row)input.readObject();
        		if(row != null){
        			_put(row.getKeyspaces(), row.getKey(), row.getValue());
        		}
			} while (row != null);
        }
        enable = true;
    }
	
	@Override
	public Object up(Message msg) {
		ClusterDatabaseHeader header = msg.getHeader(id);
		if(header != null){
			if(msg.getSrc().equals(this.localAddress)){
				return super.up(msg);
			}
			Row row = msg.getObject();
			switch (header.options) {
			case PUT:{
				_put(row.getKeyspaces(), row.getKey(), row.getValue());
				break;
			}
			case PUTALL:{
				_putAll(row.getKeyspaces(), row.getValues());
				break;
			}
			case REMOVE:{
				_remove(row.getKeyspaces(), row.getKey());
				break;
			}
			case CLEAR:{
				_clear(row.getKeyspaces());
				break;
			}
			default:
				throw new RuntimeException("bad request");
			}
		}
		return super.up(msg);
	}
	
	@Override
	public void up(MessageBatch batch) {
		for (Message msg : batch) {
			ClusterDatabaseHeader header = msg.getHeader(id);
			if(header != null){
				batch.remove(msg);
				up(msg);
			}
		}
		if(!batch.isEmpty()){
			up_prot.up(batch);
		}
	}
	
	@Override
	public Object up(Event evt) {
		switch (evt.type()) {
			case Event.VIEW_CHANGE:{
				View view = evt.getArg();
				handleView(view);
				break;
			}
			case Event.SET_LOCAL_ADDRESS:
                localAddress=evt.getArg();
                break;
			case 11501:
				break;
		}
		return super.up(evt);
	}
	
	@Override
	public Object down(Event evt) {
		switch (evt.type()) {
			case Event.VIEW_CHANGE:{
				View view = evt.getArg();
				handleView(view);
				break;
			}
			case Event.SET_LOCAL_ADDRESS:
                localAddress=evt.getArg();
                break;
			case 11501:
				break;
		}
		return super.down(evt);
	}
	
	@Override
	public Object down(Message msg) {
		System.out.println("down-message");
		System.out.println(msg);
		return super.down(msg);
	}
	
	protected void handleView(View view) {
		List mbrs = view.getMembers();
		if(mbrs.size() == 1){ this.enable = true; }
		List added = new ArrayList<>();
		mbrs.forEach(m->{
			if(!members.contains(m)) {
				added.add(m);
			}
		});
		members = mbrs;
		handleAddView(added);
	}
	/**
	 * 新增的节点
	 * @param added 新增的节点地址
	 */
	protected void handleAddView(List added) { }

	@Override
	public boolean containsKey(String keyspaces, String key) throws RocksDBException {
		return database.containsKey(keyspaces, key);
	}

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

	@Override
	public void forEach(String keyspaces, BiConsumer consumer) throws RocksDBException {
		database.forEach(keyspaces, consumer);
	}

	@Override
	public void forEach(Consumer consumer) throws RocksDBException {
		database.forEach(consumer);
	}

	@Override
	public void put(String keyspaces, String key, byte[] value) throws RocksDBException {
		boolean changed = clusterPut(keyspaces, key, value);
		if(changed && this.notification != null) {
			this.notification.put(keyspaces, key, value);
		}
	}
	
	@Override
	public boolean putIfChanged(String keyspaces, String key, byte[] value) throws RocksDBException {
		boolean changed = clusterPut(keyspaces, key, value);
		if(changed && this.notification != null) {
			this.notification.put(keyspaces, key, value);
		}
		return changed;
	}

	@Override
	public void putAll(String keyspaces, Map datas) throws RocksDBException {
		Map result = clusterPutAll(keyspaces, datas);
		if(result.size() > 0 && this.notification != null) {
			this.notification.putAll(keyspaces, result);
		}
	}
	
	@Override
	public Map putAllIfChanged(String keyspaces, Map datas) throws RocksDBException {
		Map result = clusterPutAll(keyspaces, datas);
		if(result.size() > 0 && this.notification != null) {
			this.notification.putAll(keyspaces, result);
		}
		return result;
	}
	
	@Override
	public Row getLastData(String keyspaces) throws RocksDBException {
		return this.database.getLastData(keyspaces);
	}

	@Override
	public void remove(String keyspaces, String key) throws RocksDBException {
		boolean flag = clusterRemove(keyspaces, key);
		if (flag && this.notification != null) {
			this.notification.remove(keyspaces, key);
		}
	}
	
	@Override
	public boolean removeIfExist(String keyspaces, String key) throws RocksDBException {
		boolean flag = clusterRemove(keyspaces, key);
		if (flag && this.notification != null) {
			this.notification.remove(keyspaces, key);
		}
		return flag;
	}

	@Override
	public void clear(String keyspaces) throws RocksDBException {
		clusterClear(keyspaces);
		if(this.notification != null) {
			this.notification.clear(keyspaces);
		}
	}
	
	public boolean clusterPut(String keyspaces, String key, byte[] value) throws RocksDBException {
		boolean changed = _putIfChanged(keyspaces, key, value);
		sendMessage(OptionsEnum.PUT, new Row(keyspaces, key, value));
		return changed;
	}
	
	public Map clusterPutAll(String keyspaces, Map datas) throws RocksDBException {
		Map result = _putAllIfChanged(keyspaces, datas);
		sendMessage(OptionsEnum.PUTALL, new Row(keyspaces, datas));
		return result;
	}

	public boolean clusterRemove(String keyspaces, String key) throws RocksDBException {
		boolean changed = _removeIfExist(keyspaces, key);
		sendMessage(OptionsEnum.REMOVE, new Row(keyspaces, key));
		return changed;
	}

	public void clusterClear(String keyspaces) throws RocksDBException {
		database.clear(keyspaces);
		sendMessage(OptionsEnum.CLEAR, new Row(keyspaces));
	}
	
	public void _put(String keyspaces, String key, byte[] value) {
		try {
			database.put(keyspaces, key, value);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}
	
	public boolean _putIfChanged(String keyspaces, String key, byte[] value) {
		try {
			return database.putIfChanged(keyspaces, key, value);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _putAll(String keyspaces, Map datas) {
		try {
			database.putAll(keyspaces, datas);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}
	
	public Map _putAllIfChanged(String keyspaces, Map datas) {
		try {
			return database.putAllIfChanged(keyspaces, datas);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _remove(String keyspaces, String key) {
		try {
			database.remove(keyspaces, key);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}
	
	public boolean _removeIfExist(String keyspaces, String key) {
		try {
			return database.removeIfExist(keyspaces, key);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _clear(String keyspaces) {
		try {
			database.clear(keyspaces);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}
	
	void sendMessage(OptionsEnum option, Row row){
		Message msg = new Message(null, row);
		msg.putHeader(id, new ClusterDatabaseHeader(option));
		down_prot.down(msg);
	}

	@Override
	public void close() throws Exception {
		database.close();
	}
	
	public SingleDatabase getSingleDatabase() {
		return database;
	}
	
	public static class ClusterDatabaseHeader extends Header{
	    public Supplier create() { return ClusterDatabaseHeader::new; }
	    
	    private OptionsEnum options;
	    
	    public ClusterDatabaseHeader() {
	    	super();
	    }
	    
	    public ClusterDatabaseHeader(OptionsEnum options) {
	    	this.options = options;
	 	}
	    
	    public short getMagicId() {return (short)25001;}
	    
	    public int serializedSize() { return Global.INT_SIZE; }
	    
	    public String toString() {
	        return options.name();
	    }

	    public void writeTo(DataOutput out) throws Exception {
	    	out.writeInt(options.getType());
	    }

	    public void readFrom(DataInput in) throws Exception {
	    	this.options = OptionsEnum.get(in.readInt());
	    }
	}
}