feat(server): Journal omits

src/server/db_slice.cc

@@ -4,7 +4,10 @@
 
 #include "server/db_slice.h"
 
+#include <ranges>
+
 #include "core/dense_set.h"
+#include "core/overloaded.h"
 #include "strings/human_readable.h"
 
 extern "C" {
@@ -365,7 +368,7 @@ DbStats& DbStats::operator+=(const DbStats& o) {
 }
 
 SliceEvents& SliceEvents::operator+=(const SliceEvents& o) {
-  static_assert(sizeof(SliceEvents) == 136, "You should update this function with new fields");
+  static_assert(sizeof(SliceEvents) == 144, "You should update this function with new fields");
 
   ADD(evicted_keys);
   ADD(hard_evictions);
@@ -384,6 +387,7 @@ SliceEvents& SliceEvents::operator+=(const SliceEvents& o) {
   ADD(ram_misses);
   ADD(huff_encode_total);
   ADD(huff_encode_success);
+  ADD(journal_omit);
   return *this;
 }
 
@@ -539,12 +543,11 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::FindMutable(const Context& cntx, string
 OpResult<DbSlice::ItAndUpdater> DbSlice::FindMutableInternal(const Context& cntx, string_view key,
                                                              std::optional<unsigned> req_obj_type) {
   auto res = FindInternal(cntx, key, req_obj_type, UpdateStatsMode::kMutableStats);
-  if (!res.ok()) {
-    return res.status();
-  }
+  RETURN_ON_BAD_STATUS(res);
 
   auto it = Iterator(*res, StringOrView::FromView(key));
   PreUpdateBlocking(cntx.db_index, it);
+
   // PreUpdate() might have caused a deletion of `it`
   if (res->IsOccupied()) {
     DCHECK_GE(db_arr_[cntx.db_index]->stats.obj_memory_usage, (*res)->second.MallocUsed());
@@ -667,11 +670,17 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::AddOrFindInternal(const Context& cntx,
 
   if (res.ok()) {
     Iterator it(*res, StringOrView::FromView(key));
-    PreUpdateBlocking(cntx.db_index, it);
+
+    bool omitted_journal = IsOmittableWrite(cntx, it.GetInnerIt());
+    if (!omitted_journal)
+      PreUpdateBlocking(cntx.db_index, it);
 
     // PreUpdate() might have caused a deletion of `it`
     if (res->IsOccupied()) {
-      return ItAndUpdater{.it = it, .post_updater{cntx.db_index, key, it, this}, .is_new = false};
+      return ItAndUpdater{.it = it,
+                          .post_updater{cntx.db_index, key, it, this},
+                          .is_new = false,
+                          .omitted_journal = omitted_journal};
     } else {
       res = OpStatus::KEY_NOTFOUND;
     }
@@ -683,9 +692,13 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::AddOrFindInternal(const Context& cntx,
   auto status = res.status();
   CHECK(status == OpStatus::KEY_NOTFOUND || status == OpStatus::OUT_OF_MEMORY) << status;
 
+  bool omit_journal = false;
   if (!change_cb_.empty()) {
     auto bucket_set = db.prime.CVCUponInsert(key);
-    CallChangeCallbacks(cntx.db_index, bucket_set);
+
+    omit_journal = IsOmittableWrite(cntx, {bucket_set});
+    if (!omit_journal)
+      CallChangeCallbacks(cntx.db_index, bucket_set);
 
     // Set of possible insertion buckets must be the same after possibly blocking call
     DCHECK(bucket_set == db.prime.CVCUponInsert(key));
@@ -778,7 +791,9 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::AddOrFindInternal(const Context& cntx,
   }
 
   DCHECK_EQ(it->second.MallocUsed(), 0UL);  // Make sure accounting is no-op
-  it.SetVersion(NextVersion());
+
+  if (!omit_journal)
+    it.SetVersion(NextVersion());
 
   TouchTopKeysIfNeeded(key, db.sample_top_keys);
   TouchHllIfNeeded(key, db.sample_unique_keys);
@@ -796,7 +811,8 @@ OpResult<DbSlice::ItAndUpdater> DbSlice::AddOrFindInternal(const Context& cntx,
   return ItAndUpdater{
       .it = Iterator(it, StringOrView::FromView(key)),
       .post_updater{cntx.db_index, key, Iterator(it, StringOrView::FromView(key)), this},
-      .is_new = true};
+      .is_new = true,
+      .omitted_journal = omit_journal};
 }
 
 void DbSlice::ActivateDb(DbIndex db_ind) {
@@ -929,7 +945,7 @@ void DbSlice::FlushSlots(const cluster::SlotRanges& slot_ranges) {
     }
   };
 
-  uint64_t cb_id = RegisterOnChange(std::move(on_change));
+  uint64_t cb_id = RegisterOnChange(false, std::move(on_change));
 
   fb2::Fiber("flush_slots", [this, shared_slots, next_version, cb_id]() {
     FlushSlotsFb(*shared_slots, next_version, cb_id);
@@ -1340,9 +1356,9 @@ void DbSlice::ExpireAllIfNeeded() {
   }
 }
 
-uint64_t DbSlice::RegisterOnChange(ChangeCallback cb) {
+uint64_t DbSlice::RegisterOnChange(bool replica, ChangeCallback cb) {
   DCHECK(!owner_->shard_lock()->IsFree());
-  return change_cb_.emplace_back(NextVersion(), std::move(cb)).first;
+  return std::get<0>(change_cb_.emplace_back(NextVersion(), replica, std::move(cb)));
 }
 
 // Ordering invariant (PIT mode):
@@ -1361,9 +1377,9 @@ void DbSlice::FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_
            << ", upper_bound=" << upper_bound;
 
   const size_t limit = change_cb_.size();
-  auto ccb = change_cb_.begin();
+  auto change_cb_iter = change_cb_.begin();
   for (size_t i = 0; i < limit; ++i) {
-    uint64_t cb_version = ccb->first;
+    const auto& [cb_version, _, cb] = *change_cb_iter;
     DCHECK_LE(cb_version, upper_bound);
     if (cb_version == upper_bound) {
       return;
@@ -1378,17 +1394,17 @@ void DbSlice::FlushChangeToEarlierCallbacks(DbIndex db_ind, Iterator it, uint64_
     // does not change during the serialization, therefore we allow at most one serializer
     // reading the bucket at the same time.
     if (bucket_version <= cb_version) {
-      ccb->second(db_ind, ChangeReq{it.GetInnerIt()});
+      cb(db_ind, ChangeReq{it.GetInnerIt()});
     }
-    ++ccb;
+    ++change_cb_iter;
   }
 }
 
 //! Unregisters the callback.
 void DbSlice::UnregisterOnChange(uint64_t id) {
   serialization_latch_.Wait();
   auto it = find_if(change_cb_.begin(), change_cb_.end(),
-                    [id](const auto& cb) { return cb.first == id; });
+                    [id](const auto& cb) { return std::get<0>(cb) == id; });
   CHECK(it != change_cb_.end());
   change_cb_.erase(it);
 }
@@ -1921,10 +1937,34 @@ void DbSlice::CallChangeCallbacks(DbIndex id, const ChangeReq& cr) const {
   const size_t limit = change_cb_.size();
   auto ccb = change_cb_.begin();
   for (size_t i = 0; i < limit; ++i) {
-    CHECK(ccb->second);
-    ccb->second(id, cr);
+    const auto& [_v, _r, cb] = *ccb;
+    CHECK(cb);
+    cb(id, cr);
     ++ccb;
   }
 }
 
+// We can omit the journal write if:
+// 1. it supports mutation hints and uses only a single key
+// 2. there is a single eventually-consistent snapshot (i.e. replica full sync)
+// 3. there are no other journal consumers
+// 4. the snapshot did not reach the bucket yet
+bool DbSlice::IsOmittableWrite(const Context& cntx, ChangeReq req) {
+  auto cb1 = [](PrimeTable::bucket_iterator it) { return it.GetVersion(); };
+  auto cb2 = [cb1](const PrimeTable::BucketSet& bs) -> uint64_t {
+    DCHECK(!std::ranges::empty(bs.buckets()));
+    return std::ranges::max(bs.buckets(), {}, cb1).GetVersion();
+  };
+
+  bool omit_update = false;
+  if (cntx.is_omittable_operation) {
+    uint64_t max_v = std::visit(Overloaded{cb1, cb2}, req);
+    bool allowed_snapshot = change_cb_.size() == 1 && std::get<1>(change_cb_.front()) &&
+                            max_v == 0 && max_v < std::get<0>(change_cb_.front());
+    omit_update = allowed_snapshot && journal::GetCallbackCount() == 1;
+    events_.journal_omit += unsigned(omit_update);
+  }
+  return omit_update;
+}
+
 }  // namespace dfly

src/server/db_slice.h

@@ -74,6 +74,9 @@ struct SliceEvents {
   // how many updates and insertions of keys between snapshot intervals
   size_t update = 0;
 
+  // how many journal omit optimizations were performed
+  size_t journal_omit = 0;
+
   uint64_t huff_encode_total = 0, huff_encode_success = 0;
 
   SliceEvents& operator+=(const SliceEvents& o);
@@ -248,6 +251,10 @@ class DbSlice {
     Iterator it;
     AutoUpdater post_updater;
     bool is_new = false;
+
+    // Set if DbContext::is_omittable_operation was set and the conditions were met.
+    // Means that the journal write should NOT be performed.
+    bool omitted_journal = false;
   };
 
   ItAndUpdater FindMutable(const Context& cntx, std::string_view key);
@@ -393,7 +400,7 @@ class DbSlice {
   //! Registers the callback to be called for each change.
   //! Returns the registration id which is also the unique version of the dbslice
   //! at a time of the call.
-  uint64_t RegisterOnChange(ChangeCallback cb);
+  uint64_t RegisterOnChange(bool replica, ChangeCallback cb);
 
   bool HasRegisteredCallbacks() const {
     return !change_cb_.empty();
@@ -543,6 +550,9 @@ class DbSlice {
 
   void CreateDb(DbIndex index);
 
+  // Returns true if this write could be ignored during replication without losing consistency
+  bool IsOmittableWrite(const Context& cntx, ChangeReq req);
+
   enum class UpdateStatsMode : uint8_t {
     kReadStats,
     kMutableStats,
@@ -608,7 +618,7 @@ class DbSlice {
   mutable absl::flat_hash_set<uint64_t, FpHasher> uniq_fps_;
 
   // ordered from the smallest to largest version.
-  std::list<std::pair<uint64_t, ChangeCallback>> change_cb_;
+  std::list<std::tuple<uint64_t, bool /* replica */, ChangeCallback>> change_cb_;
 
   // Used in temporary computations in Find item and CbFinish
   // This set is used to hold fingerprints of key accessed during the run of

src/server/dragonfly_test.cc

@@ -711,7 +711,7 @@ TEST_F(DflyEngineTest, Bug496) {
     // RegisterOnChange requires the shard lock to be held (see #7153).
     shard->shard_lock()->Acquire(IntentLock::EXCLUSIVE);
     uint32_t cb_id =
-        db.RegisterOnChange([&cb_hits](DbIndex, const DbSlice::ChangeReq&) { cb_hits++; });
+        db.RegisterOnChange(false, [&cb_hits](DbIndex, const DbSlice::ChangeReq&) { cb_hits++; });
     shard->shard_lock()->Release(IntentLock::EXCLUSIVE);
 
     {

src/server/generic_family.cc

@@ -837,7 +837,7 @@ OpStatus OpExpire(const OpArgs& op_args, string_view key, const DbSlice::ExpireP
 OpResult<vector<long>> OpFieldExpire(const OpArgs& op_args, string_view key, uint32_t ttl_sec,
                                      CmdArgList values) {
   auto& db_slice = op_args.GetDbSlice();
-  auto [it, auto_updater, is_new] = db_slice.FindMutable(op_args.db_cntx, key);
+  auto [it, auto_updater, is_new, _] = db_slice.FindMutable(op_args.db_cntx, key);
 
   if (!IsValid(it) || (it->second.ObjType() != OBJ_SET && it->second.ObjType() != OBJ_HASH)) {
     std::vector<long> res(values.size(), -2);

src/server/journal/journal.cc

@@ -56,8 +56,8 @@ error_code Close() {
   return {};
 }
 
-bool HasRegisteredCallbacks() {
-  return journal_slice.HasRegisteredCallbacks();
+unsigned GetCallbackCount() {
+  return journal_slice.OnChangeCbCount();
 }
 
 bool IsLSNInBuffer(LSN lsn) {

src/server/journal/journal.h

@@ -24,7 +24,10 @@ std::error_code Close();
 
 //******* The following functions must be called in the context of the owning shard *********//
 
-bool HasRegisteredCallbacks();
+unsigned GetCallbackCount();
+inline bool HasRegisteredCallbacks() {
+  return GetCallbackCount() > 0;
+}
 
 bool IsLSNInBuffer(LSN lsn);
 

src/server/journal/journal_slice.h

@@ -41,8 +41,8 @@ class JournalSlice {
   uint32_t RegisterOnChange(JournalConsumerInterface* consumer);
   void UnregisterOnChange(uint32_t);
 
-  bool HasRegisteredCallbacks() const {
-    return !journal_consumers_arr_.empty();
+  unsigned OnChangeCbCount() const {
+    return journal_consumers_arr_.size();
   }
 
   /// Returns whether the journal entry with this LSN is available

src/server/journal/streamer.cc

@@ -432,7 +432,7 @@ void RestoreStreamer::Start(util::FiberSocketBase* dest) {
     return;
 
   VLOG(1) << "RestoreStreamer start";
-  SerializerBase::RegisterChangeListener();
+  SerializerBase::RegisterChangeListener(true);
   JournalStreamer::Start(dest);
 }
 

src/server/serializer_base.cc

@@ -129,12 +129,12 @@ SerializerBase::~SerializerBase() {
 //   same fiber, so the baseline is serialized first. big_value_mu_ prevents this callback path
 //   from interleaving with the traversal fiber's bucket serialization, which may preempt while
 //   emitting large values.
-void SerializerBase::RegisterChangeListener() {
+void SerializerBase::RegisterChangeListener(bool replica) {
   db_array_ = db_slice_->databases();  // copy pointers to survive flush
   auto cb = [this](DbIndex dbid, const ChangeReq& req) {
     std::visit([&](auto it) { OnChangeBlocking(dbid, it); }, req);
   };
-  snapshot_version_ = db_slice_->RegisterOnChange(cb);
+  snapshot_version_ = db_slice_->RegisterOnChange(replica, cb);
 }
 
 void SerializerBase::UnregisterChangeListener() {
@@ -145,7 +145,7 @@ void SerializerBase::UnregisterChangeListener() {
 bool SerializerBase::ProcessBucket(DbIndex db_index, PrimeTable::bucket_iterator it,
                                    bool on_update) {
   // Check if this bucket is stale
-  if (it.is_done() || it.GetVersion() >= snapshot_version_) {
+  if (it.GetVersion() >= snapshot_version_) {
     stats_.buckets_skipped++;
 
     // Update versions for empty buckets
@@ -161,6 +161,12 @@ bool SerializerBase::ProcessBucket(DbIndex db_index, PrimeTable::bucket_iterator
     return false;
   }
 
+  // TODO: Flushing to earlier callbacks
+  if (it.is_done()) {
+    it.SetVersion(snapshot_version_);
+    return false;
+  }
+
   // For non updates (traversal flow), flush change to earlier snapshots and
   // acquire serialization latch.
   // We must make sure that earlier snapshots serialized this bucket before we update its

src/server/serializer_base.h

@@ -93,7 +93,7 @@ class SerializerBase : public BucketDependencies, public DelayedEntryHandler {
   virtual ~SerializerBase();
 
   // Register db_slice change listener and save snapshot version.
-  void RegisterChangeListener();
+  void RegisterChangeListener(bool replica);
 
   // Unregisters the callback.  Safe to call if already unregistered.
   void UnregisterChangeListener();

src/server/server_family.cc

@@ -3373,6 +3373,8 @@ string ServerFamily::FormatInfoMetrics(const Metrics& m, std::string_view sectio
     append("lua_force_gc_calls", m.lua_stats.force_gc_calls);
     append("lua_gc_freed_memory_total", m.lua_stats.gc_freed_memory);
     append("lua_gc_duration_total_sec", m.lua_stats.gc_duration_ns * 1e-9);
+
+    append("total_journal_omits", m.events.journal_omit);
   };
 
   auto add_tiered_info = [&] {

src/server/snapshot.cc

@@ -77,7 +77,7 @@ void SliceSnapshot::Start(bool stream_journal, SnapshotFlush allow_flush) {
   DCHECK(!snapshot_fb_.IsJoinable());
 
   use_background_mode_ = absl::GetFlag(FLAGS_background_snapshotting);
-  SerializerBase::RegisterChangeListener();
+  SerializerBase::RegisterChangeListener(stream_journal);
 
   if (stream_journal) {
     journal_cb_id_ = journal::RegisterConsumer(this);

src/server/string_family.cc

@@ -110,8 +110,9 @@ class SetCmd {
 
   OpStatus CachePrevIfNeeded(const SetParams& params, DbSlice::Iterator it);
 
-  const OpArgs op_args_;
+  OpArgs op_args_;
   bool explicit_journal_;  // call RecordJournal (auto journaling disabled)
+  bool skip_journal_ = false;
 };
 
 size_t SetRangeInternal(std::string* value, size_t start, std::string_view range) {
@@ -836,9 +837,13 @@ OpStatus SetCmd::Set(const SetParams& params, string_view key, string_view value
     }
   }
 
+  // Enable journal omits for this operation
+  op_args_.db_cntx.is_omittable_operation = true;
+
   // We can use std::nullopt here because SET command can change the key type to string
   auto op_res = db_slice.AddOrFind(op_args_.db_cntx, key, std::nullopt);
   RETURN_ON_BAD_STATUS(op_res);
+  skip_journal_ = op_res->omitted_journal;
 
   if (!op_res->is_new) {
     if (auto status = CachePrevIfNeeded(params, op_res->it); status != OpStatus::OK)
@@ -934,7 +939,7 @@ void SetCmd::PostEdit(const SetParams& params, std::string_view key, std::string
     StashPrimeValue(op_args_.db_cntx.db_index, key, pv, ts, params.backpressure);
   }
 
-  if (explicit_journal_ && op_args_.shard->journal()) {
+  if (!skip_journal_ && explicit_journal_ && op_args_.shard->journal()) {
     RecordJournal(params, key, value);
   }
 }