Add support for snapshots.

rocksdb/options/readopts.nim

@@ -9,7 +9,9 @@
 
 {.push raises: [].}
 
-import ../lib/librocksdb
+import ../lib/librocksdb, ../snapshot
+
+export snapshot.SnapshotRef, snapshot.isClosed, snapshot.getSequenceNumber
 
 type
   ReadOptionsPtr* = ptr rocksdb_readoptions_t
@@ -51,6 +53,10 @@ opt ignoreRangeDeletions, bool, uint8
 opt deadline, int, uint64
 opt ioTimeout, int, uint64
 
+proc setSnapshot*(readOpts: ReadOptionsRef, snapshot: SnapshotRef) =
+  doAssert not readOpts.isClosed()
+  rocksdb_readoptions_set_snapshot(readOpts.cPtr, snapshot.cPtr)
+
 proc defaultReadOptions*(autoClose = false): ReadOptionsRef {.inline.} =
   let readOpts = createReadOptions(autoClose)
 

rocksdb/rocksdb.nim

@@ -31,11 +31,12 @@ import
   ./options/[dbopts, readopts, writeopts],
   ./columnfamily/[cfopts, cfdescriptor, cfhandle],
   ./internal/[cftable, utils],
-  ./[rocksiterator, rocksresult, writebatch, writebatchwi]
+  ./[rocksiterator, rocksresult, writebatch, writebatchwi, snapshot]
 
 export
   rocksresult, dbopts, readopts, writeopts, cfdescriptor, cfhandle, rocksiterator,
-  writebatch, writebatchwi
+  writebatch, writebatchwi, snapshot.SnapshotRef, snapshot.isClosed,
+  snapshot.getSequenceNumber
 
 type
   RocksDbPtr* = ptr rocksdb_t
@@ -355,16 +356,17 @@ proc delete*(
   ok()
 
 proc openIterator*(
-    db: RocksDbRef, cfHandle = db.defaultCfHandle
+    db: RocksDbRef,
+    readOpts = defaultReadOptions(autoClose = true),
+    cfHandle = db.defaultCfHandle,
 ): RocksDBResult[RocksIteratorRef] =
   ## Opens an `RocksIteratorRef` for the specified column family.
   ## The iterator should be closed using the `close` method after usage.
   doAssert not db.isClosed()
 
-  let rocksIterPtr =
-    rocksdb_create_iterator_cf(db.cPtr, db.readOpts.cPtr, cfHandle.cPtr)
+  let rocksIterPtr = rocksdb_create_iterator_cf(db.cPtr, readOpts.cPtr, cfHandle.cPtr)
 
-  ok(newRocksIterator(rocksIterPtr))
+  ok(newRocksIterator(rocksIterPtr, readOpts))
 
 proc openWriteBatch*(
     db: RocksDbReadWriteRef, cfHandle = db.defaultCfHandle
@@ -442,6 +444,28 @@ proc ingestExternalFile*(
 
   ok()
 
+proc getSnapshot*(db: RocksDbRef): RocksDBResult[SnapshotRef] =
+  ## Return a handle to the current DB state. Iterators created with this handle
+  ## will all observe a stable snapshot of the current DB state. The caller must
+  ## call ReleaseSnapshot(result) when the snapshot is no longer needed.
+  doAssert not db.isClosed()
+
+  let sHandle = rocksdb_create_snapshot(db.cPtr)
+  if sHandle.isNil():
+    err("rocksdb: failed to create snapshot")
+  else:
+    ok(newSnapshot(sHandle, SnapshotType.rocksdb))
+
+proc releaseSnapshot*(db: RocksDbRef, snapshot: SnapshotRef) =
+  ## Release a previously acquired snapshot. The caller must not use "snapshot"
+  ## after this call.
+  doAssert not db.isClosed()
+  doAssert snapshot.kind == SnapshotType.rocksdb
+
+  if not snapshot.isClosed():
+    rocksdb_release_snapshot(db.cPtr, snapshot.cPtr)
+    snapshot.setClosed()
+
 proc close*(db: RocksDbRef) =
   ## Close the `RocksDbRef` which will release the connection to the database
   ## and free the memory associated with it. `close` is idempotent and can

rocksdb/rocksiterator.nim

@@ -12,7 +12,7 @@
 
 {.push raises: [].}
 
-import ./lib/librocksdb, ./internal/utils, ./rocksresult
+import ./lib/librocksdb, ./internal/utils, ./options/readopts, ./rocksresult
 
 export rocksresult
 
@@ -21,10 +21,13 @@ type
 
   RocksIteratorRef* = ref object
     cPtr: RocksIteratorPtr
+    readOpts: ReadOptionsRef
 
-proc newRocksIterator*(cPtr: RocksIteratorPtr): RocksIteratorRef =
+proc newRocksIterator*(
+    cPtr: RocksIteratorPtr, readOpts: ReadOptionsRef
+): RocksIteratorRef =
   doAssert not cPtr.isNil()
-  RocksIteratorRef(cPtr: cPtr)
+  RocksIteratorRef(cPtr: cPtr, readOpts: readOpts)
 
 proc isClosed*(iter: RocksIteratorRef): bool {.inline.} =
   ## Returns `true` if the iterator is closed and `false` otherwise.
@@ -128,6 +131,8 @@ proc close*(iter: RocksIteratorRef) =
     rocksdb_iter_destroy(iter.cPtr)
     iter.cPtr = nil
 
+    autoCloseNonNil(iter.readOpts)
+
 iterator pairs*(iter: RocksIteratorRef): tuple[key: seq[byte], value: seq[byte]] =
   ## Iterates over the key value pairs in the column family yielding them in
   ## the form of a tuple. The iterator is automatically closed after the

rocksdb/snapshot.nim

@@ -0,0 +1,51 @@
+# Nim-RocksDB
+# Copyright 2024 Status Research & Development GmbH
+# Licensed under either of
+#
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0)
+#  * GPL license, version 2.0, ([LICENSE-GPLv2](LICENSE-GPLv2) or https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)
+#
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
+## A `SnapshotRef` represents a view of the state of the database at some point in time.
+
+{.push raises: [].}
+
+import ./lib/librocksdb
+
+type
+  SnapshotPtr* = ptr rocksdb_snapshot_t
+
+  SnapshotType* = enum
+    rocksdb
+    transactiondb
+
+  SnapshotRef* = ref object
+    cPtr: SnapshotPtr
+    kind: SnapshotType
+
+proc newSnapshot*(cPtr: SnapshotPtr, kind: SnapshotType): SnapshotRef =
+  doAssert not cPtr.isNil()
+  SnapshotRef(cPtr: cPtr, kind: kind)
+
+proc isClosed*(snapshot: SnapshotRef): bool {.inline.} =
+  ## Returns `true` if the `SnapshotRef` has been closed and `false` otherwise.
+  snapshot.cPtr.isNil()
+
+proc cPtr*(snapshot: SnapshotRef): SnapshotPtr =
+  ## Get the underlying database pointer.
+  doAssert not snapshot.isClosed()
+  snapshot.cPtr
+
+proc kind*(snapshot: SnapshotRef): SnapshotType =
+  ## Get the kind of the `SnapshotRef`.
+  snapshot.kind
+
+proc getSequenceNumber*(snapshot: SnapshotRef): uint64 =
+  ## Return the associated sequence number.
+  doAssert not snapshot.isClosed()
+  rocksdb_snapshot_get_sequence_number(snapshot.cPtr).uint64
+
+proc setClosed*(snapshot: SnapshotRef) =
+  # The snapshot should be released from `RocksDbRef` or `TransactionDbRef`
+  snapshot.cPtr = nil

rocksdb/transactiondb.nim

@@ -22,10 +22,12 @@ import
   ./transactions/[transaction, txdbopts, txopts],
   ./columnfamily/[cfopts, cfdescriptor, cfhandle],
   ./internal/[cftable, utils],
-  ./rocksresult
+  ./[rocksiterator, rocksresult, snapshot]
 
 export
-  dbopts, txdbopts, cfdescriptor, readopts, writeopts, txopts, transaction, rocksresult
+  dbopts, txdbopts, cfdescriptor, readopts, writeopts, txopts, transaction,
+  rocksiterator, rocksresult, snapshot.SnapshotRef, snapshot.isClosed,
+  snapshot.getSequenceNumber
 
 type
   TransactionDbPtr* = ptr rocksdb_transactiondb_t
@@ -120,6 +122,42 @@ proc beginTransaction*(
 
   newTransaction(txPtr, readOpts, writeOpts, txOpts, nil, cfHandle)
 
+proc openIterator*(
+    db: TransactionDbRef,
+    readOpts = defaultReadOptions(autoClose = true),
+    cfHandle = db.defaultCfHandle,
+): RocksDBResult[RocksIteratorRef] =
+  ## Opens an `RocksIteratorRef` for the specified column family.
+  ## The iterator should be closed using the `close` method after usage.
+  doAssert not db.isClosed()
+
+  let rocksIterPtr =
+    rocksdb_transactiondb_create_iterator_cf(db.cPtr, readOpts.cPtr, cfHandle.cPtr)
+
+  ok(newRocksIterator(rocksIterPtr, readOpts))
+
+proc getSnapshot*(db: TransactionDbRef): RocksDBResult[SnapshotRef] =
+  ## Return a handle to the current DB state. Iterators created with this handle
+  ## will all observe a stable snapshot of the current DB state. The caller must
+  ## call ReleaseSnapshot(result) when the snapshot is no longer needed.
+  doAssert not db.isClosed()
+
+  let sHandle = rocksdb_transactiondb_create_snapshot(db.cPtr)
+  if sHandle.isNil():
+    err("rocksdb: failed to create snapshot")
+  else:
+    ok(newSnapshot(sHandle, SnapshotType.transactiondb))
+
+proc releaseSnapshot*(db: TransactionDbRef, snapshot: SnapshotRef) =
+  ## Release a previously acquired snapshot. The caller must not use "snapshot"
+  ## after this call.
+  doAssert not db.isClosed()
+  doAssert snapshot.kind == SnapshotType.transactiondb
+
+  if not snapshot.isClosed():
+    rocksdb_transactiondb_release_snapshot(db.cPtr, snapshot.cPtr)
+    snapshot.setClosed()
+
 proc close*(db: TransactionDbRef) =
   ## Close the `TransactionDbRef`.
 

rocksdb/writebatchwi.nim

@@ -46,7 +46,7 @@ proc isClosed*(batch: WriteBatchWIRef): bool {.inline.} =
   batch.cPtr.isNil()
 
 proc cPtr*(batch: WriteBatchWIRef): WriteBatchWIPtr =
-  ## Get the underlying database pointer.
+  ## Get the underlying write batch pointer.
   doAssert not batch.isClosed()
   batch.cPtr
 

tests/test_rocksdb.nim

@@ -460,3 +460,38 @@ suite "RocksDbRef Tests":
 
     cfOpts.close()
     removeDir($dbPath)
+
+  test "Create and restore snapshot":
+    check:
+      db.put(key, val).isOk()
+      db.keyExists(key).get() == true
+      db.keyMayExist(otherKey).get() == false
+
+    let snapshot = db.getSnapshot().get()
+    check:
+      snapshot.getSequenceNumber() > 0
+      not snapshot.isClosed()
+
+    check:
+      db.delete(key).isOk()
+      db.put(otherKey, val).isOk()
+      db.keyMayExist(key).get() == false
+      db.keyExists(otherKey).get() == true
+
+    let readOpts = defaultReadOptions(autoClose = true)
+    readOpts.setSnapshot(snapshot)
+    let iter = db.openIterator(readOpts = readOpts).get()
+    defer:
+      iter.close()
+
+    iter.seekToKey(key)
+    check:
+      iter.isValid() == true
+      iter.key() == key
+      iter.value() == val
+    iter.seekToKey(otherKey)
+    check:
+      iter.isValid() == false
+
+    db.releaseSnapshot(snapshot)
+    check snapshot.isClosed()

tests/test_rocksiterator.nim

@@ -46,7 +46,7 @@ suite "RocksIteratorRef Tests":
     removeDir($dbPath)
 
   test "Iterate forwards using default column family":
-    let res = db.openIterator(defaultCfHandle)
+    let res = db.openIterator(cfHandle = defaultCfHandle)
     check res.isOk()
 
     var iter = res.get()
@@ -72,7 +72,7 @@ suite "RocksIteratorRef Tests":
     check expected == byte(4)
 
   test "Iterate backwards using other column family":
-    let res = db.openIterator(otherCfHandle)
+    let res = db.openIterator(cfHandle = otherCfHandle)
     check res.isOk()
 
     var iter = res.get()
@@ -106,12 +106,12 @@ suite "RocksIteratorRef Tests":
     iter.close()
 
   test "Open two iterators on the same column family":
-    let res1 = db.openIterator(defaultCfHandle)
+    let res1 = db.openIterator(cfHandle = defaultCfHandle)
     check res1.isOk()
     var iter1 = res1.get()
     defer:
       iter1.close()
-    let res2 = db.openIterator(defaultCfHandle)
+    let res2 = db.openIterator(cfHandle = defaultCfHandle)
     check res2.isOk()
     var iter2 = res2.get()
     defer:
@@ -129,12 +129,12 @@ suite "RocksIteratorRef Tests":
       iter2.value() == @[byte(3)]
 
   test "Open two iterators on different column families":
-    let res1 = db.openIterator(defaultCfHandle)
+    let res1 = db.openIterator(cfHandle = defaultCfHandle)
     check res1.isOk()
     var iter1 = res1.get()
     defer:
       iter1.close()
-    let res2 = db.openIterator(otherCfHandle)
+    let res2 = db.openIterator(cfHandle = otherCfHandle)
     check res2.isOk()
     var iter2 = res2.get()
     defer:
@@ -152,7 +152,7 @@ suite "RocksIteratorRef Tests":
       iter2.value() == @[byte(3)]
 
   test "Iterate forwards using seek to key":
-    let res = db.openIterator(defaultCfHandle)
+    let res = db.openIterator(cfHandle = defaultCfHandle)
     check res.isOk()
 
     var iter = res.get()
@@ -180,7 +180,7 @@ suite "RocksIteratorRef Tests":
       iter.value() == val1
 
   test "Empty column family":
-    let res = db.openIterator(emptyCfHandle)
+    let res = db.openIterator(cfHandle = emptyCfHandle)
     check res.isOk()
     var iter = res.get()
     defer:
@@ -193,7 +193,7 @@ suite "RocksIteratorRef Tests":
     check not iter.isValid()
 
   test "Test status":
-    let res = db.openIterator(emptyCfHandle)
+    let res = db.openIterator(cfHandle = emptyCfHandle)
     check res.isOk()
     var iter = res.get()
     defer:
@@ -204,7 +204,7 @@ suite "RocksIteratorRef Tests":
     check iter.status().isOk()
 
   test "Test pairs iterator":
-    let res = db.openIterator(defaultCfHandle)
+    let res = db.openIterator(cfHandle = defaultCfHandle)
     check res.isOk()
     var iter = res.get()