lsm: add support for read-your-own-writes of uncommitted data

Now that we have unified the memtable and write batch to a single
structure, we can provide iteration and `get` semantics over the
pending write batch to support use cases where you want to explicitly
flush pending operations or provide some kind of transactional support
with read-your-own-writes guarentees.

For prior art, see the following in RocksDB:
https://github.com/facebook/rocksdb/wiki/Write-Batch-With-Index#write-batch-with-index

Author: rockwotj

src/v/lsm/db/impl.cc

@@ -186,13 +186,22 @@ ss::future<lookup_result> impl::get(internal::key_view key) {
     co_return result;
 }
 
-ss::future<std::unique_ptr<internal::iterator>> impl::create_iterator() {
-    auto iter = co_await create_internal_iterator();
+ss::future<std::unique_ptr<internal::iterator>>
+impl::create_iterator(ss::optimized_optional<ss::lw_shared_ptr<memtable>> wb) {
+    auto iter = co_await create_internal_iterator(wb->get());
+    internal::sequence_number iter_seqno = max_applied_seqno();
+    if (wb) {
+        auto wb_seqno = (*wb)->last_seqno().value_or(iter_seqno);
+        vassert(
+          wb_seqno > iter_seqno,
+          "write memtable seqno must be greater than the current "
+          "max_applied_seqno: {} > {}",
+          wb_seqno,
+          iter_seqno);
+        iter_seqno = wb_seqno;
+    }
     co_return create_db_iterator(
-      std::move(iter),
-      max_applied_seqno(),
-      _opts,
-      [this](internal::key_view key) {
+      std::move(iter), iter_seqno, _opts, [this](internal::key_view key) {
           return _versions->current()->record_read_sample(key).then(
             [this](bool compaction_needed) {
                 if (compaction_needed) {
@@ -203,8 +212,11 @@ ss::future<std::unique_ptr<internal::iterator>> impl::create_iterator() {
 }
 
 ss::future<std::unique_ptr<internal::iterator>>
-impl::create_internal_iterator() {
+impl::create_internal_iterator(ss::optimized_optional<memtable*> wb) {
     chunked_vector<std::unique_ptr<internal::iterator>> list;
+    if (wb) {
+        list.push_back((*wb)->create_iterator());
+    }
     list.push_back(_mem->create_iterator());
     if (_imm) {
         list.push_back((*_imm)->create_iterator());

src/v/lsm/db/impl.h

@@ -57,10 +57,12 @@ class impl {
     // Get a key from the database
     ss::future<lookup_result> get(internal::key_view);
 
-    // Create an iterator over the database. Note that this iterator
-    // results in ALL entries from the database, a deduplicating iterator
-    // needs to be added on top to give a traditional iterator view.
-    ss::future<std::unique_ptr<internal::iterator>> create_iterator();
+    // Create an interator over the database.
+    //
+    // If a non-null memtable is passed in, then a frozen state of the memtable
+    // is applied ontop of the existing database.
+    ss::future<std::unique_ptr<internal::iterator>>
+      create_iterator(ss::optimized_optional<ss::lw_shared_ptr<memtable>>);
 
     // Flush any pending state in memtables to disk.
     ss::future<> flush();
@@ -72,7 +74,14 @@ class impl {
     ss::future<> close();
 
 private:
-    ss::future<std::unique_ptr<internal::iterator>> create_internal_iterator();
+    // Create an iterator over the database. Note that this iterator
+    // results in ALL entries from the database, a deduplicating iterator
+    // needs to be added on top to give a traditional iterator view.
+    //
+    // If a non-null memtable is passed in, then a frozen state of the memtable
+    // is applied ontop of the existing database.
+    ss::future<std::unique_ptr<internal::iterator>>
+      create_internal_iterator(ss::optimized_optional<memtable*>);
 
     ss::future<> recover();
 

src/v/lsm/db/memtable.h

@@ -68,7 +68,8 @@ class memtable : public ss::enable_lw_shared_from_this<memtable> {
 
     // Create an iterator for this memtable.
     //
-    // This iterator is safe to use in face of concurrent updates
+    // This iterator is safe to use in face of concurrent updates, but users can
+    // see newer values that are added since the iterator is created.
     std::unique_ptr<internal::iterator> create_iterator();
 
     // The approximate amount of memory used for this memtable.

src/v/lsm/db/tests/db_bench.cc

@@ -475,7 +475,7 @@ class benchmark {
         stats.bytes_read = 0;
 
         // Perform sequential iteration
-        auto iter = co_await _db->create_iterator();
+        auto iter = co_await _db->create_iterator(std::nullopt);
         co_await iter->seek_to_first();
 
         size_t count = 0;
@@ -667,7 +667,7 @@ class benchmark {
 
         // Build a set of all keys from iterator
         std::map<ss::sstring, iobuf> db_keys;
-        auto iter = co_await _db->create_iterator();
+        auto iter = co_await _db->create_iterator(std::nullopt);
         co_await iter->seek_to_first();
 
         while (iter->valid()) {

src/v/lsm/db/tests/impl_test.cc

@@ -118,7 +118,7 @@ class ImplTest : public testing::Test {
     }
 
     testing::AssertionResult matches_shadow() {
-        auto iter = _db->create_iterator().get();
+        auto iter = _db->create_iterator(std::nullopt).get();
         auto it = _shadow.begin();
         std::vector<std::string> errors;
         for (iter->seek_to_first().get(); iter->valid(); iter->next().get()) {
@@ -220,4 +220,83 @@ TEST_F(ImplTest, Randomized) {
     }
 }
 
+TEST_F(ImplTest, ReadYourOwnWrites) {
+    // Write some initial data to the database
+    auto batch1 = ss::make_lw_shared<lsm::db::memtable>();
+    auto seqno = _db->max_applied_seqno();
+    auto key1 = lsm::internal::key::encode({
+      .key = "key1",
+      .seqno = ++seqno,
+    });
+    auto value1 = iobuf::from("value1");
+    batch1->put(key1, value1.share());
+
+    auto key2 = lsm::internal::key::encode({
+      .key = "key2",
+      .seqno = ++seqno,
+    });
+    auto value2 = iobuf::from("value2");
+    batch1->put(key2, value2.share());
+
+    _db->apply(batch1).get();
+
+    // Create a pending write batch with new keys (not yet applied)
+    auto pending_batch = ss::make_lw_shared<lsm::db::memtable>();
+    seqno = _db->max_applied_seqno();
+
+    auto key3 = lsm::internal::key::encode({
+      .key = "key3",
+      .seqno = ++seqno,
+    });
+    auto value3 = iobuf::from("value3");
+    pending_batch->put(key3, value3.share());
+
+    // Verify new key is visible through iterator with write batch
+    auto iter = _db->create_iterator(pending_batch).get();
+    std::map<ss::sstring, iobuf> seen;
+    for (iter->seek_to_first().get(); iter->valid(); iter->next().get()) {
+        seen.insert_or_assign(
+          ss::sstring(iter->key().user_key()), iter->value());
+    }
+    EXPECT_EQ(seen.size(), 3);
+    EXPECT_EQ(seen["key1"], value1);
+    EXPECT_EQ(seen["key2"], value2);
+    EXPECT_EQ(seen["key3"], value3);
+
+    // Create another pending write batch that updates an existing key
+    auto update_batch = ss::make_lw_shared<lsm::db::memtable>();
+    seqno = _db->max_applied_seqno();
+
+    auto key2_updated = lsm::internal::key::encode({
+      .key = "key2",
+      .seqno = ++seqno,
+    });
+    auto value2_updated = iobuf::from("value2_updated");
+    update_batch->put(key2_updated, value2_updated.share());
+
+    // Verify updated value is visible through iterator with write batch
+    auto iter2 = _db->create_iterator(update_batch).get();
+    seen.clear();
+    for (iter2->seek_to_first().get(); iter2->valid(); iter2->next().get()) {
+        seen.insert_or_assign(
+          ss::sstring(iter2->key().user_key()), iter2->value());
+    }
+    EXPECT_EQ(seen.size(), 2);
+    EXPECT_EQ(seen["key1"], value1);
+    EXPECT_EQ(seen["key2"], value2_updated);
+
+    // Apply the update batch and verify all data is now committed
+    _db->apply(update_batch).get();
+
+    auto iter3 = _db->create_iterator(std::nullopt).get();
+    seen.clear();
+    for (iter3->seek_to_first().get(); iter3->valid(); iter3->next().get()) {
+        seen.insert_or_assign(
+          ss::sstring(iter3->key().user_key()), iter3->value());
+    }
+    EXPECT_EQ(seen.size(), 2);
+    EXPECT_EQ(seen["key1"], value1);
+    EXPECT_EQ(seen["key2"], value2_updated);
+}
+
 } // namespace

src/v/lsm/lsm.cc

@@ -105,12 +105,15 @@ ss::future<std::optional<iobuf>> database::get(std::string_view target) {
 }
 
 ss::future<iterator> database::create_iterator() {
-    auto iter = co_await _impl->create_iterator();
+    auto iter = co_await _impl->create_iterator(std::nullopt);
     co_return iterator(std::move(iter));
 }
 
-write_batch::write_batch()
-  : _batch(ss::make_lw_shared<db::memtable>()) {}
+write_batch database::create_write_batch() { return write_batch{_impl.get()}; }
+
+write_batch::write_batch(db::impl* db)
+  : _batch(ss::make_lw_shared<db::memtable>())
+  , _db(db) {}
 write_batch::write_batch(write_batch&&) noexcept = default;
 write_batch& write_batch::operator=(write_batch&&) noexcept = default;
 write_batch::~write_batch() noexcept = default;
@@ -133,4 +136,22 @@ void write_batch::remove(std::string_view key, model::offset offset) {
     _batch->remove(std::move(k));
 }
 
+ss::future<std::optional<iobuf>> write_batch::get(std::string_view target) {
+    auto key = internal::key::encode({
+      .key = target,
+      .seqno = internal::sequence_number::max(),
+      .type = internal::value_type::value,
+    });
+    auto result = _batch->get(key);
+    if (result.is_missing()) {
+        result = co_await _db->get(key);
+    }
+    co_return std::move(result).take_value();
+}
+
+ss::future<iterator> write_batch::create_iterator() {
+    auto iter = co_await _db->create_iterator(_batch);
+    co_return iterator(std::move(iter));
+}
+
 } // namespace lsm

src/v/lsm/lsm.h

@@ -59,19 +59,32 @@ class database {
 
     // The maximum offset that has been persisted to durable storage.
     model::offset max_persisted_offset() const;
+
     // The maximum offset that has been applied to the database (persisted or
     // not).
     model::offset max_applied_offset() const;
 
     // Apply a batch of data atomically to the database.
+    //
+    // Any gets/iteration on the write batch must be finished before calling
+    // this method.
     ss::future<> apply(write_batch);
 
     // Lookup a value in the database
     ss::future<std::optional<iobuf>> get(std::string_view key);
 
     // Create an iterator over the database.
+    //
+    // This iterator gives snapshot isolation, so writes applied after
+    // this future completes will not be seen by the iterator.
     ss::future<iterator> create_iterator();
 
+    // Create a write batch that can be applied to the the database.
+    //
+    // This write batch can only be used with this database and it's lifetime is
+    // tied to this database as well.
+    write_batch create_write_batch();
+
 private:
     std::unique_ptr<db::impl> _impl;
 };
@@ -85,9 +98,11 @@ class iterator {
     iterator& operator=(const iterator&) = delete;
     iterator& operator=(iterator&&) noexcept;
     ~iterator() noexcept;
+
     // An iterator is either positioned at a key/value pair, or
     // not valid. This method returns true iff the iterator is valid.
     bool valid() const;
+
     // Position at the first key in the source. The iterator is valid()
     // after this call iff the source is not empty.
     ss::future<> seek_to_first();
@@ -124,10 +139,14 @@ class iterator {
     std::unique_ptr<internal::iterator> _impl;
 };
 
-// A batch of data that can be applied to the database.
+// A batch of data that can be applied atomically to the database.
+//
+// In addition to being a staging ground for write operations,
+// the write batch also supports reading the staged data as a view over the
+// existing database so that you can Read-Your-Own-Writes before data is
+// committed to the LSM tree.
 class write_batch {
 public:
-    write_batch();
     write_batch(const write_batch&) = delete;
     write_batch(write_batch&&) noexcept;
     write_batch& operator=(const write_batch&) = delete;
@@ -144,9 +163,32 @@ class write_batch {
     // REQUIRES: offsets must be monotonically increasing as added to the batch.
     void remove(std::string_view key, model::offset);
 
+    // Lookup a value in the database as if the current write batch was applied.
+    //
+    // The returned future must finish resolving before the write_batch is
+    // applied to the database.
+    ss::future<std::optional<iobuf>> get(std::string_view key);
+
+    // Create an iterator over the database as if the current write batch was
+    // applied.
+    //
+    // This iterator gives snapshot isolation, so  writes applied after
+    // this future completes will not be seen by the iterator as long as
+    // these writes do not have the same offset.
+    //
+    // The resulting iterator is safe to use concurrently with additional writes
+    // being applied to the batch.
+    //
+    // The returned iterator must not be used after the write_batch is applied
+    // to the database.
+    ss::future<iterator> create_iterator();
+
 private:
+    explicit write_batch(db::impl*);
+
     friend class database;
     ss::lw_shared_ptr<db::memtable> _batch;
+    db::impl* _db;
 };
 
 } // namespace lsm