Java Code Examples for org.rocksdb.RocksDB#open()

/**
 * Create a new rocks DB at a particular location on disk.
 * @param pathToIndex A path to the directory where the index will be created.
 * @param columnFamilies Column families to create in the DB.
 * @param <T> A type (probably an enum) that represents a set of column families.
 * @return A DB and map of column family labels (as T) to enums.
 * @throws RocksDBException
 */
public static <T extends ColumnFamilyEnumeration<T>> RocksDBAndHandles<T> createNewRocksDB(
    File pathToIndex, T[] columnFamilies) throws RocksDBException {
  RocksDB db = null; // Not auto-closable.
  Map<T, ColumnFamilyHandle> columnFamilyHandles = new HashMap<>();

  db = RocksDB.open(ROCKS_DB_CREATE_OPTIONS, pathToIndex.getAbsolutePath());

  for (T cf : columnFamilies) {
    LOGGER.info("Creating column family %s", cf.getName());
    ColumnFamilyHandle cfh =
        db.createColumnFamily(new ColumnFamilyDescriptor(cf.getName().getBytes(UTF8)));
    columnFamilyHandles.put(cf, cfh);
  }

  return new RocksDBAndHandles<T>(db, columnFamilyHandles);
}
 

@Override
protected void before() throws Throwable {
	this.temporaryFolder = new TemporaryFolder();
	this.temporaryFolder.create();
	final File rocksFolder = temporaryFolder.newFolder();
	this.dbOptions = optionsFactory.createDBOptions(PredefinedOptions.DEFAULT.createDBOptions()).
		setCreateIfMissing(true);
	this.columnFamilyOptions = optionsFactory.createColumnOptions(PredefinedOptions.DEFAULT.createColumnOptions());
	this.writeOptions = new WriteOptions();
	this.writeOptions.disableWAL();
	this.readOptions = new ReadOptions();
	this.columnFamilyHandles = new ArrayList<>(1);
	this.rocksDB = RocksDB.open(
		dbOptions,
		rocksFolder.getAbsolutePath(),
		Collections.singletonList(new ColumnFamilyDescriptor("default".getBytes(), columnFamilyOptions)),
		columnFamilyHandles);
	this.batchWrapper = new RocksDBWriteBatchWrapper(rocksDB, writeOptions);
}
 

@Override
public void readSnapshot(String snapshotDir) {
    try {
        // copy snapshot dir to data dir
        if (db != null) {
            db.close();
        }
        String dataDir = raftDataDir + File.separator + "rocksdb_data";
        File dataFile = new File(dataDir);
        if (dataFile.exists()) {
            FileUtils.deleteDirectory(dataFile);
        }
        File snapshotFile = new File(snapshotDir);
        if (snapshotFile.exists()) {
            FileUtils.copyDirectory(snapshotFile, dataFile);
        }
        // open rocksdb data dir
        Options options = new Options();
        options.setCreateIfMissing(true);
        db = RocksDB.open(options, dataDir);
    } catch (Exception ex) {
        LOG.warn("meet exception, msg={}", ex.getMessage());
    }
}
 

@SuppressWarnings("resource")
public static KnowledgeBase getInstance(final String kbDir, final String kbMetadataPathname, final boolean readOnly) throws RocksDBException, ClassNotFoundException, IOException {
	final boolean metadataExists = new File(kbMetadataPathname).exists();
	final boolean kbDirExists = new File(kbDir).exists();
	if (metadataExists != kbDirExists) throw new IllegalArgumentException("Either both or none of the knowledge-base directory and metadata must exist");

	RocksDB.loadLibrary();
	final ColumnFamilyOptions cfOptions = new ColumnFamilyOptions().setCompressionType(CompressionType.LZ4_COMPRESSION);
	final DBOptions dbOptions = new DBOptions().setCreateIfMissing(true).setCreateMissingColumnFamilies(true);
	final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, cfOptions), new ColumnFamilyDescriptor(GID2URI, cfOptions), new ColumnFamilyDescriptor(URI2GID, cfOptions));

	final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
	final RocksDB db = readOnly ? RocksDB.openReadOnly(dbOptions, kbDir, cfDescriptors, columnFamilyHandles) : RocksDB.open(dbOptions, kbDir, cfDescriptors, columnFamilyHandles);

	final KnowledgeBase kb;
	if (metadataExists) {
		kb = (KnowledgeBase) BinIO.loadObject(kbMetadataPathname);
		kb.readOnly = readOnly;
		kb.callGraphDB = db;
		kb.defaultHandle = columnFamilyHandles.get(0);
		kb.gid2uriFamilyHandle = columnFamilyHandles.get(1);
		kb.uri2gidFamilyHandle = columnFamilyHandles.get(2);
	} else kb = new KnowledgeBase(db, columnFamilyHandles.get(0), columnFamilyHandles.get(1), columnFamilyHandles.get(2), kbMetadataPathname, readOnly);
	return kb;
}
 

/**
 * Tests that {@link RocksDBWriteBatchWrapper} flushes after the memory consumed exceeds the preconfigured value.
 */
@Test
public void testWriteBatchWrapperFlushAfterMemorySizeExceed() throws Exception {
	try (RocksDB db = RocksDB.open(folder.newFolder().getAbsolutePath());
		WriteOptions options = new WriteOptions().setDisableWAL(true);
		ColumnFamilyHandle handle = db.createColumnFamily(new ColumnFamilyDescriptor("test".getBytes()));
		RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(db, options, 200, 50)) {

		long initBatchSize = writeBatchWrapper.getDataSize();
		byte[] dummy = new byte[6];
		ThreadLocalRandom.current().nextBytes(dummy);
		// will add 1 + 1 + 1 + 6 + 1 + 6 = 16 bytes for each KV
		// format is [handleType|kvType|keyLen|key|valueLen|value]
		// more information please ref write_batch.cc in RocksDB
		writeBatchWrapper.put(handle, dummy, dummy);
		assertEquals(initBatchSize + 16, writeBatchWrapper.getDataSize());
		writeBatchWrapper.put(handle, dummy, dummy);
		assertEquals(initBatchSize + 32, writeBatchWrapper.getDataSize());
		writeBatchWrapper.put(handle, dummy, dummy);
		// will flush all, then an empty write batch
		assertEquals(initBatchSize, writeBatchWrapper.getDataSize());
	}
}
 

@Test
public void deleteBackup() throws RocksDBException {
    // Open empty database.
    try (final Options opt = new Options().setCreateIfMissing(true);
         final RocksDB db = RocksDB.open(opt,
                 dbFolder.getRoot().getAbsolutePath())) {
        // Fill database with some test values
        prepareDatabase(db);
        // Create two backups
        try (final BackupableDBOptions bopt = new BackupableDBOptions(
                backupFolder.getRoot().getAbsolutePath());
             final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
            be.createNewBackup(db, false);
            be.createNewBackup(db, true);
            final List<BackupInfo> backupInfo =
                    verifyNumberOfValidBackups(be, 2);
            // Delete the first backup
            be.deleteBackup(backupInfo.get(0).backupId());
            final List<BackupInfo> newBackupInfo =
                    verifyNumberOfValidBackups(be, 1);

            // The second backup must remain.
            assertThat(newBackupInfo.get(0).backupId()).
                    isEqualTo(backupInfo.get(1).backupId());
        }
    }
}
 

@Test
public void purgeOldBackups() throws RocksDBException {
    // Open empty database.
    try (final Options opt = new Options().setCreateIfMissing(true);
         final RocksDB db = RocksDB.open(opt,
                 dbFolder.getRoot().getAbsolutePath())) {
        // Fill database with some test values
        prepareDatabase(db);
        // Create four backups
        try (final BackupableDBOptions bopt = new BackupableDBOptions(
                backupFolder.getRoot().getAbsolutePath());
             final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
            be.createNewBackup(db, false);
            be.createNewBackup(db, true);
            be.createNewBackup(db, true);
            be.createNewBackup(db, true);
            final List<BackupInfo> backupInfo =
                    verifyNumberOfValidBackups(be, 4);
            // Delete everything except the latest backup
            be.purgeOldBackups(1);
            final List<BackupInfo> newBackupInfo =
                    verifyNumberOfValidBackups(be, 1);
            // The latest backup must remain.
            assertThat(newBackupInfo.get(0).backupId()).
                    isEqualTo(backupInfo.get(3).backupId());
        }
    }
}
 

/**
 * Constructor of RocksDao (Database Access Object).
 *
 * @param dbDir Directory where RocksDB data will be stored
 * @throws RocksDBException if there is an error loading or opening RocksDB instance
 */
public RocksDao(final String dbDir) throws RocksDBException {
    RocksDB.loadLibrary();
    final ColumnFamilyOptions cfOptions = new ColumnFamilyOptions();
    final DBOptions dbOptions = new DBOptions()
            .setCreateIfMissing(true)
            .setCreateMissingColumnFamilies(true);
    final List<ColumnFamilyDescriptor> cfDescriptors = Collections.singletonList(
            new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, cfOptions));
    final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
    this.rocksDb = RocksDB.open(dbOptions, dbDir, cfDescriptors, columnFamilyHandles);
    this.defaultHandle = columnFamilyHandles.get(0);
    initKryo();
}
 

@Override
public RocksDB open(File path, Options options) throws IOException {
	try {
		return RocksDB.open(options, path.getAbsolutePath());
	} catch (RocksDBException e) {
		throw new IOException(e);
	}
}
 

public void init(String[] cfs) throws Exception {
    logger.debug("Checking streaming dict local store for {} at {}.", cubeName, String.join(", ", cfs));
    if (!dictPath.exists() && dictPath.mkdirs()) {
        logger.warn("Create {} failed.", dictPath);
    }
    // maybe following options is naive, should improve in the future
    try (DBOptions options = new DBOptions()
            .setCreateIfMissing(true)
            .setCreateMissingColumnFamilies(true)
            .setMaxBackgroundCompactions(5)
            .setWritableFileMaxBufferSize(400 * SizeUnit.KB)) {
        String dataPath = dictPath.getAbsolutePath() + "/data";
        List<ColumnFamilyDescriptor> columnFamilyDescriptorList = new ArrayList<>();
        List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>(); // to be fill in
        for (String family : cfs) {
            ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
                    family.getBytes(StandardCharsets.UTF_8));
            columnFamilyDescriptorList.add(columnFamilyDescriptor);
        }
        logger.debug("Try to open rocksdb {}.", dataPath);
        db = RocksDB.open(options, dataPath, columnFamilyDescriptorList, columnFamilyHandleList);
        Preconditions.checkNotNull(db, "RocksDB cannot created for some reasons.");
        for (int i = 0; i < columnFamilyHandleList.size(); i++) {
            columnFamilyHandleMap.put(new ByteArray(cfs[i].getBytes(StandardCharsets.UTF_8)),
                    columnFamilyHandleList.get(i));
        }
    } catch (Exception e) {
        logger.error("Init rocks db failed.", e);
        throw e;
    }
    logger.debug("Init local dict succeed.");
}
 

@BeforeClass
static public void createRocksDb() throws IOException, RocksDBException {
  if (Files.exists(dirPath)) {
    removeRecursiveDirectory(dirPath);
  }
  Files.createDirectories(dirPath);
  Options options = new Options().setCreateIfMissing(true);
  db = RocksDB.open(options, dirPath.toString());
  db.put("testString".getBytes(), "this is string".getBytes());
  db.put(ByteBuffer.allocate(4).putInt(123).array(), ByteBuffer.allocate(4).putInt(456).array());
}
 

public static RocksDB openDB(
	String path,
	List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
	List<ColumnFamilyHandle> stateColumnFamilyHandles,
	ColumnFamilyOptions columnFamilyOptions,
	DBOptions dbOptions) throws IOException {
	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		new ArrayList<>(1 + stateColumnFamilyDescriptors.size());

	// we add the required descriptor for the default CF in FIRST position, see
	// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
	columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnFamilyOptions));
	columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);

	RocksDB dbRef;

	try {
		dbRef = RocksDB.open(
			Preconditions.checkNotNull(dbOptions),
			Preconditions.checkNotNull(path),
			columnFamilyDescriptors,
			stateColumnFamilyHandles);
	} catch (RocksDBException e) {
		IOUtils.closeQuietly(columnFamilyOptions);
		columnFamilyDescriptors.forEach((cfd) -> IOUtils.closeQuietly(cfd.getOptions()));
		throw new IOException("Error while opening RocksDB instance.", e);
	}

	// requested + default CF
	Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
		"Not all requested column family handles have been created");
	return dbRef;
}
 

public RocksdbMap(String name, StorageDataType keyType, StorageDataType valueType, RocksdbStorage storage) {
    super(name, keyType, valueType, storage);

    Options options = new Options();
    options.setCreateIfMissing(true);
    BlockBasedTableConfig config = new BlockBasedTableConfig();
    options.setTableFormatConfig(config);
    dbPath = storage.getStoragePath() + File.separator + name;
    try {
        db = RocksDB.open(options, dbPath);
    } catch (RocksDBException e) {
        throw ioe(e, "Failed to open " + dbPath);
    }
    setMaxKey(lastKey());
}
 

public void init(String[] cfs) throws Exception {
    logger.debug("Checking streaming dict local store for {} at {}.", cubeName, String.join(", ", cfs));
    if (!dictPath.exists() && dictPath.mkdirs()) {
        logger.warn("Create {} failed.", dictPath);
    }
    // maybe following options is naive, should improve in the future
    try (DBOptions options = new DBOptions()
            .setCreateIfMissing(true)
            .setCreateMissingColumnFamilies(true)
            .setMaxBackgroundCompactions(5)
            .setWritableFileMaxBufferSize(400 * SizeUnit.KB)) {
        String dataPath = dictPath.getAbsolutePath() + "/data";
        List<ColumnFamilyDescriptor> columnFamilyDescriptorList = new ArrayList<>();
        List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>(); // to be fill in
        for (String family : cfs) {
            ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
                    family.getBytes(StandardCharsets.UTF_8));
            columnFamilyDescriptorList.add(columnFamilyDescriptor);
        }
        logger.debug("Try to open rocksdb {}.", dataPath);
        db = RocksDB.open(options, dataPath, columnFamilyDescriptorList, columnFamilyHandleList);
        Preconditions.checkNotNull(db, "RocksDB cannot created for some reasons.");
        for (int i = 0; i < columnFamilyHandleList.size(); i++) {
            columnFamilyHandleMap.put(new ByteArray(cfs[i].getBytes(StandardCharsets.UTF_8)),
                    columnFamilyHandleList.get(i));
        }
    } catch (Exception e) {
        logger.error("Init rocks db failed.", e);
        throw e;
    }
    logger.debug("Init local dict succeed.");
}
 

private RocksDBMetadataTable loadMetadataTable(final CacheKey key) throws RocksDBException {
  final File dir = new File(key.directory);
  if (!dir.exists() && !dir.mkdirs()) {
    LOGGER.error("Unable to create directory for rocksdb store '" + key.directory + "'");
  }
  return new RocksDBMetadataTable(
      RocksDB.open(metadataOptions, key.directory),
      key.requiresTimestamp,
      visibilityEnabled,
      compactOnWrite);
}
 

@Test
    public void backupDb() throws RocksDBException {
//        String originPath = dbFolder.getRoot().getAbsolutePath();
//        String backupPath = backupFolder.getRoot().getAbsolutePath();
        String originPath = "/tmp/rocksdb";
        String backupPath = "/tmp/rocksdb_backup";
        System.out.println("originPath=" + originPath);
        System.out.println("backupPath=" + backupPath);
        // Open empty database.
        try (final Options opt = new Options().setCreateIfMissing(true);
             final RocksDB db = RocksDB.open(opt, originPath)) {

            // Fill database with some test values
            prepareDatabase(db);

            try (RocksIterator it = db.newIterator()) {
                for (it.seekToFirst(); it.isValid(); it.next()) {
                    System.out.println(originPath + ":" + new String(it.key()) + ":" + new String(it.value()));
                }
            }

            // Create two backups
            try (final BackupableDBOptions bopt = new BackupableDBOptions(backupPath);
                 final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
                be.createNewBackup(db, false);
                be.createNewBackup(db, true);
                verifyNumberOfValidBackups(be, 2);
            }
        }
    }
 

@Override
protected void before() throws Throwable {
  dbPath = temporaryFolder.newFolder().getPath();
  db = RocksDB.open(dbPath);
}
 

public static RestoreState restore() throws RocksDBException {
    if (restoring) {
        LOGGER.info("is restoring, return");
        return RestoreState.BEING_RESTORE;
    }

    LOGGER.info("start restore");
    restoring = true;
    RocksDB restoreDB = null;
    try (final BackupableDBOptions bopt = new BackupableDBOptions(DB_PATH_BACKUP);
         final BackupEngine be = BackupEngine.open(Env.getDefault(), bopt)) {
        // restore db from first backup

        /**
         * @param keepLogFiles If true, restore won't overwrite the existing log files
         *   in wal_dir. It will also move all log files from archive directory to
         *   wal_dir. Use this option in combination with
         *   BackupableDBOptions::backup_log_files = false for persisting in-memory
         *   databases.
         *   Default: false
         */
        boolean keepLogFiles = false;
        be.restoreDbFromLatestBackup(DB_PATH_RESTORE, DB_PATH_RESTORE, new RestoreOptions(keepLogFiles));
        // open database again.
        restoreDB = RocksDB.open(OptionsConfig.DB_OPTIONS, DB_PATH_RESTORE, CFManager.CF_DESCRIPTORS, CFManager.CF_HANDLES);

        int i = 0;
        try (RocksIterator it = restoreDB.newIterator()) {
            for (it.seekToFirst(); it.isValid(); it.next()) {
                LOGGER.info("i:{}, key:{}, value:{}", i++, new String(it.key()), new String(it.value()));
                if (i == 10) {
                    break;
                }
            }
        }

        return RestoreState.SUCCESS;
    } catch (RocksDBException e) {
        LOGGER.error("error while restore, path:{}, err:{}", DB_PATH_RESTORE, e.getMessage(), e);
        return RestoreState.FAIL;
    } finally {
        if (restoreDB != null) {
            restoreDB.close();
        }

        restoring = false;
        LOGGER.info("end restore");
    }
}
 

/**
 * Open an existing RocksDB index.
 * @param pathToIndex A path to the RocksDB index directory to use.
 * @param columnFamilies A list of column familities to open.  Must be exhaustive, non-empty, and non-null.
 * @return A DB and map of column family labels (as T) to enums.
 * @throws RocksDBException
 */
public static <T extends ColumnFamilyEnumeration<T>> RocksDBAndHandles<T> openExistingRocksDB(
    File pathToIndex, T[] columnFamilies) throws RocksDBException {
  if (columnFamilies == null || columnFamilies.length == 0) {
    throw new RuntimeException("Cannot open a RocksDB with an empty list of column families.");
  }

  List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>(columnFamilies.length + 1);
  // Must also open the "default" family or RocksDB will probably choke.
  columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_ROCKSDB_COLUMN_FAMILY.getBytes()));
  for (T family : columnFamilies) {
    columnFamilyDescriptors.add(new ColumnFamilyDescriptor(family.getName().getBytes()));
  }
  List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(columnFamilyDescriptors.size());

  DBOptions dbOptions = ROCKS_DB_OPEN_OPTIONS;
  dbOptions.setCreateIfMissing(false);
  RocksDB db = RocksDB.open(dbOptions, pathToIndex.getAbsolutePath(),
      columnFamilyDescriptors, columnFamilyHandles);
  Map<T, ColumnFamilyHandle> columnFamilyHandleMap = new HashMap<>(columnFamilies.length);
  // TODO: can we zip these together more easily w/ Java 8?

  for (int i = 0; i < columnFamilyDescriptors.size(); i++) {
    ColumnFamilyDescriptor cfd = columnFamilyDescriptors.get(i);
    ColumnFamilyHandle cfh = columnFamilyHandles.get(i);
    String familyName = new String(cfd.columnFamilyName(), UTF8);
    T descriptorFamily = columnFamilies[0].getFamilyByName(familyName); // Use any instance to get the next family.
    if (descriptorFamily == null) {
      if (!DEFAULT_ROCKSDB_COLUMN_FAMILY.equals(familyName)) {
        String msg = String.format("Found unexpected family name '%s' when trying to open RocksDB at %s",
            familyName, pathToIndex.getAbsolutePath());
        LOGGER.error(msg);
        // Crash if we don't recognize the contents of this DB.
        throw new RuntimeException(msg);
      }
      // Just skip this column family if it doesn't map to something we know but is expected.
      continue;
    }

    columnFamilyHandleMap.put(descriptorFamily, cfh);
  }

  return new RocksDBAndHandles<T>(db, columnFamilyHandleMap);
}
 

@Test
public void restoreLatestBackup() throws RocksDBException {
    try (final Options opt = new Options().setCreateIfMissing(true)) {
        // Open empty database.
        RocksDB db = null;
        try {
            db = RocksDB.open(opt,
                    dbFolder.getRoot().getAbsolutePath());
            // Fill database with some test values
            prepareDatabase(db);

            try (final BackupableDBOptions bopt = new BackupableDBOptions(
                    backupFolder.getRoot().getAbsolutePath());
                 final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
                be.createNewBackup(db, true);
                verifyNumberOfValidBackups(be, 1);
                db.put("key1".getBytes(), "valueV2".getBytes());
                db.put("key2".getBytes(), "valueV2".getBytes());
                be.createNewBackup(db, true);
                verifyNumberOfValidBackups(be, 2);
                db.put("key1".getBytes(), "valueV3".getBytes());
                db.put("key2".getBytes(), "valueV3".getBytes());
                assertThat(new String(db.get("key1".getBytes()))).endsWith("V3");
                assertThat(new String(db.get("key2".getBytes()))).endsWith("V3");

                db.close();
                db = null;

                verifyNumberOfValidBackups(be, 2);
                // restore db from latest backup
                try (final RestoreOptions ropts = new RestoreOptions(false)) {
                    be.restoreDbFromLatestBackup(dbFolder.getRoot().getAbsolutePath(),
                            dbFolder.getRoot().getAbsolutePath(), ropts);
                }

                // Open database again.
                db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath());

                // Values must have suffix V2 because of restoring latest backup.
                assertThat(new String(db.get("key1".getBytes()))).endsWith("V2");
                assertThat(new String(db.get("key2".getBytes()))).endsWith("V2");
            }
        } finally {
            if (db != null) {
                db.close();
            }
        }
    }
}