Java RocksDB类代码示例

import org.rocksdb.RocksDB; //导入依赖的package包/类
public RocksDBStore(String name, ColumnFamilyDescriptor family, ColumnFamilyHandle handle, RocksDB db, int stripes) {
  super();
  this.family = family;
  this.name = name;
  this.db = db;
  this.parallel = stripes;
  this.handle = handle;
  this.sharedLocks = new AutoCloseableLock[stripes];
  this.exclusiveLocks = new AutoCloseableLock[stripes];

  for (int i = 0; i < stripes; i++) {
    ReadWriteLock core = new ReentrantReadWriteLock();
    sharedLocks[i] = new AutoCloseableLock(core.readLock());
    exclusiveLocks[i] = new AutoCloseableLock(core.writeLock());
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
public FindByRangeIterator(RocksDB db, ColumnFamilyHandle handle, FindByRange<byte[]> range) {
  this.iter = db.newIterator(handle);
  this.end = range == null ? null : range.getEnd();
  this.endInclusive = range == null ? false : range.isEndInclusive();

  // position at beginning of cursor.
  if (range != null && range.getStart() != null) {
    iter.seek(range.getStart());
    if (!range.isStartInclusive() && Arrays.equals(iter.key(), range.getStart())) {
      iter.next();
    }
  } else {
    iter.seekToFirst();
  }

  populateNext();

}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
public static void main(String[] args) {
    RocksDB.loadLibrary();
    Options rockopts = new Options().setCreateIfMissing(false);
    RocksDB db = null;

    try {
        db = RocksDB.open(rockopts, args[0]);
        RocksIterator iter = db.newIterator();
        iter.seekToFirst();
        while (iter.isValid()) {
            String key = new String(iter.key(), StandardCharsets.UTF_8);
            String val = new String(iter.value(), StandardCharsets.UTF_8);
            System.out.println(key + " -> " + val);
            iter.next();
        }

        db.close();
    } catch (RocksDBException rdbe) {
        rdbe.printStackTrace(System.err);
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
public static void main(String[] args) {
    RocksDB.loadLibrary();
    Options rockopts = new Options().setCreateIfMissing(false);
    RocksDB db = null;

    try {
        db = RocksDB.open(rockopts, args[0]);
        String query = args[1];
        byte[] val = db.get(query.getBytes(StandardCharsets.UTF_8));
        if (val != null) {
            System.out.println(new String(val, StandardCharsets.UTF_8));
        }
        db.close();
    } catch (RocksDBException rdbe) {
        rdbe.printStackTrace(System.err);
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
public StreamStore(String bootstrapServer, String kafkaTopic, int nodeID, int[] partitionToNode) {
    RocksDB.loadLibrary();

    if (bootstrapServer == null) bootstrapServer = DEFAULT_BOOTSTRAP_SERVER;
    Properties consumerConfig = PartitionHandler.consumerConfig();
    Properties producerConfig = PartitionHandler.producerConfig();
    consumerConfig.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
    producerConfig.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);

    producer = new KafkaProducer<>(producerConfig);

    if (kafkaTopic == null) kafkaTopic = DEFAULT_KAFKA_TOPIC;
    if (partitionToNode.length != NUM_PARTITIONS) {
        throw new IllegalArgumentException("Expected " + NUM_PARTITIONS + " partitions, got " +
                    partitionToNode.length);
    }

    this.handlers = new PartitionHandler[NUM_PARTITIONS];
    for (int i = 0; i < NUM_PARTITIONS; i++) {
        if (partitionToNode[i] == nodeID) {
            TopicPartition topicPartition = new TopicPartition(kafkaTopic, i);
            this.handlers[i] = new PartitionHandler(topicPartition, consumerConfig, producer);
        } else {
            throw new UnsupportedOperationException("TODO multi-node support");
        }
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
@Override
public void close() {
    if (this.closed.compareAndSet(false, true)) {
        RocksDB db = this.database.get();
        if (db != null) {
            db.close();
            this.database.set(null);
        }

        if (this.writeOptions != null) {
            this.writeOptions.close();
        }

        if (this.databaseOptions != null) {
            this.databaseOptions.close();
        }

        clear(false);
        log.info("{}: Closed.", this.logId);

        Consumer<String> callback = this.closeCallback;
        if (callback != null) {
            Callbacks.invokeSafely(callback, this.id, null);
        }
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
@Override
public SuffixDataStore openSuffixDataStore(int id) {
  RocksDB.loadLibrary();
  try {
    LOGGER.info("Opening TnT suffix model: {}", id);
    RocksDB rocksDB = RocksDB.openReadOnly(dbPath.resolve(getSuffixesName(id)).toString());
    RocksDbSuffixDataStore rocksDbSuffixDataStore = new RocksDbSuffixDataStore(rocksDB);
    if (inMemory) {
      LOGGER.info("Loading TnT suffix model into memory: {}", id);
      InMemorySuffixDataStore inMemorySuffixDataStore = rocksDbSuffixDataStore.inMemory();
      LOGGER.info("Done loading TnT suffix model into memory: {}", id);
      rocksDB.close();
      return inMemorySuffixDataStore;
    }
    rocksDBS.add(rocksDB);
    return rocksDbSuffixDataStore;
  } catch (RocksDBException e) {
    throw new RuntimeException(e);
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
@Override
public KnownWordsDataStore openKnownWordDataStore(int id) {
  RocksDB.loadLibrary();
  try {
    LOGGER.info("Opening TnT model known word model: {}", id);
    RocksDB rocksDB = RocksDB.openReadOnly(dbPath.resolve(getWordsName(id)).toString());
    RocksDB candidatesDB = RocksDB.openReadOnly(dbPath.resolve(getCandidatesName(id)).toString());

    RocksDbKnownWordsDataStore rocksDbKnownWordsDataStore = new RocksDbKnownWordsDataStore(
        rocksDB, candidatesDB);
    if (inMemory) {
      LOGGER.info("Loading TnT known word model into memory: {}", id);
      InMemoryKnownWordDataStore inMemoryKnownWordDataStore = rocksDbKnownWordsDataStore
          .inMemory();
      LOGGER.info("Done loading TnT known word model into memory: {}", id);
      rocksDB.close();
      candidatesDB.close();
      return inMemoryKnownWordDataStore;
    }
    rocksDBS.add(rocksDB);
    rocksDBS.add(candidatesDB);
    return rocksDbKnownWordsDataStore;
  } catch (RocksDBException e) {
    throw new RuntimeException(e);
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
private void verifyRocksObjectsReleased() {
	//Ensure every RocksObject was closed exactly once
	for (RocksObject rocksCloseable : allCreatedCloseables) {
		verify(rocksCloseable, times(1)).close();
	}

	assertNotNull(null, keyedStateBackend.db);
	RocksDB spyDB = keyedStateBackend.db;

	if (!enableIncrementalCheckpointing) {
		verify(spyDB, times(1)).getSnapshot();
		verify(spyDB, times(1)).releaseSnapshot(any(Snapshot.class));
	}

	keyedStateBackend.dispose();
	verify(spyDB, times(1)).close();
	assertEquals(null, keyedStateBackend.db);
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
/**
 * Construct the <code>RocksDbKeyValueReader</code> with store's name,
 * database's path and Samza's config
 *
 * @param storeName name of the RocksDb defined in the config file
 * @param dbPath path to the db directory
 * @param config Samza's config
 */
public RocksDbKeyValueReader(String storeName, String dbPath, Config config) {
  // get the key serde and value serde from the config
  JavaStorageConfig storageConfig = new JavaStorageConfig(config);
  JavaSerializerConfig serializerConfig = new JavaSerializerConfig(config);

  keySerde = getSerdeFromName(storageConfig.getStorageKeySerde(storeName), serializerConfig);
  valueSerde = getSerdeFromName(storageConfig.getStorageMsgSerde(storeName), serializerConfig);

  // get db options
  ArrayList<TaskName> taskNameList = new ArrayList<TaskName>();
  taskNameList.add(new TaskName("read-rocks-db"));
  SamzaContainerContext samzaContainerContext =
      new SamzaContainerContext("0",  config, taskNameList, new MetricsRegistryMap());
  Options options = RocksDbOptionsHelper.options(config, samzaContainerContext);

  // open the db
  RocksDB.loadLibrary();
  try {
    db = RocksDB.openReadOnly(options, dbPath);
  } catch (RocksDBException e) {
    throw new SamzaException("can not open the rocksDb in " + dbPath, e);
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
public AccessControl(RocksDB db, ColumnFamilyHandle ruleCf, ColumnFamilyHandle policyCf) throws RocksDBException {
    this.db = db;
    this.ruleCf = ruleCf;
    this.policyCf = policyCf;

    rules = loadRules(db, ruleCf);
    policies = loadPolicies(db, policyCf);

    rulesBySurt = new RulesBySsurt(rules.values());

    nextRuleId = new AtomicLong(calculateNextId(db, ruleCf));
    nextPolicyId = new AtomicLong(calculateNextId(db, policyCf));

    if (policies.isEmpty()) {
        // create some default policies
        put(new AccessPolicy("Public", "public", "staff"));
        put(new AccessPolicy("Staff Only", "staff"));
        put(new AccessPolicy("No Access"));
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
private void prepareColumnFamilyDescriptors() throws RocksDBException {
    if (columnFamilies == null || columnFamilies.size() == 0) {
        if (columnFamilies == null) {
            columnFamilies = new HashSet<>();
        }
        String[] cfList = RocksDbUtils.getColumnFamilyList(directory.getAbsolutePath());
        for (String cf : cfList) {
            columnFamilies.add(RocksDbUtils.buildColumnFamilyDescriptor(cf));
        }
    }

    boolean hasDefaultCf = false;
    for (ColumnFamilyDescriptor cfd : columnFamilies) {
        if (Arrays.equals(cfd.columnFamilyName(), RocksDB.DEFAULT_COLUMN_FAMILY)) {
            hasDefaultCf = true;
            break;
        }
    }
    if (!hasDefaultCf) {
        columnFamilies.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY));
    }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
private void executeWrite(Query q)
{
  MyConnection conn = null;
  try
  {
    byte[] key = q.getKey().getBytes();
    conn = _connPool.getConnection(q.getDbName());
    _log.debug(Thread.currentThread().getName() + ": Got conn");
    ((RocksDB) conn.getConn()).put(key, q.getValue());
    _log.debug(Thread.currentThread().getName() + ": Performed put");

    q.setResult(QueryResult.OK);
  }
  catch(Exception e)
  {
    _log.error(Thread.currentThread().getName() + ": Write query failed", e);
    q.setResult(QueryResult.FAIL);
  }
  finally
  {
    if(conn != null)
      tryRelease(conn);
    _log.debug(Thread.currentThread().getName() + ": Released conn");
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
private void executeDelete(Query q)
{
  MyConnection conn = null;
  try
  {
    byte[] key = q.getKey().getBytes();
    conn = _connPool.getConnection(q.getDbName());
    _log.debug(Thread.currentThread().getName() + ": Got conn");
    ((RocksDB) conn.getConn()).remove(key);
    _log.debug(Thread.currentThread().getName() + ": Performed remove");

    q.setResult(QueryResult.OK);
  }
  catch(Exception e)
  {
    _log.error(Thread.currentThread().getName() + ": Delete query failed", e);
    q.setResult(QueryResult.FAIL);
  }
  finally
  {
    if(conn != null)
      tryRelease(conn);
    _log.debug(Thread.currentThread().getName() + ": Released conn");
  }
}
 

import org.rocksdb.RocksDB; //导入依赖的package包/类
protected TupleIndexR(RocksDB db, int N, TupleMap indexMapping, String name) {
    super(N, indexMapping, name) ;
    this.rdb = db;
    NodeId n[] = new NodeId[N] ;
    Arrays.fill(n, NodeId.NodeIdAny);
    allPattern = TupleFactory.create(n);
}