RdbStorageEngine.java

package org.yamcs.yarch.rocksdb;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yamcs.TimeInterval;
import org.yamcs.archive.TagDb;
import org.yamcs.utils.FileUtils;
import org.yamcs.yarch.AbstractStream;
import org.yamcs.yarch.HistogramRecord;
import org.yamcs.yarch.Partition;
import org.yamcs.yarch.StorageEngine;
import org.yamcs.yarch.TableDefinition;
import org.yamcs.yarch.TableDefinition.PartitionStorage;
import org.yamcs.yarch.TableWriter;
import org.yamcs.yarch.TableWriter.InsertMode;
import org.yamcs.yarch.YarchDatabase;
import org.yamcs.yarch.YarchException;

/**
 * Storage Engine based on RocksDB.
 * 
 * Tables are mapped to multiple RocksDB databases - one for each time based partition.
 * Value based partitions are mapped to column families.
 * 
 * 
 */
public class RdbStorageEngine implements StorageEngine {
    Map<TableDefinition, RdbPartitionManager> partitionManagers = new HashMap<>();
    final YarchDatabase ydb;
    static Map<YarchDatabase, RdbStorageEngine> instances = new HashMap<>();
    static {
        RocksDB.loadLibrary();
    }
    static Logger log = LoggerFactory.getLogger(RdbStorageEngine.class.getName());
    RdbTagDb rdbTagDb = null;
    boolean ignoreVersionIncompatibility = false;

    public RdbStorageEngine(YarchDatabase ydb, boolean ignoreVersionIncompatibility) throws YarchException {
        this.ydb = ydb;
        instances.put(ydb, this);
        this.ignoreVersionIncompatibility = ignoreVersionIncompatibility;
    }

    public RdbStorageEngine(YarchDatabase ydb) throws YarchException {
        this(ydb, false);
    }

    @Override
    public void loadTable(TableDefinition tbl) throws YarchException {
        if(tbl.hasPartitioning()) {
            RdbPartitionManager pm = new RdbPartitionManager(ydb, tbl);
            pm.readPartitionsFromDisk();
            partitionManagers.put(tbl, pm);
        }
    }

    @Override
    public void dropTable(TableDefinition tbl) throws YarchException {
        RdbPartitionManager pm = partitionManagers.remove(tbl);

        for(Partition p:pm.getPartitions()) {
            RdbPartition rdbp = (RdbPartition)p;
            File f=new File(tbl.getDataDir()+"/"+rdbp.dir);
            RDBFactory rdbFactory = RDBFactory.getInstance(ydb.getName());
            rdbFactory.closeIfOpen(f.getAbsolutePath());
            try {
                if(f.exists()) {
                    log.debug("Recursively removing {}", f);
                    FileUtils.deleteRecursively(f.toPath());
                }
            } catch (IOException e) {
                throw new YarchException("Cannot remove "+f, e);
            }
        }

    }

    @Override
    public TableWriter newTableWriter(TableDefinition tblDef, InsertMode insertMode) throws YarchException {
        if(!partitionManagers.containsKey(tblDef)) {
            throw new IllegalArgumentException("Do not have a partition manager for this table");
        }
        checkFormatVersion(tblDef);

        try {
            if(tblDef.isPartitionedByValue()) {
                if(tblDef.getPartitionStorage()==PartitionStorage.COLUMN_FAMILY) {
                    return new CfTableWriter(ydb, tblDef, insertMode, partitionManagers.get(tblDef));
                } else if(tblDef.getPartitionStorage()==PartitionStorage.IN_KEY) {
                    return new InKeyTableWriter(ydb, tblDef, insertMode, partitionManagers.get(tblDef));
                } else {
                    throw new IllegalArgumentException("Unknwon partition storage: "+tblDef.getPartitionStorage());
                }
            } else {
                return new CfTableWriter(ydb, tblDef, insertMode, partitionManagers.get(tblDef));
            }
        } catch (IOException e) {
            throw new YarchException("Failed to create writer", e);
        } 
    }


    @Override
    public AbstractStream newTableReaderStream(TableDefinition tbl, boolean ascending, boolean follow) {
        if(!partitionManagers.containsKey(tbl)) {
            throw new IllegalArgumentException("Do not have a partition manager for this table");
        }
        if(tbl.isPartitionedByValue()) {
            if(tbl.getPartitionStorage()==PartitionStorage.COLUMN_FAMILY) {
                return new CfTableReaderStream(ydb, tbl, partitionManagers.get(tbl), ascending, follow);
            } else if(tbl.getPartitionStorage()==PartitionStorage.IN_KEY) {
                return new InkeyTableReaderStream(ydb, tbl, partitionManagers.get(tbl), ascending, follow);
            } else {
                throw new RuntimeException("Unknwon partition storage: "+tbl.getPartitionStorage());
            }
        } else {
            return new CfTableReaderStream(ydb, tbl, partitionManagers.get(tbl), ascending, follow);
        }
    }

    @Override
    public void createTable(TableDefinition def) {		
        RdbPartitionManager pm = new RdbPartitionManager(ydb, def);
        partitionManagers.put(def, pm);
    }

    public static synchronized RdbStorageEngine getInstance(YarchDatabase ydb) {
        return instances.get(ydb);
    }

    public RdbPartitionManager getPartitionManager(TableDefinition tdef) {      
        return partitionManagers.get(tdef);
    }




    @Override
    public synchronized TagDb getTagDb() throws YarchException {
        if(rdbTagDb==null) {
            try {
                rdbTagDb = new RdbTagDb(ydb);
            } catch (RocksDBException e) {
                throw new YarchException("Cannot create tag db",e);
            }
        }
        return rdbTagDb;
    }

    /** 
     * Called from Unit tests to cleanup before the next test
     */
    public void shutdown() {
        RDBFactory rdbFactory = RDBFactory.getInstance(ydb.getName());
        rdbFactory.shutdown();
    }

    /**
     * Called from unit tests to cleanup before the next test
     * @param ydb
     */
    public static synchronized void removeInstance(YarchDatabase ydb) {
        RdbStorageEngine rse = instances.remove(ydb);
        if(rse!=null) {
            rse.shutdown();
        }
    }

    /**
     * set to ignore version incompatibility - only used from the version upgrading functions to allow loading old tables.
     * 
     * @param b
     */
    public void setIgnoreVersionIncompatibility(boolean b) {
        this.ignoreVersionIncompatibility = b;
    }

    private void checkFormatVersion(TableDefinition tblDef) throws YarchException {
        if(ignoreVersionIncompatibility) {
            return;
        }

        if(tblDef.getFormatVersion()!=TableDefinition.CURRENT_FORMAT_VERSION) {
            throw new YarchException("Table "+ydb.getName()+"/"+tblDef.getName()+" format version is "+tblDef.getFormatVersion()
            + " instead of "+TableDefinition.CURRENT_FORMAT_VERSION+", please upgrade (use the \"yamcs archive upgrade\" command).");
        }
    }
    @Override
    public Iterator<HistogramRecord> getHistogramIterator(TableDefinition tblDef, String columnName, TimeInterval interval, long mergeTime) throws YarchException {
        checkFormatVersion(tblDef);
        try {
            return new RdbHistogramIterator(ydb, tblDef, columnName, interval, mergeTime);
        } catch (RocksDBException e) {
            throw new YarchException(e);
        }
    }
}