content: /*
   *
   * 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 org.apache.bookkeeper.bookie.storage.ldb;
  
  
  import static com.google.common.base.Preconditions.checkState;
  //CHECKSTYLE.OFF: IllegalImport
  //CHECKSTYLE.OFF: ImportOrder
  import static io.netty.util.internal.PlatformDependent.maxDirectMemory;
  //CHECKSTYLE.ON: IllegalImport
  //CHECKSTYLE.ON: ImportOrder
  
  import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
  import java.io.IOException;
  import java.nio.file.FileSystems;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.nio.file.Paths;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map.Entry;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.locks.ReentrantReadWriteLock;
  import org.apache.bookkeeper.bookie.storage.ldb.KeyValueStorageFactory.DbConfigType;
  import org.apache.bookkeeper.conf.ServerConfiguration;
  import org.apache.commons.lang3.StringUtils;
  import org.rocksdb.BlockBasedTableConfig;
  import org.rocksdb.BloomFilter;
  import org.rocksdb.Cache;
  import org.rocksdb.ChecksumType;
  import org.rocksdb.ColumnFamilyDescriptor;
  import org.rocksdb.ColumnFamilyHandle;
  import org.rocksdb.CompressionType;
  import org.rocksdb.ConfigOptions;
  import org.rocksdb.DBOptions;
  import org.rocksdb.Env;
  import org.rocksdb.InfoLogLevel;
  import org.rocksdb.LRUCache;
  import org.rocksdb.LiveFileMetaData;
  import org.rocksdb.Options;
  import org.rocksdb.OptionsUtil;
  import org.rocksdb.ReadOptions;
  import org.rocksdb.RocksDB;
  import org.rocksdb.RocksDBException;
  import org.rocksdb.RocksIterator;
  import org.rocksdb.RocksObject;
  import org.rocksdb.Slice;
  import org.rocksdb.WriteBatch;
  import org.rocksdb.WriteOptions;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * RocksDB based implementation of the KeyValueStorage.
   */
  public class KeyValueStorageRocksDB implements KeyValueStorage {
  
      static KeyValueStorageFactory factory = (defaultBasePath, subPath, dbConfigType, conf) ->
              new KeyValueStorageRocksDB(defaultBasePath, subPath, dbConfigType, conf);
  
      private volatile boolean closed = true;
      private final ReentrantReadWriteLock closedLock = new ReentrantReadWriteLock();
      private final RocksDB db;
      private RocksObject options;
      private List<ColumnFamilyDescriptor> columnFamilyDescriptors;
  
      private final WriteOptions optionSync;
      private final WriteOptions optionDontSync;
      private Cache cache;
  
      private final ReadOptions optionCache;
      private final ReadOptions optionDontCache;
      private final WriteBatch emptyBatch;
      private final int writeBatchMaxSize;
  
      private String dbPath;
  
      private static final String ROCKSDB_LOG_PATH = "dbStorage_rocksDB_logPath";
      private static final String ROCKSDB_LOG_LEVEL = "dbStorage_rocksDB_logLevel";
      private static final String ROCKSDB_LZ4_COMPRESSION_ENABLED = "dbStorage_rocksDB_lz4CompressionEnabled";
      private static final String ROCKSDB_WRITE_BUFFER_SIZE_MB = "dbStorage_rocksDB_writeBufferSizeMB";
      private static final String ROCKSDB_SST_SIZE_MB = "dbStorage_rocksDB_sstSizeInMB";
      private static final String ROCKSDB_BLOCK_SIZE = "dbStorage_rocksDB_blockSize";
      private static final String ROCKSDB_BLOOM_FILTERS_BITS_PER_KEY = "dbStorage_rocksDB_bloomFilterBitsPerKey";
      private static final String ROCKSDB_BLOCK_CACHE_SIZE = "dbStorage_rocksDB_blockCacheSize";
      private static final String ROCKSDB_NUM_LEVELS = "dbStorage_rocksDB_numLevels";
      private static final String ROCKSDB_NUM_FILES_IN_LEVEL0 = "dbStorage_rocksDB_numFilesInLevel0";
      private static final String ROCKSDB_MAX_SIZE_IN_LEVEL1_MB = "dbStorage_rocksDB_maxSizeInLevel1MB";
      private static final String ROCKSDB_FORMAT_VERSION = "dbStorage_rocksDB_format_version";
      private static final String ROCKSDB_CHECKSUM_TYPE = "dbStorage_rocksDB_checksum_type";
  
      public KeyValueStorageRocksDB(String basePath, String subPath, DbConfigType dbConfigType, ServerConfiguration conf)
              throws IOException {
          this(basePath, subPath, dbConfigType, conf, false);
      }
  
      public KeyValueStorageRocksDB(String basePath, String subPath, DbConfigType dbConfigType, ServerConfiguration conf,
                                    boolean readOnly)
              throws IOException {
          try {
              RocksDB.loadLibrary();
          } catch (Throwable t) {
              throw new IOException("Failed to load RocksDB JNI library", t);
          }
  
          this.optionSync = new WriteOptions();
          this.optionDontSync = new WriteOptions();
          this.optionCache = new ReadOptions();
          this.optionDontCache = new ReadOptions();
          this.emptyBatch = new WriteBatch();
  
          String dbFilePath = "";
          if (dbConfigType == DbConfigType.EntryLocation) {
              dbFilePath = conf.getEntryLocationRocksdbConf();
          } else if (dbConfigType == DbConfigType.LedgerMetadata) {
              dbFilePath = conf.getLedgerMetadataRocksdbConf();
          } else {
              dbFilePath = conf.getDefaultRocksDBConf();
          }
          log.info("Searching for a RocksDB configuration file in {}", dbFilePath);
          if (StringUtils.isNotBlank(dbFilePath) && Paths.get(dbFilePath).toFile().exists()) {
              log.info("Found a RocksDB configuration file and using it to initialize the RocksDB");
              db = initializeRocksDBWithConfFile(basePath, subPath, dbConfigType, conf, readOnly, dbFilePath);
          } else {
              log.info("Haven't found the file and read the configuration from the main bookkeeper configuration");
              db = initializeRocksDBWithBookieConf(basePath, subPath, dbConfigType, conf, readOnly);
          }
  
          optionSync.setSync(true);
          optionDontSync.setSync(false);
  
          optionCache.setFillCache(true);
          optionDontCache.setFillCache(false);
  
          this.writeBatchMaxSize = conf.getMaxOperationNumbersInSingleRocksDBBatch();
          this.closed = false;
      }
  
      private RocksDB initializeRocksDBWithConfFile(String basePath, String subPath, DbConfigType dbConfigType,
                                                 ServerConfiguration conf, boolean readOnly,
                                                 String dbFilePath) throws IOException {
          DBOptions dbOptions = new DBOptions();
          final List<ColumnFamilyDescriptor> cfDescs = new ArrayList<>();
          final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
          try (final ConfigOptions cfgOpts = new ConfigOptions()
                  .setIgnoreUnknownOptions(false)
                  .setEnv(Env.getDefault())) {
              OptionsUtil.loadOptionsFromFile(cfgOpts, dbFilePath, dbOptions, cfDescs);
              // Configure file path
              String logPath = conf.getString(ROCKSDB_LOG_PATH, "");
              if (!logPath.isEmpty()) {
                  Path logPathSetting = FileSystems.getDefault().getPath(logPath, subPath);
                  Files.createDirectories(logPathSetting);
                  log.info("RocksDB<{}> log path: {}", subPath, logPathSetting);
                  dbOptions.setDbLogDir(logPathSetting.toString());
              }
              this.dbPath = FileSystems.getDefault().getPath(basePath, subPath).toFile().toString();
              this.options = dbOptions;
              this.columnFamilyDescriptors = cfDescs;
              if (readOnly) {
                  return RocksDB.openReadOnly(dbOptions, dbPath, cfDescs, cfHandles);
              } else {
                  return RocksDB.open(dbOptions, dbPath, cfDescs, cfHandles);
              }
          } catch (RocksDBException e) {
              throw new IOException("Error open RocksDB database", e);
          }
      }
  
      private RocksDB initializeRocksDBWithBookieConf(String basePath, String subPath, DbConfigType dbConfigType,
                                             ServerConfiguration conf, boolean readOnly) throws IOException {
          Options options = new Options();
          options.setCreateIfMissing(true);
          ChecksumType checksumType = ChecksumType.valueOf(conf.getString(ROCKSDB_CHECKSUM_TYPE, "kxxHash"));
          int formatVersion = conf.getInt(ROCKSDB_FORMAT_VERSION, 5);
  
          if (dbConfigType == DbConfigType.EntryLocation) {
              /* Set default RocksDB block-cache size to 10% / numberOfLedgers of direct memory, unless override */
              int ledgerDirsSize = conf.getLedgerDirNames().length;
              long defaultRocksDBBlockCacheSizeBytes = maxDirectMemory() / ledgerDirsSize / 10;
              long blockCacheSize = DbLedgerStorage.getLongVariableOrDefault(conf, ROCKSDB_BLOCK_CACHE_SIZE,
                  defaultRocksDBBlockCacheSizeBytes);
  
              long writeBufferSizeMB = conf.getInt(ROCKSDB_WRITE_BUFFER_SIZE_MB, 64);
              long sstSizeMB = conf.getInt(ROCKSDB_SST_SIZE_MB, 64);
              int numLevels = conf.getInt(ROCKSDB_NUM_LEVELS, -1);
              int numFilesInLevel0 = conf.getInt(ROCKSDB_NUM_FILES_IN_LEVEL0, 4);
              long maxSizeInLevel1MB = conf.getLong(ROCKSDB_MAX_SIZE_IN_LEVEL1_MB, 256);
              int blockSize = conf.getInt(ROCKSDB_BLOCK_SIZE, 64 * 1024);
              int bloomFilterBitsPerKey = conf.getInt(ROCKSDB_BLOOM_FILTERS_BITS_PER_KEY, 10);
              boolean lz4CompressionEnabled = conf.getBoolean(ROCKSDB_LZ4_COMPRESSION_ENABLED, true);
  
              if (lz4CompressionEnabled) {
                  options.setCompressionType(CompressionType.LZ4_COMPRESSION);
              }
              options.setWriteBufferSize(writeBufferSizeMB * 1024 * 1024);
              options.setMaxWriteBufferNumber(4);
              if (numLevels > 0) {
                  options.setNumLevels(numLevels);
              }
              options.setLevelZeroFileNumCompactionTrigger(numFilesInLevel0);
              options.setMaxBytesForLevelBase(maxSizeInLevel1MB * 1024 * 1024);
              options.setMaxBackgroundJobs(32);
              options.setIncreaseParallelism(32);
              options.setMaxTotalWalSize(512 * 1024 * 1024);
              options.setMaxOpenFiles(-1);
              options.setTargetFileSizeBase(sstSizeMB * 1024 * 1024);
              options.setDeleteObsoleteFilesPeriodMicros(TimeUnit.HOURS.toMicros(1));
  
              this.cache = new LRUCache(blockCacheSize);
              BlockBasedTableConfig tableOptions = new BlockBasedTableConfig();
              tableOptions.setBlockSize(blockSize);
              tableOptions.setBlockCache(cache);
              tableOptions.setFormatVersion(formatVersion);
              tableOptions.setChecksumType(checksumType);
              if (bloomFilterBitsPerKey > 0) {
                  tableOptions.setFilterPolicy(new BloomFilter(bloomFilterBitsPerKey, false));
              }
  
              // Options best suited for HDDs
              tableOptions.setCacheIndexAndFilterBlocks(true);
              options.setLevelCompactionDynamicLevelBytes(true);
  
              options.setTableFormatConfig(tableOptions);
          } else {
              this.cache = null;
              BlockBasedTableConfig tableOptions = new BlockBasedTableConfig();
              tableOptions.setFormatVersion(formatVersion);
              tableOptions.setChecksumType(checksumType);
              options.setTableFormatConfig(tableOptions);
          }
  
              // Configure file path
          String logPath = conf.getString(ROCKSDB_LOG_PATH, "");
          if (!logPath.isEmpty()) {
              Path logPathSetting = FileSystems.getDefault().getPath(logPath, subPath);
              Files.createDirectories(logPathSetting);
              log.info("RocksDB<{}> log path: {}", subPath, logPathSetting);
              options.setDbLogDir(logPathSetting.toString());
          }
          this.dbPath = FileSystems.getDefault().getPath(basePath, subPath).toFile().toString();
  
          // Configure log level
          String logLevel = conf.getString(ROCKSDB_LOG_LEVEL, "info");
          switch (logLevel) {
              case "debug":
                  options.setInfoLogLevel(InfoLogLevel.DEBUG_LEVEL);
                  break;
              case "info":
                  options.setInfoLogLevel(InfoLogLevel.INFO_LEVEL);
                  break;
              case "warn":
                  options.setInfoLogLevel(InfoLogLevel.WARN_LEVEL);
                  break;
              case "error":
                  options.setInfoLogLevel(InfoLogLevel.ERROR_LEVEL);
                  break;
              default:
                  log.warn("Unrecognized RockDB log level: {}", logLevel);
          }
  
              // Keep log files for 1month
          options.setKeepLogFileNum(30);
          options.setLogFileTimeToRoll(TimeUnit.DAYS.toSeconds(1));
          this.options = options;
          try {
              if (readOnly) {
                  return RocksDB.openReadOnly(options, dbPath);
              } else {
                  return RocksDB.open(options, dbPath);
              }
          } catch (RocksDBException e) {
              throw new IOException("Error open RocksDB database", e);
          }
      }
  
      @Override
      public void close() throws IOException {
          try {
              closedLock.writeLock().lock();
              closed = true;
              db.close();
          } finally {
              closedLock.writeLock().unlock();
          }
          if (cache != null) {
              cache.close();
          }
          if (options != null) {
              options.close();
          }
          optionSync.close();
          optionDontSync.close();
          optionCache.close();
          optionDontCache.close();
          emptyBatch.close();
      }
  
      @Override
      public void put(byte[] key, byte[] value) throws IOException {
          try {
              db.put(optionDontSync, key, value);
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB put", e);
          }
      }
  
      @Override
      public byte[] get(byte[] key) throws IOException {
          try {
              return db.get(key);
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB get", e);
          }
      }
  
      @Override
      public int get(byte[] key, byte[] value) throws IOException {
          try {
              int res = db.get(key, value);
              if (res == RocksDB.NOT_FOUND) {
                  return -1;
              } else if (res > value.length) {
                  throw new IOException("Value array is too small to fit the result");
              } else {
                  return res;
              }
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB get", e);
          }
      }
  
      @Override
      @SuppressFBWarnings("RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE")
      public Entry<byte[], byte[]> getFloor(byte[] key) throws IOException {
          try (Slice upperBound = new Slice(key);
                   ReadOptions option = new ReadOptions(optionCache).setIterateUpperBound(upperBound);
                   RocksIterator iterator = db.newIterator(option)) {
              iterator.seekToLast();
              if (iterator.isValid()) {
                  return new EntryWrapper(iterator.key(), iterator.value());
              }
          }
          return null;
      }
  
      @Override
      @SuppressFBWarnings("RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE")
      public Entry<byte[], byte[]> getCeil(byte[] key) throws IOException {
          try (RocksIterator iterator = db.newIterator(optionCache)) {
              // Position the iterator on the record whose key is >= to the supplied key
              iterator.seek(key);
  
              if (iterator.isValid()) {
                  return new EntryWrapper(iterator.key(), iterator.value());
              } else {
                  return null;
              }
          }
      }
  
      @Override
      public void delete(byte[] key) throws IOException {
          try {
              db.delete(optionDontSync, key);
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB delete", e);
          }
      }
  
      @Override
      public String getDBPath() {
          return dbPath;
      }
  
      @Override
      public void compact(byte[] firstKey, byte[] lastKey) throws IOException {
          try {
              db.compactRange(firstKey, lastKey);
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB compact", e);
          }
      }
  
      @Override
      public void compact() throws IOException {
          try {
              final long start = System.currentTimeMillis();
              final int oriRocksDBFileCount = db.getLiveFilesMetaData().size();
              final long oriRocksDBSize = getRocksDBSize();
              log.info("Starting RocksDB {} compact, current RocksDB hold {} files and {} Bytes.",
                      db.getName(), oriRocksDBFileCount, oriRocksDBSize);
  
              db.compactRange();
  
              final long end = System.currentTimeMillis();
              final int rocksDBFileCount = db.getLiveFilesMetaData().size();
              final long rocksDBSize = getRocksDBSize();
              log.info("RocksDB {} compact finished {} ms, space reduced {} Bytes, current hold {} files and {} Bytes.",
                      db.getName(), end - start, oriRocksDBSize - rocksDBSize, rocksDBFileCount, rocksDBSize);
          } catch (RocksDBException e) {
              throw new IOException("Error in RocksDB compact", e);
          }
      }
  
      private long getRocksDBSize() {
          List<LiveFileMetaData> liveFilesMetaData = db.getLiveFilesMetaData();
          long rocksDBFileSize = 0L;
          for (LiveFileMetaData fileMetaData : liveFilesMetaData) {
              rocksDBFileSize += fileMetaData.size();
          }
          return rocksDBFileSize;
      }
  
      @Override
      public void sync() throws IOException {
          try {
              db.write(optionSync, emptyBatch);
          } catch (RocksDBException e) {
              throw new IOException(e);
          }
      }
  
      @Override
      public CloseableIterator<byte[]> keys() {
          final RocksIterator iterator = db.newIterator(optionCache);
          iterator.seekToFirst();
  
          return new CloseableIterator<byte[]>() {
              @Override
              public boolean hasNext() {
                  return iterator.isValid();
              }
  
              @Override
              public byte[] next() {
                  checkState(iterator.isValid());
                  byte[] key = iterator.key();
                  iterator.next();
                  return key;
              }
  
              @Override
              public void close() {
                  iterator.close();
              }
          };
      }
  
      @Override
      public CloseableIterator<byte[]> keys(byte[] firstKey, byte[] lastKey) {
          final Slice upperBound = new Slice(lastKey);
          final ReadOptions option = new ReadOptions(optionCache).setIterateUpperBound(upperBound);
          final RocksIterator iterator = db.newIterator(option);
          iterator.seek(firstKey);
  
          return new CloseableIterator<byte[]>() {
              @Override
              public boolean hasNext() {
                  return iterator.isValid();
              }
  
              @Override
              public byte[] next() {
                  checkState(iterator.isValid());
                  byte[] key = iterator.key();
                  iterator.next();
                  return key;
              }
  
              @Override
              public void close() {
                  iterator.close();
                  option.close();
                  upperBound.close();
              }
          };
      }
  
      @Override
      public CloseableIterator<Entry<byte[], byte[]>> iterator() {
          final RocksIterator iterator = db.newIterator(optionDontCache);
          iterator.seekToFirst();
          final EntryWrapper entryWrapper = new EntryWrapper();
  
          return new CloseableIterator<Entry<byte[], byte[]>>() {
              @Override
              public boolean hasNext() {
                  return iterator.isValid();
              }
  
              @Override
              public Entry<byte[], byte[]> next() {
                  checkState(iterator.isValid());
                  entryWrapper.key = iterator.key();
                  entryWrapper.value = iterator.value();
                  iterator.next();
                  return entryWrapper;
              }
  
              @Override
              public void close() {
                  iterator.close();
              }
          };
      }
  
      @Override
      public long count() throws IOException {
          try {
              if (closed) {
                  throw new IOException("RocksDB is closed");
              }
              try {
                  closedLock.readLock().lock();
                  if (!closed) {
                      return db.getLongProperty("rocksdb.estimate-num-keys");
                  }
                  throw new IOException("RocksDB is closed");
              } finally {
                  closedLock.readLock().unlock();
              }
          } catch (RocksDBException e) {
              throw new IOException("Error in getting records count", e);
          }
      }
  
      @Override
      public Batch newBatch() {
          return new RocksDBBatch(writeBatchMaxSize);
      }
  
      private class RocksDBBatch implements Batch {
          private final WriteBatch writeBatch = new WriteBatch();
          private final int batchSize;
          private int batchCount = 0;
  
          RocksDBBatch(int batchSize) {
              this.batchSize = batchSize;
          }
  
          @Override
          public void close() {
              writeBatch.close();
              batchCount = 0;
          }
  
          @Override
          public void put(byte[] key, byte[] value) throws IOException {
              try {
                  writeBatch.put(key, value);
                  countBatchAndFlushIfNeeded();
              } catch (RocksDBException e) {
                  throw new IOException("Failed to flush RocksDB batch", e);
              }
          }
  
          @Override
          public void remove(byte[] key) throws IOException {
              try {
                  writeBatch.delete(key);
                  countBatchAndFlushIfNeeded();
              } catch (RocksDBException e) {
                  throw new IOException("Failed to flush RocksDB batch", e);
              }
          }
  
          @Override
          public void clear() {
              writeBatch.clear();
              batchCount = 0;
          }
  
          @Override
          public void deleteRange(byte[] beginKey, byte[] endKey) throws IOException {
              try {
                  writeBatch.deleteRange(beginKey, endKey);
                  countBatchAndFlushIfNeeded();
              } catch (RocksDBException e) {
                  throw new IOException("Failed to flush RocksDB batch", e);
              }
          }
  
          private void countBatchAndFlushIfNeeded() throws IOException {
              if (++batchCount >= batchSize) {
                  flush();
                  clear();
              }
          }
  
          @Override
          public int batchCount() {
              return batchCount;
          }
  
          @Override
          public void flush() throws IOException {
              try {
                  db.write(optionSync, writeBatch);
              } catch (RocksDBException e) {
                  throw new IOException("Failed to flush RocksDB batch", e);
              }
          }
      }
  
      private static final class EntryWrapper implements Entry<byte[], byte[]> {
          // This is not final since the iterator will reuse the same EntryWrapper
          // instance at each step
          private byte[] key;
          private byte[] value;
  
          public EntryWrapper() {
              this.key = null;
              this.value = null;
          }
  
          public EntryWrapper(byte[] key, byte[] value) {
              this.key = key;
              this.value = value;
          }
  
          @Override
          public byte[] setValue(byte[] value) {
              throw new UnsupportedOperationException();
          }
  
          @Override
          public byte[] getValue() {
              return value;
          }
  
          @Override
          public byte[] getKey() {
              return key;
          }
      }
  
      RocksDB db() {
          return db;
      }
  
      List<ColumnFamilyDescriptor> getColumnFamilyDescriptors() {
          return columnFamilyDescriptors;
      }
  
      RocksObject getOptions() {
          return options;
      }
  
      private static final Logger log = LoggerFactory.getLogger(KeyValueStorageRocksDB.class);
  }
  