Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.alibaba.jstorm.cache.rocksdb; import backtype.storm.utils.Utils; import com.alibaba.jstorm.client.ConfigExtension; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import org.apache.commons.io.FileUtils; import org.rocksdb.*; import org.slf4j.LoggerFactory; import storm.trident.operation.builtin.Max; import java.io.File; import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; /** * Factory class for {@link RocksDB}. */ public class RocksDbFactory { private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(RocksDbFactory.class); private static final String DEFAULT_COLUMN_FAMILY = "default"; /** * Returns a {@link RocksDB} */ public static RocksDB createDB(Map conf, String rocksDbDir) throws IOException { return create(conf, rocksDbDir, -1); } /** * Returns a {@link RocksDB} with ttl. */ public static RocksDB createTtlDB(Map conf, String rocksDbDir, int ttlTimeSec) throws IOException { return create(conf, rocksDbDir, ttlTimeSec); } public static RocksDB create(Map conf, String rocksDbDir, int ttlTimeSec) throws IOException { Options options = getOptions(conf); try { RocksDB rocksDb = ttlTimeSec > 0 ? TtlDB.open(options, rocksDbDir, ttlTimeSec, false) : RocksDB.open(options, rocksDbDir); LOG.info("Finished loading RocksDB"); // enable compaction rocksDb.compactRange(); return rocksDb; } catch (RocksDBException e) { throw new IOException("Failed to initialize RocksDb.", e); } } public static RocksDB createDBWithColumnFamily(Map conf, String dbPath, final Map<String, ColumnFamilyHandle> columnFamilyHandleMap) throws IOException { return createWithColumnFamily(conf, dbPath, columnFamilyHandleMap, -1); } public static RocksDB createTtlDBWithColumnFamily(Map conf, String dbPath, final Map<String, ColumnFamilyHandle> columnFamilyHandleMap, int ttlSec) throws IOException { return createWithColumnFamily(conf, dbPath, columnFamilyHandleMap, ttlSec); } public static RocksDB createWithColumnFamily(Map conf, String rocksDbDir, final Map<String, ColumnFamilyHandle> columnFamilyHandleMap, int ttlTimeSec) throws IOException { List<ColumnFamilyDescriptor> columnFamilyDescriptors = getExistingColumnFamilyDesc(conf, rocksDbDir); List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(); DBOptions dbOptions = getDBOptions(conf); try { RocksDB rocksDb = ttlTimeSec > 0 ? TtlDB.open(dbOptions, rocksDbDir, columnFamilyDescriptors, columnFamilyHandles, getTtlValues(ttlTimeSec, columnFamilyDescriptors), false) : RocksDB.open(dbOptions, rocksDbDir, columnFamilyDescriptors, columnFamilyHandles); int n = Math.min(columnFamilyDescriptors.size(), columnFamilyHandles.size()); // skip default column columnFamilyHandleMap.put(DEFAULT_COLUMN_FAMILY, rocksDb.getDefaultColumnFamily()); for (int i = 1; i < n; i++) { ColumnFamilyDescriptor descriptor = columnFamilyDescriptors.get(i); columnFamilyHandleMap.put(new String(descriptor.columnFamilyName()), columnFamilyHandles.get(i)); } LOG.info("Finished loading RocksDB with existing column family={}, dbPath={}, ttlSec={}", columnFamilyHandleMap.keySet(), rocksDbDir, ttlTimeSec); // enable compaction rocksDb.compactRange(); return rocksDb; } catch (RocksDBException e) { throw new IOException("Failed to initialize RocksDb.", e); } } private static List<Integer> getTtlValues(int ttlSec, List<ColumnFamilyDescriptor> descriptors) { List<Integer> ttlValues = Lists.newArrayList(); for (ColumnFamilyDescriptor descriptor : descriptors) { ttlValues.add(ttlSec); } return ttlValues; } private static List<ColumnFamilyDescriptor> getExistingColumnFamilyDesc(Map conf, String dbPath) throws IOException { try { List<byte[]> families = Lists.newArrayList(); List<byte[]> existingFamilies = RocksDB.listColumnFamilies(getOptions(conf), dbPath); if (existingFamilies != null) { families.addAll(existingFamilies); } else { families.add(RocksDB.DEFAULT_COLUMN_FAMILY); } List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>(); for (byte[] bytes : families) { columnFamilyDescriptors.add(new ColumnFamilyDescriptor(bytes, getColumnFamilyOptions(conf))); LOG.info("Load column family of {}", new String(bytes)); } return columnFamilyDescriptors; } catch (RocksDBException e) { throw new IOException("Failed to retrieve existing column families.", e); } } private RocksDbFactory() { } public static Options getOptions(Map conf) { Options options = (new RocksDbOptionsFactory.Defaults()).createOptions(null); String optionsFactoryClass = (String) conf.get(ConfigExtension.ROCKSDB_OPTIONS_FACTORY_CLASS); if (optionsFactoryClass != null) { RocksDbOptionsFactory udfOptionFactory = (RocksDbOptionsFactory) Utils.newInstance(optionsFactoryClass); options = udfOptionFactory.createOptions(options); } return options; } public static DBOptions getDBOptions(Map conf) { DBOptions dbOptions = (new RocksDbOptionsFactory.Defaults()).createDbOptions(null); String optionsFactoryClass = (String) conf.get(ConfigExtension.ROCKSDB_OPTIONS_FACTORY_CLASS); if (optionsFactoryClass != null) { RocksDbOptionsFactory udfOptionFactory = (RocksDbOptionsFactory) Utils.newInstance(optionsFactoryClass); dbOptions = udfOptionFactory.createDbOptions(dbOptions); } return dbOptions; } public static ColumnFamilyOptions getColumnFamilyOptions(Map conf) { ColumnFamilyOptions cfOptions = (new RocksDbOptionsFactory.Defaults()).createColumnFamilyOptions(null); String optionsFactoryClass = (String) conf.get(ConfigExtension.ROCKSDB_OPTIONS_FACTORY_CLASS); if (optionsFactoryClass != null) { RocksDbOptionsFactory udfOptionFactory = (RocksDbOptionsFactory) Utils.newInstance(optionsFactoryClass); cfOptions = udfOptionFactory.createColumnFamilyOptions(cfOptions); } return cfOptions; } public static void cleanRocksDbLocalDir(String rocksDbDir) throws IOException { File file = new File(rocksDbDir); if (file.exists()) { FileUtils.cleanDirectory(file); } else { FileUtils.forceMkdir(file); } } }