/**
* 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);
}
}
}