stream/statelib/src/main/java/org/apache/bookkeeper/statelib/impl/kv/RocksdbKVStore.java - bookkeeper - Git at Google

 /*
  * 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.statelib.impl.kv;

 import static com.google.common.base.Preconditions.checkNotNull;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.BLOCK_CACHE_SIZE;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.BLOCK_SIZE;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_CHECKSUM_TYPE;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_COMPACTION_STYLE;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_COMPRESSION_TYPE;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_LOG_LEVEL;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_PARALLELISM;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.MAX_WRITE_BUFFERS;
 import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.WRITE_BUFFER_SIZE;

 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.google.common.io.MoreFiles;
 import com.google.common.io.RecursiveDeleteOption;
 import com.google.common.primitives.SignedBytes;
 import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicLongFieldUpdater;
 import lombok.AccessLevel;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.common.coder.Coder;
 import org.apache.bookkeeper.common.kv.KV;
 import org.apache.bookkeeper.common.kv.KVImpl;
 import org.apache.bookkeeper.statelib.api.StateStoreSpec;
 import org.apache.bookkeeper.statelib.api.checkpoint.CheckpointStore;
 import org.apache.bookkeeper.statelib.api.exceptions.InvalidStateStoreException;
 import org.apache.bookkeeper.statelib.api.exceptions.StateStoreException;
 import org.apache.bookkeeper.statelib.api.exceptions.StateStoreRuntimeException;
 import org.apache.bookkeeper.statelib.api.kv.KVIterator;
 import org.apache.bookkeeper.statelib.api.kv.KVMulti;
 import org.apache.bookkeeper.statelib.api.kv.KVStore;
 import org.apache.bookkeeper.statelib.impl.Bytes;
 import org.apache.bookkeeper.statelib.impl.rocksdb.RocksUtils;
 import org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint.CheckpointInfo;
 import org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint.RocksCheckpointer;
 import org.apache.commons.lang3.tuple.Pair;
 import org.rocksdb.BlockBasedTableConfig;
 import org.rocksdb.Cache;
 import org.rocksdb.ColumnFamilyDescriptor;
 import org.rocksdb.ColumnFamilyHandle;
 import org.rocksdb.ColumnFamilyOptions;
 import org.rocksdb.DBOptions;
 import org.rocksdb.FlushOptions;
 import org.rocksdb.LRUCache;
 import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import org.rocksdb.RocksIterator;
 import org.rocksdb.WriteBatch;
 import org.rocksdb.WriteOptions;

 /**
  * A key/value store implemented using {@link http://rocksdb.org/}.
  *
  * @param <K> key type
  * @param <V> value type
  */
 @Slf4j
 public class RocksdbKVStore<K, V> implements KVStore<K, V> {

     private static final byte[] METADATA_CF = ".meta".getBytes(UTF_8);
     private static final byte[] DATA_CF = "default".getBytes(UTF_8);
     private static final byte[] LAST_REVISION = ".lrev".getBytes(UTF_8);

     private static final AtomicLongFieldUpdater<RocksdbKVStore> lastRevisionUpdater =
         AtomicLongFieldUpdater.newUpdater(RocksdbKVStore.class, "lastRevision");

     // parameters for the store
     protected String name;
     protected Coder<K> keyCoder;
     protected Coder<V> valCoder;

     // rocksdb state
     protected File dbDir;
     protected RocksDB db;
     protected ColumnFamilyHandle metaCfHandle;
     protected ColumnFamilyHandle dataCfHandle;

     // iterators
     protected final Set<KVIterator<K, V>> kvIters;

     // options used by rocksdb
     protected DBOptions dbOpts;
     protected ColumnFamilyOptions cfOpts;
     protected WriteOptions writeOpts;
     protected FlushOptions flushOpts;

     // states of the store
     protected volatile boolean isInitialized = false;
     protected volatile boolean closed = false;
     protected volatile long lastRevision = -1L;
     private final byte[] lastRevisionBytes = new byte[Long.BYTES];

     // checkpointer store
     private CheckpointStore checkpointStore;
     private ScheduledExecutorService checkpointScheduler;
     // rocksdb checkpointer
     private RocksCheckpointer checkpointer;

     static {
         RocksDB.loadLibrary();
     }

     private boolean cleanupLocalStoreDirEnable;

     public RocksdbKVStore() {
         // initialize the iterators set
         this.kvIters = Collections.synchronizedSet(Sets.newHashSet());
     }

     protected void checkStoreOpen() {
         if (closed) {
             throw new InvalidStateStoreException("State store " + name + " is already closed");
         }
         if (!isInitialized) {
             throw new InvalidStateStoreException("State Store " + name + " is not initialized yet");
         }
     }

     @VisibleForTesting
     public synchronized RocksDB getDb() {
         return db;
     }

     @Override
     public synchronized String name() {
         return this.name;
     }

     private void loadRocksdbFromCheckpointStore(StateStoreSpec spec) {
         checkNotNull(spec.getCheckpointIOScheduler(),
             "checkpoint io scheduler is not configured");
         checkNotNull(spec.getCheckpointDuration(),
             "checkpoint duration is not configured");

         String dbName = spec.getName();
         File localStorePath = spec.getLocalStateStoreDir();

         List<CheckpointInfo> checkpoints = RocksCheckpointer.getCheckpoints(dbName, spec.getCheckpointStore());
         for (CheckpointInfo cpi : checkpoints) {
             try {
                 cpi.restore(dbName, localStorePath, spec.getCheckpointStore());
                 openRocksdb(spec);
                 checkpoints.stream()
                     .filter(cp -> cp != cpi) // ignore the current restored checkpoint
                     .forEach(cp -> cp.remove(localStorePath)); // delete everything else
                 break;
             } catch (StateStoreException e) {
                 // Got an exception. Log and try the next checkpoint
                 log.error("Failed to restore checkpoint: {}", cpi, e);
             }
         }
     }

     @Override
     public synchronized void checkpoint() {
         log.info("Checkpoint local state store {} at revision {}", name, getLastRevision());
         byte[] checkpointAtRevisionBytes = new byte[Long.BYTES];
         System.arraycopy(lastRevisionBytes, 0, checkpointAtRevisionBytes, 0, checkpointAtRevisionBytes.length);
         checkpointScheduler.submit(() -> {
             try {
                 // TODO: move create checkpoint to the checkpoint method
                 checkpointer.checkpointAtTxid(checkpointAtRevisionBytes);
             } catch (StateStoreException e) {
                 log.error("Failed to checkpoint state store {} at revision {}",
                     name, Bytes.toLong(checkpointAtRevisionBytes, 0), e);
             }
         });
     }

     private void readLastRevision() throws StateStoreException {
         byte[] revisionBytes;
         try {
             revisionBytes = db.get(metaCfHandle, LAST_REVISION);
         } catch (RocksDBException e) {
             throw new StateStoreException("Failed to read last revision from state store " + name(), e);
         }
         if (null == revisionBytes) {
             return;
         }
         long revision = Bytes.toLong(revisionBytes, 0);
         lastRevisionUpdater.set(this, revision);
     }

     @Override
     public long getLastRevision() {
         return lastRevisionUpdater.get(this);
     }

     private void setLastRevision(long lastRevision) {
         lastRevisionUpdater.set(this, lastRevision);
         Bytes.toBytes(lastRevision, lastRevisionBytes, 0);
     }

     private void updateLastRevision(long revision) {
         if (revision >= 0) { // k/v comes from log stream
             if (getLastRevision() >= revision) { // these k/v pairs are duplicates
                 return;
             }
             // update revision
             setLastRevision(revision);
         }
     }

     protected void updateLastRevision(WriteBatch batch, long revision) {
         if (revision >= 0) { // k/v comes from log stream
             if (getLastRevision() >= revision) { // these k/v pairs are duplicates
                 return;
             }
             try {
                 // update revision
                 setLastRevision(revision);
                 batch.put(metaCfHandle, LAST_REVISION, lastRevisionBytes);
             } catch (RocksDBException e) {
                 throw new StateStoreRuntimeException(
                         "Error while updating last revision " + revision + " from store " + name, e);
             }
         }
     }

     @Override
     @SuppressWarnings("unchecked")
     public synchronized void init(StateStoreSpec spec) throws StateStoreException {
         checkNotNull(spec.getLocalStateStoreDir(),
             "local state store directory is not configured");

         this.name = spec.getName();
         this.cleanupLocalStoreDirEnable = spec.isLocalStorageCleanupEnable();

         // initialize the coders
         this.keyCoder = (Coder<K>) spec.getKeyCoder();
         this.valCoder = (Coder<V>) spec.getValCoder();

         cleanupLocalStoreDir(spec.getLocalStateStoreDir());

         checkpointStore = spec.getCheckpointStore();
         if (null != checkpointStore) {
             // load checkpoint from checkpoint store
             loadRocksdbFromCheckpointStore(spec);
         } else {
             // open the rocksdb
             openRocksdb(spec);
         }


         // once the rocksdb is opened, read the last revision
         readLastRevision();

         if (null != checkpointStore) {
             checkpointer = new RocksCheckpointer(
                 name(),
                 dbDir,
                 db,
                 checkpointStore,
                 true,
                 true,
                 spec.isCheckpointChecksumEnable(),
                 spec.isCheckpointChecksumCompatible());
             checkpointScheduler = spec.getCheckpointIOScheduler();
         }

         this.isInitialized = true;
     }

     protected void openRocksdb(StateStoreSpec spec) throws StateStoreException {

         // initialize the db options

         final BlockBasedTableConfig tableConfig = new BlockBasedTableConfig();
         final Cache cache = new LRUCache(BLOCK_CACHE_SIZE);
         tableConfig.setBlockCache(cache);
         tableConfig.setBlockSize(BLOCK_SIZE);
         tableConfig.setChecksumType(DEFAULT_CHECKSUM_TYPE);

         dbOpts = new DBOptions();
         dbOpts.setCreateIfMissing(true);
         dbOpts.setErrorIfExists(false);
         dbOpts.setInfoLogLevel(DEFAULT_LOG_LEVEL);
         dbOpts.setIncreaseParallelism(DEFAULT_PARALLELISM);
         dbOpts.setCreateMissingColumnFamilies(true);

         cfOpts = new ColumnFamilyOptions();
         cfOpts.setTableFormatConfig(tableConfig);
         cfOpts.setWriteBufferSize(WRITE_BUFFER_SIZE);
         cfOpts.setCompressionType(DEFAULT_COMPRESSION_TYPE);
         cfOpts.setCompactionStyle(DEFAULT_COMPACTION_STYLE);
         cfOpts.setMaxWriteBufferNumber(MAX_WRITE_BUFFERS);

         // initialize the write options

         writeOpts = new WriteOptions();
         writeOpts.setDisableWAL(true); // disable wal, since the source of truth will be on distributedlog

         // initialize the flush options

         flushOpts = new FlushOptions();
         flushOpts.setWaitForFlush(true);

         // open the rocksdb

         this.dbDir = spec.getLocalStateStoreDir();
         Pair<RocksDB, List<ColumnFamilyHandle>> dbPair = openLocalDB(dbDir, dbOpts, cfOpts);
         this.db = dbPair.getLeft();
         this.metaCfHandle = dbPair.getRight().get(0);
         this.dataCfHandle = dbPair.getRight().get(1);
     }

     protected Pair<RocksDB, List<ColumnFamilyHandle>> openLocalDB(File dir,
                                                                   DBOptions options,
                                                                   ColumnFamilyOptions cfOpts)
         throws StateStoreException {
         return openRocksdb(dir, options, cfOpts);
     }

     protected static Pair<RocksDB, List<ColumnFamilyHandle>> openRocksdb(
         File dir, DBOptions options, ColumnFamilyOptions cfOpts)
         throws StateStoreException {
         // make sure the db directory's parent dir is created
         ColumnFamilyDescriptor metaDesc = new ColumnFamilyDescriptor(METADATA_CF, cfOpts);
         ColumnFamilyDescriptor dataDesc = new ColumnFamilyDescriptor(DATA_CF, cfOpts);

         try {
             Files.createDirectories(dir.toPath());
             File dbDir = new File(dir, "current");

             if (!dbDir.exists()) {
                 // empty state
                 String uuid = UUID.randomUUID().toString();
                 Path checkpointPath = Paths.get(dir.getAbsolutePath(), "checkpoints", uuid);
                 Files.createDirectories(checkpointPath);
                 Files.createSymbolicLink(
                     Paths.get(dbDir.getAbsolutePath()),
                     checkpointPath);
             }

             List<ColumnFamilyHandle> cfHandles = Lists.newArrayListWithExpectedSize(2);
             RocksDB db = RocksDB.open(
                 options,
                 dbDir.getAbsolutePath(),
                 Lists.newArrayList(metaDesc, dataDesc),
                 cfHandles);
             return Pair.of(db, cfHandles);
         } catch (IOException ioe) {
             log.error("Failed to create parent directory {} for opening rocksdb", dir.getParentFile().toPath(), ioe);
             throw new StateStoreException(ioe);
         } catch (RocksDBException dbe) {
             log.error("Failed to open rocksdb at dir {}", dir.getAbsolutePath(), dbe);
             throw new StateStoreException(dbe);
         }
     }

     @Override
     public synchronized void flush() throws StateStoreException {
         if (null == db) {
             return;
         }
         try {
             db.flush(flushOpts);
         } catch (RocksDBException e) {
             throw new StateStoreException("Exception on flushing rocksdb from store " + name, e);
         }
     }

     @Override
     public synchronized void close() {
         if (closed) {
             return;
         }
         closed = true;

         if (null != checkpointer) {
             checkpointer.close();
         }

         // close iterators
         closeIters();

         // close db
         closeLocalDB();

         // release options
         RocksUtils.close(dbOpts);
         RocksUtils.close(writeOpts);
         RocksUtils.close(flushOpts);
         RocksUtils.close(cfOpts);

         cleanupLocalStoreDir(dbDir);
     }

     private void cleanupLocalStoreDir(File dbDir) {
         if (cleanupLocalStoreDirEnable) {
             if (dbDir.exists()) {
                 try {
                     MoreFiles.deleteRecursively(dbDir.toPath(), RecursiveDeleteOption.ALLOW_INSECURE);
                 } catch (IOException e) {
                     log.error("Failed to cleanup localStoreDir", e);
                 }
             }
         }
     }

     protected void closeLocalDB() {
         try {
             flush();
         } catch (StateStoreException e) {
             // flush() already logs this exception.
         }

         RocksUtils.close(metaCfHandle);
         RocksUtils.close(dataCfHandle);
         RocksUtils.close(db);
     }

     private void closeIters() {
         Set<KVIterator> iterators;
         synchronized (kvIters) {
             iterators = Sets.newHashSet(kvIters);
         }
         iterators.forEach(KVIterator::close);
     }

     @Override
     public synchronized V get(K key) {
         checkNotNull(key, "key cannot be null");
         checkStoreOpen();

         byte[] keyBytes = keyCoder.encode(key);
         return getRaw(key, keyBytes);
     }

     private V getRaw(K key, byte[] keyBytes) {
         byte[] valBytes = getRawBytes(key, keyBytes);
         if (null == valBytes) {
             return null;
         }
         return valCoder.decode(valBytes);
     }

     protected byte[] getRawBytes(K key, byte[] keyBytes) {
         try {
             return this.db.get(dataCfHandle, keyBytes);
         } catch (RocksDBException e) {
             throw new StateStoreRuntimeException("Error while getting value for key " + key + " from store " + name, e);
         }
     }

     @Override
     public synchronized KVIterator<K, V> range(K from, K to) {
         checkStoreOpen();

         RocksIterator rocksIter = db.newIterator(dataCfHandle);
         if (null == from) {
             rocksIter.seekToFirst();
         } else {
             byte[] fromBytes = keyCoder.encode(from);
             rocksIter.seek(fromBytes);
         }
         KVIterator<K, V> kvIter;
         if (null == to) {
             kvIter = new RocksdbKVIterator(name, rocksIter, keyCoder, valCoder);
         } else {
             kvIter = new RocksdbRangeIterator(name, rocksIter, keyCoder, valCoder, to);
         }
         kvIters.add(kvIter);
         return kvIter;
     }

     @Override
     public synchronized void put(K key, V value) {
         put(key, value, -1);
     }

     synchronized void put(K key, V value, long revision) {
         checkNotNull(key, "key cannot be null");
         checkStoreOpen();

         updateLastRevision(revision);

         byte[] keyBytes = keyCoder.encode(key);
         putRaw(key, keyBytes, value, revision);
     }

     private void putRaw(K key, byte[] keyBytes, V value, long revision) {
         try {
             WriteBatch batch = new WriteBatch();
             if (revision > 0) {
                 // last revision has been set to revision bytes
                 batch.put(metaCfHandle, LAST_REVISION, lastRevisionBytes);
             }
             if (null == value) {
                 // delete a key if value is null
                 batch.delete(dataCfHandle, keyBytes);
             } else {
                 byte[] valBytes = valCoder.encode(value);
                 batch.put(dataCfHandle, keyBytes, valBytes);
             }

             db.write(writeOpts, batch);
         } catch (RocksDBException e) {
             throw new StateStoreRuntimeException("Error while updating key " + key
                 + " to value " + value + " from store " + name, e);
         }
     }

     @Override
     public V putIfAbsent(K key, V value) {
         return putIfAbsent(key, value, -1L);
     }

     synchronized V putIfAbsent(K key, V value, long revision) {
         checkNotNull(key, "key cannot be null");
         checkStoreOpen();

         updateLastRevision(revision);

         byte[] keyBytes = keyCoder.encode(key);
         V oldVal = getRaw(key, keyBytes);
         if (null != oldVal) {
             return oldVal;
         }

         if (value == null) {
             return null;
         }

         putRaw(key, keyBytes, value, revision);
         return null;
     }

     @Override
     public synchronized KVMulti<K, V> multi() {
         checkStoreOpen();

         return new KVMultiImpl();
     }

     @Override
     public synchronized V delete(K key) {
         return delete(key, -1L);
     }

     synchronized V delete(K key, long revision) {
         checkNotNull(key, "key cannot be null");
         checkStoreOpen();

         updateLastRevision(revision);

         byte[] keyBytes = keyCoder.encode(key);
         V val = getRaw(key, keyBytes);
         putRaw(key, keyBytes, null, revision);

         return val;
     }

     //
     // Multi
     //

     /**
      * A rocksdb based multi operation.
      */
     class KVMultiImpl implements KVMulti<K, V> {

         private final WriteBatch batch = new WriteBatch();
         private volatile boolean executed = false;

         private void checkExecuted() {
             if (executed) {
                 throw new StateStoreRuntimeException("KVMulti#execute() has been called");
             }
         }

         @Override
         public void put(K key, V value) {
             checkNotNull(key, "key cannot be null");
             checkExecuted();

             byte[] keyBytes = keyCoder.encode(key);
             if (null == value) {
                 deleteRaw(keyBytes);
             } else {
                 putRaw(keyBytes, value);
             }
         }

         private void putRaw(byte[] keyBytes, V value) {
             try {
                 batch.put(dataCfHandle, keyBytes, valCoder.encode(value));
             } catch (RocksDBException e) {
                 throw new StateStoreRuntimeException(e);
             }
         }

         @Override
         public void delete(K key) {
             checkNotNull(key, "key cannot be null");
             checkExecuted();

             byte[] keyBytes = keyCoder.encode(key);
             deleteRaw(keyBytes);
         }

         private void deleteRaw(byte[] keyBytes) {
             try {
                 batch.delete(dataCfHandle, keyBytes);
             } catch (RocksDBException e) {
                 throw new StateStoreRuntimeException(e);
             }
         }

         @Override
         public void deleteRange(K from, K to) {
             checkNotNull(from, "from key cannot be null");
             checkNotNull(to, "to key cannot be null");
             checkExecuted();

             byte[] fromBytes = keyCoder.encode(from);
             byte[] toBytes = keyCoder.encode(to);
             try {
                 batch.deleteRange(dataCfHandle, fromBytes, toBytes);
             } catch (RocksDBException e) {
                 throw new StateStoreRuntimeException(e);
             }
         }

         @Override
         public synchronized void execute() {
             if (executed) {
                 return;
             }

             checkStoreOpen();
             executed = true;

             try {
                 getDb().write(writeOpts, batch);
             } catch (RocksDBException e) {
                 throw new StateStoreRuntimeException("Error while executing a multi operation from store " + name, e);
             } finally {
                 RocksUtils.close(batch);
             }
         }
     }

     //
     // Iterators
     //

     /**
      * KV iterator over a rocksdb instance.
      */
     @RequiredArgsConstructor(access = AccessLevel.PRIVATE)
     class RocksdbKVIterator implements KVIterator<K, V> {

         final String name;
         final RocksIterator iterator;
         final Coder<K> keyCoder;
         final Coder<V> valCoder;

         private volatile boolean closed = false;

         /**
          * Ensure an iterator is open.
          *
          * @throws InvalidStateStoreException when the store is in closed state.
          */
         private void ensureIteratorOpen() {
             if (closed) {
                 throw new InvalidStateStoreException("Rocksdb state store " + name + " is already closed");
             }
         }

         @Override
         public void close() {
             kvIters.remove(this);
             iterator.close();
             closed = true;
         }

         @Override
         public boolean hasNext() {
             ensureIteratorOpen();
             return iterator.isValid();
         }

         private KV<K, V> getKvPair() {
             return new KVImpl<>(keyCoder.decode(iterator.key()), valCoder.decode(iterator.value()));
         }

         @Override
         public KV<K, V> next() {
             if (!hasNext()) {
                 throw new NoSuchElementException();
             }

             KV<K, V> kv = getKvPair();
             iterator.next();
             return kv;
         }
     }

     /**
      * KV iterator over a rocksdb, which ends at the provided <i>endKey</i>.
      */
     class RocksdbRangeIterator extends RocksdbKVIterator {

         private final Comparator<byte[]> comparator = SignedBytes.lexicographicalComparator();
         private final byte[] endKeyBytes;

         private RocksdbRangeIterator(String name,
                                      RocksIterator iterator,
                                      Coder<K> keyCoder,
                                      Coder<V> valCoder,
                                      K endKey) {
             super(name, iterator, keyCoder, valCoder);
             checkNotNull(endKey, "End key cannot be null");
             this.endKeyBytes = keyCoder.encode(endKey);
         }


         @Override
         public boolean hasNext() {
             return super.hasNext()
                 && comparator.compare(iterator.key(), endKeyBytes) <= 0;
         }
     }


 }
	/*
	* 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.statelib.impl.kv;

	import static com.google.common.base.Preconditions.checkNotNull;
	import static java.nio.charset.StandardCharsets.UTF_8;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.BLOCK_CACHE_SIZE;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.BLOCK_SIZE;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_CHECKSUM_TYPE;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_COMPACTION_STYLE;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_COMPRESSION_TYPE;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_LOG_LEVEL;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.DEFAULT_PARALLELISM;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.MAX_WRITE_BUFFERS;
	import static org.apache.bookkeeper.statelib.impl.rocksdb.RocksConstants.WRITE_BUFFER_SIZE;

	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Sets;
	import com.google.common.io.MoreFiles;
	import com.google.common.io.RecursiveDeleteOption;
	import com.google.common.primitives.SignedBytes;
	import java.io.File;
	import java.io.IOException;
	import java.nio.file.Files;
	import java.nio.file.Path;
	import java.nio.file.Paths;
	import java.util.Collections;
	import java.util.Comparator;
	import java.util.List;
	import java.util.NoSuchElementException;
	import java.util.Set;
	import java.util.UUID;
	import java.util.concurrent.ScheduledExecutorService;
	import java.util.concurrent.atomic.AtomicLongFieldUpdater;
	import lombok.AccessLevel;
	import lombok.RequiredArgsConstructor;
	import lombok.extern.slf4j.Slf4j;
	import org.apache.bookkeeper.common.coder.Coder;
	import org.apache.bookkeeper.common.kv.KV;
	import org.apache.bookkeeper.common.kv.KVImpl;
	import org.apache.bookkeeper.statelib.api.StateStoreSpec;
	import org.apache.bookkeeper.statelib.api.checkpoint.CheckpointStore;
	import org.apache.bookkeeper.statelib.api.exceptions.InvalidStateStoreException;
	import org.apache.bookkeeper.statelib.api.exceptions.StateStoreException;
	import org.apache.bookkeeper.statelib.api.exceptions.StateStoreRuntimeException;
	import org.apache.bookkeeper.statelib.api.kv.KVIterator;
	import org.apache.bookkeeper.statelib.api.kv.KVMulti;
	import org.apache.bookkeeper.statelib.api.kv.KVStore;
	import org.apache.bookkeeper.statelib.impl.Bytes;
	import org.apache.bookkeeper.statelib.impl.rocksdb.RocksUtils;
	import org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint.CheckpointInfo;
	import org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint.RocksCheckpointer;
	import org.apache.commons.lang3.tuple.Pair;
	import org.rocksdb.BlockBasedTableConfig;
	import org.rocksdb.Cache;
	import org.rocksdb.ColumnFamilyDescriptor;
	import org.rocksdb.ColumnFamilyHandle;
	import org.rocksdb.ColumnFamilyOptions;
	import org.rocksdb.DBOptions;
	import org.rocksdb.FlushOptions;
	import org.rocksdb.LRUCache;
	import org.rocksdb.RocksDB;
	import org.rocksdb.RocksDBException;
	import org.rocksdb.RocksIterator;
	import org.rocksdb.WriteBatch;
	import org.rocksdb.WriteOptions;

	/**
	* A key/value store implemented using {@link http://rocksdb.org/}.
	*
	* @param <K> key type
	* @param <V> value type
	*/
	@Slf4j
	public class RocksdbKVStore<K, V> implements KVStore<K, V> {

	private static final byte[] METADATA_CF = ".meta".getBytes(UTF_8);
	private static final byte[] DATA_CF = "default".getBytes(UTF_8);
	private static final byte[] LAST_REVISION = ".lrev".getBytes(UTF_8);

	private static final AtomicLongFieldUpdater<RocksdbKVStore> lastRevisionUpdater =
	AtomicLongFieldUpdater.newUpdater(RocksdbKVStore.class, "lastRevision");

	// parameters for the store
	protected String name;
	protected Coder<K> keyCoder;
	protected Coder<V> valCoder;

	// rocksdb state
	protected File dbDir;
	protected RocksDB db;
	protected ColumnFamilyHandle metaCfHandle;
	protected ColumnFamilyHandle dataCfHandle;

	// iterators
	protected final Set<KVIterator<K, V>> kvIters;

	// options used by rocksdb
	protected DBOptions dbOpts;
	protected ColumnFamilyOptions cfOpts;
	protected WriteOptions writeOpts;
	protected FlushOptions flushOpts;

	// states of the store
	protected volatile boolean isInitialized = false;
	protected volatile boolean closed = false;
	protected volatile long lastRevision = -1L;
	private final byte[] lastRevisionBytes = new byte[Long.BYTES];

	// checkpointer store
	private CheckpointStore checkpointStore;
	private ScheduledExecutorService checkpointScheduler;
	// rocksdb checkpointer
	private RocksCheckpointer checkpointer;

	static {
	RocksDB.loadLibrary();
	}

	private boolean cleanupLocalStoreDirEnable;

	public RocksdbKVStore() {
	// initialize the iterators set
	this.kvIters = Collections.synchronizedSet(Sets.newHashSet());
	}

	protected void checkStoreOpen() {
	if (closed) {
	throw new InvalidStateStoreException("State store " + name + " is already closed");
	}
	if (!isInitialized) {
	throw new InvalidStateStoreException("State Store " + name + " is not initialized yet");
	}
	}

	@VisibleForTesting
	public synchronized RocksDB getDb() {
	return db;
	}

	@Override
	public synchronized String name() {
	return this.name;
	}

	private void loadRocksdbFromCheckpointStore(StateStoreSpec spec) {
	checkNotNull(spec.getCheckpointIOScheduler(),
	"checkpoint io scheduler is not configured");
	checkNotNull(spec.getCheckpointDuration(),
	"checkpoint duration is not configured");

	String dbName = spec.getName();
	File localStorePath = spec.getLocalStateStoreDir();

	List<CheckpointInfo> checkpoints = RocksCheckpointer.getCheckpoints(dbName, spec.getCheckpointStore());
	for (CheckpointInfo cpi : checkpoints) {
	try {
	cpi.restore(dbName, localStorePath, spec.getCheckpointStore());
	openRocksdb(spec);
	checkpoints.stream()
	.filter(cp -> cp != cpi) // ignore the current restored checkpoint
	.forEach(cp -> cp.remove(localStorePath)); // delete everything else
	break;
	} catch (StateStoreException e) {
	// Got an exception. Log and try the next checkpoint
	log.error("Failed to restore checkpoint: {}", cpi, e);
	}
	}
	}

	@Override
	public synchronized void checkpoint() {
	log.info("Checkpoint local state store {} at revision {}", name, getLastRevision());
	byte[] checkpointAtRevisionBytes = new byte[Long.BYTES];
	System.arraycopy(lastRevisionBytes, 0, checkpointAtRevisionBytes, 0, checkpointAtRevisionBytes.length);
	checkpointScheduler.submit(() -> {
	try {
	// TODO: move create checkpoint to the checkpoint method
	checkpointer.checkpointAtTxid(checkpointAtRevisionBytes);
	} catch (StateStoreException e) {
	log.error("Failed to checkpoint state store {} at revision {}",
	name, Bytes.toLong(checkpointAtRevisionBytes, 0), e);
	}
	});
	}

	private void readLastRevision() throws StateStoreException {
	byte[] revisionBytes;
	try {
	revisionBytes = db.get(metaCfHandle, LAST_REVISION);
	} catch (RocksDBException e) {
	throw new StateStoreException("Failed to read last revision from state store " + name(), e);
	}
	if (null == revisionBytes) {
	return;
	}
	long revision = Bytes.toLong(revisionBytes, 0);
	lastRevisionUpdater.set(this, revision);
	}

	@Override
	public long getLastRevision() {
	return lastRevisionUpdater.get(this);
	}

	private void setLastRevision(long lastRevision) {
	lastRevisionUpdater.set(this, lastRevision);
	Bytes.toBytes(lastRevision, lastRevisionBytes, 0);
	}

	private void updateLastRevision(long revision) {
	if (revision >= 0) { // k/v comes from log stream
	if (getLastRevision() >= revision) { // these k/v pairs are duplicates
	return;
	}
	// update revision
	setLastRevision(revision);
	}
	}

	protected void updateLastRevision(WriteBatch batch, long revision) {
	if (revision >= 0) { // k/v comes from log stream
	if (getLastRevision() >= revision) { // these k/v pairs are duplicates
	return;
	}
	try {
	// update revision
	setLastRevision(revision);
	batch.put(metaCfHandle, LAST_REVISION, lastRevisionBytes);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException(
	"Error while updating last revision " + revision + " from store " + name, e);
	}
	}
	}

	@Override
	@SuppressWarnings("unchecked")
	public synchronized void init(StateStoreSpec spec) throws StateStoreException {
	checkNotNull(spec.getLocalStateStoreDir(),
	"local state store directory is not configured");

	this.name = spec.getName();
	this.cleanupLocalStoreDirEnable = spec.isLocalStorageCleanupEnable();

	// initialize the coders
	this.keyCoder = (Coder<K>) spec.getKeyCoder();
	this.valCoder = (Coder<V>) spec.getValCoder();

	cleanupLocalStoreDir(spec.getLocalStateStoreDir());

	checkpointStore = spec.getCheckpointStore();
	if (null != checkpointStore) {
	// load checkpoint from checkpoint store
	loadRocksdbFromCheckpointStore(spec);
	} else {
	// open the rocksdb
	openRocksdb(spec);
	}


	// once the rocksdb is opened, read the last revision
	readLastRevision();

	if (null != checkpointStore) {
	checkpointer = new RocksCheckpointer(
	name(),
	dbDir,
	db,
	checkpointStore,
	true,
	true,
	spec.isCheckpointChecksumEnable(),
	spec.isCheckpointChecksumCompatible());
	checkpointScheduler = spec.getCheckpointIOScheduler();
	}

	this.isInitialized = true;
	}

	protected void openRocksdb(StateStoreSpec spec) throws StateStoreException {

	// initialize the db options

	final BlockBasedTableConfig tableConfig = new BlockBasedTableConfig();
	final Cache cache = new LRUCache(BLOCK_CACHE_SIZE);
	tableConfig.setBlockCache(cache);
	tableConfig.setBlockSize(BLOCK_SIZE);
	tableConfig.setChecksumType(DEFAULT_CHECKSUM_TYPE);

	dbOpts = new DBOptions();
	dbOpts.setCreateIfMissing(true);
	dbOpts.setErrorIfExists(false);
	dbOpts.setInfoLogLevel(DEFAULT_LOG_LEVEL);
	dbOpts.setIncreaseParallelism(DEFAULT_PARALLELISM);
	dbOpts.setCreateMissingColumnFamilies(true);

	cfOpts = new ColumnFamilyOptions();
	cfOpts.setTableFormatConfig(tableConfig);
	cfOpts.setWriteBufferSize(WRITE_BUFFER_SIZE);
	cfOpts.setCompressionType(DEFAULT_COMPRESSION_TYPE);
	cfOpts.setCompactionStyle(DEFAULT_COMPACTION_STYLE);
	cfOpts.setMaxWriteBufferNumber(MAX_WRITE_BUFFERS);

	// initialize the write options

	writeOpts = new WriteOptions();
	writeOpts.setDisableWAL(true); // disable wal, since the source of truth will be on distributedlog

	// initialize the flush options

	flushOpts = new FlushOptions();
	flushOpts.setWaitForFlush(true);

	// open the rocksdb

	this.dbDir = spec.getLocalStateStoreDir();
	Pair<RocksDB, List<ColumnFamilyHandle>> dbPair = openLocalDB(dbDir, dbOpts, cfOpts);
	this.db = dbPair.getLeft();
	this.metaCfHandle = dbPair.getRight().get(0);
	this.dataCfHandle = dbPair.getRight().get(1);
	}

	protected Pair<RocksDB, List<ColumnFamilyHandle>> openLocalDB(File dir,
	DBOptions options,
	ColumnFamilyOptions cfOpts)
	throws StateStoreException {
	return openRocksdb(dir, options, cfOpts);
	}

	protected static Pair<RocksDB, List<ColumnFamilyHandle>> openRocksdb(
	File dir, DBOptions options, ColumnFamilyOptions cfOpts)
	throws StateStoreException {
	// make sure the db directory's parent dir is created
	ColumnFamilyDescriptor metaDesc = new ColumnFamilyDescriptor(METADATA_CF, cfOpts);
	ColumnFamilyDescriptor dataDesc = new ColumnFamilyDescriptor(DATA_CF, cfOpts);

	try {
	Files.createDirectories(dir.toPath());
	File dbDir = new File(dir, "current");

	if (!dbDir.exists()) {
	// empty state
	String uuid = UUID.randomUUID().toString();
	Path checkpointPath = Paths.get(dir.getAbsolutePath(), "checkpoints", uuid);
	Files.createDirectories(checkpointPath);
	Files.createSymbolicLink(
	Paths.get(dbDir.getAbsolutePath()),
	checkpointPath);
	}

	List<ColumnFamilyHandle> cfHandles = Lists.newArrayListWithExpectedSize(2);
	RocksDB db = RocksDB.open(
	options,
	dbDir.getAbsolutePath(),
	Lists.newArrayList(metaDesc, dataDesc),
	cfHandles);
	return Pair.of(db, cfHandles);
	} catch (IOException ioe) {
	log.error("Failed to create parent directory {} for opening rocksdb", dir.getParentFile().toPath(), ioe);
	throw new StateStoreException(ioe);
	} catch (RocksDBException dbe) {
	log.error("Failed to open rocksdb at dir {}", dir.getAbsolutePath(), dbe);
	throw new StateStoreException(dbe);
	}
	}

	@Override
	public synchronized void flush() throws StateStoreException {
	if (null == db) {
	return;
	}
	try {
	db.flush(flushOpts);
	} catch (RocksDBException e) {
	throw new StateStoreException("Exception on flushing rocksdb from store " + name, e);
	}
	}

	@Override
	public synchronized void close() {
	if (closed) {
	return;
	}
	closed = true;

	if (null != checkpointer) {
	checkpointer.close();
	}

	// close iterators
	closeIters();

	// close db
	closeLocalDB();

	// release options
	RocksUtils.close(dbOpts);
	RocksUtils.close(writeOpts);
	RocksUtils.close(flushOpts);
	RocksUtils.close(cfOpts);

	cleanupLocalStoreDir(dbDir);
	}

	private void cleanupLocalStoreDir(File dbDir) {
	if (cleanupLocalStoreDirEnable) {
	if (dbDir.exists()) {
	try {
	MoreFiles.deleteRecursively(dbDir.toPath(), RecursiveDeleteOption.ALLOW_INSECURE);
	} catch (IOException e) {
	log.error("Failed to cleanup localStoreDir", e);
	}
	}
	}
	}

	protected void closeLocalDB() {
	try {
	flush();
	} catch (StateStoreException e) {
	// flush() already logs this exception.
	}

	RocksUtils.close(metaCfHandle);
	RocksUtils.close(dataCfHandle);
	RocksUtils.close(db);
	}

	private void closeIters() {
	Set<KVIterator> iterators;
	synchronized (kvIters) {
	iterators = Sets.newHashSet(kvIters);
	}
	iterators.forEach(KVIterator::close);
	}

	@Override
	public synchronized V get(K key) {
	checkNotNull(key, "key cannot be null");
	checkStoreOpen();

	byte[] keyBytes = keyCoder.encode(key);
	return getRaw(key, keyBytes);
	}

	private V getRaw(K key, byte[] keyBytes) {
	byte[] valBytes = getRawBytes(key, keyBytes);
	if (null == valBytes) {
	return null;
	}
	return valCoder.decode(valBytes);
	}

	protected byte[] getRawBytes(K key, byte[] keyBytes) {
	try {
	return this.db.get(dataCfHandle, keyBytes);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException("Error while getting value for key " + key + " from store " + name, e);
	}
	}

	@Override
	public synchronized KVIterator<K, V> range(K from, K to) {
	checkStoreOpen();

	RocksIterator rocksIter = db.newIterator(dataCfHandle);
	if (null == from) {
	rocksIter.seekToFirst();
	} else {
	byte[] fromBytes = keyCoder.encode(from);
	rocksIter.seek(fromBytes);
	}
	KVIterator<K, V> kvIter;
	if (null == to) {
	kvIter = new RocksdbKVIterator(name, rocksIter, keyCoder, valCoder);
	} else {
	kvIter = new RocksdbRangeIterator(name, rocksIter, keyCoder, valCoder, to);
	}
	kvIters.add(kvIter);
	return kvIter;
	}

	@Override
	public synchronized void put(K key, V value) {
	put(key, value, -1);
	}

	synchronized void put(K key, V value, long revision) {
	checkNotNull(key, "key cannot be null");
	checkStoreOpen();

	updateLastRevision(revision);

	byte[] keyBytes = keyCoder.encode(key);
	putRaw(key, keyBytes, value, revision);
	}

	private void putRaw(K key, byte[] keyBytes, V value, long revision) {
	try {
	WriteBatch batch = new WriteBatch();
	if (revision > 0) {
	// last revision has been set to revision bytes
	batch.put(metaCfHandle, LAST_REVISION, lastRevisionBytes);
	}
	if (null == value) {
	// delete a key if value is null
	batch.delete(dataCfHandle, keyBytes);
	} else {
	byte[] valBytes = valCoder.encode(value);
	batch.put(dataCfHandle, keyBytes, valBytes);
	}

	db.write(writeOpts, batch);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException("Error while updating key " + key
	+ " to value " + value + " from store " + name, e);
	}
	}

	@Override
	public V putIfAbsent(K key, V value) {
	return putIfAbsent(key, value, -1L);
	}

	synchronized V putIfAbsent(K key, V value, long revision) {
	checkNotNull(key, "key cannot be null");
	checkStoreOpen();

	updateLastRevision(revision);

	byte[] keyBytes = keyCoder.encode(key);
	V oldVal = getRaw(key, keyBytes);
	if (null != oldVal) {
	return oldVal;
	}

	if (value == null) {
	return null;
	}

	putRaw(key, keyBytes, value, revision);
	return null;
	}

	@Override
	public synchronized KVMulti<K, V> multi() {
	checkStoreOpen();

	return new KVMultiImpl();
	}

	@Override
	public synchronized V delete(K key) {
	return delete(key, -1L);
	}

	synchronized V delete(K key, long revision) {
	checkNotNull(key, "key cannot be null");
	checkStoreOpen();

	updateLastRevision(revision);

	byte[] keyBytes = keyCoder.encode(key);
	V val = getRaw(key, keyBytes);
	putRaw(key, keyBytes, null, revision);

	return val;
	}

	//
	// Multi
	//

	/**
	* A rocksdb based multi operation.
	*/
	class KVMultiImpl implements KVMulti<K, V> {

	private final WriteBatch batch = new WriteBatch();
	private volatile boolean executed = false;

	private void checkExecuted() {
	if (executed) {
	throw new StateStoreRuntimeException("KVMulti#execute() has been called");
	}
	}

	@Override
	public void put(K key, V value) {
	checkNotNull(key, "key cannot be null");
	checkExecuted();

	byte[] keyBytes = keyCoder.encode(key);
	if (null == value) {
	deleteRaw(keyBytes);
	} else {
	putRaw(keyBytes, value);
	}
	}

	private void putRaw(byte[] keyBytes, V value) {
	try {
	batch.put(dataCfHandle, keyBytes, valCoder.encode(value));
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException(e);
	}
	}

	@Override
	public void delete(K key) {
	checkNotNull(key, "key cannot be null");
	checkExecuted();

	byte[] keyBytes = keyCoder.encode(key);
	deleteRaw(keyBytes);
	}

	private void deleteRaw(byte[] keyBytes) {
	try {
	batch.delete(dataCfHandle, keyBytes);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException(e);
	}
	}

	@Override
	public void deleteRange(K from, K to) {
	checkNotNull(from, "from key cannot be null");
	checkNotNull(to, "to key cannot be null");
	checkExecuted();

	byte[] fromBytes = keyCoder.encode(from);
	byte[] toBytes = keyCoder.encode(to);
	try {
	batch.deleteRange(dataCfHandle, fromBytes, toBytes);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException(e);
	}
	}

	@Override
	public synchronized void execute() {
	if (executed) {
	return;
	}

	checkStoreOpen();
	executed = true;

	try {
	getDb().write(writeOpts, batch);
	} catch (RocksDBException e) {
	throw new StateStoreRuntimeException("Error while executing a multi operation from store " + name, e);
	} finally {
	RocksUtils.close(batch);
	}
	}
	}

	//
	// Iterators
	//

	/**
	* KV iterator over a rocksdb instance.
	*/
	@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
	class RocksdbKVIterator implements KVIterator<K, V> {

	final String name;
	final RocksIterator iterator;
	final Coder<K> keyCoder;
	final Coder<V> valCoder;

	private volatile boolean closed = false;

	/**
	* Ensure an iterator is open.
	*
	* @throws InvalidStateStoreException when the store is in closed state.
	*/
	private void ensureIteratorOpen() {
	if (closed) {
	throw new InvalidStateStoreException("Rocksdb state store " + name + " is already closed");
	}
	}

	@Override
	public void close() {
	kvIters.remove(this);
	iterator.close();
	closed = true;
	}

	@Override
	public boolean hasNext() {
	ensureIteratorOpen();
	return iterator.isValid();
	}

	private KV<K, V> getKvPair() {
	return new KVImpl<>(keyCoder.decode(iterator.key()), valCoder.decode(iterator.value()));
	}

	@Override
	public KV<K, V> next() {
	if (!hasNext()) {
	throw new NoSuchElementException();
	}

	KV<K, V> kv = getKvPair();
	iterator.next();
	return kv;
	}
	}

	/**
	* KV iterator over a rocksdb, which ends at the provided <i>endKey</i>.
	*/
	class RocksdbRangeIterator extends RocksdbKVIterator {

	private final Comparator<byte[]> comparator = SignedBytes.lexicographicalComparator();
	private final byte[] endKeyBytes;

	private RocksdbRangeIterator(String name,
	RocksIterator iterator,
	Coder<K> keyCoder,
	Coder<V> valCoder,
	K endKey) {
	super(name, iterator, keyCoder, valCoder);
	checkNotNull(endKey, "End key cannot be null");
	this.endKeyBytes = keyCoder.encode(endKey);
	}


	@Override
	public boolean hasNext() {
	return super.hasNext()
	&& comparator.compare(iterator.key(), endKeyBytes) <= 0;
	}
	}


	}