bug(server): fix lns mismatch in replication

src/server/db_slice.cc

@@ -1167,7 +1167,6 @@ DbSlice::PrimeItAndExp DbSlice::ExpireIfNeeded(const Context& cntx, PrimeIterato
   }
 
   auto& db = db_arr_[cntx.db_index];
-
   auto expire_it = db->expire.Find(it->first);
 
   if (IsValid(expire_it)) {
@@ -1184,6 +1183,9 @@ DbSlice::PrimeItAndExp DbSlice::ExpireIfNeeded(const Context& cntx, PrimeIterato
                 << ", prime table size: " << db->prime.size() << util::fb2::GetStacktrace();
   }
 
+  DCHECK(shard_owner()->shard_lock()->Check(IntentLock::Mode::EXCLUSIVE))
+      << util::fb2::GetStacktrace();
+
   string scratch;
   string_view key = it->first.GetSlice(&scratch);
 

src/server/journal/streamer.cc

@@ -57,21 +57,14 @@ void JournalStreamer::Start(util::FiberSocketBase* dest, bool send_lsn) {
   dest_ = dest;
   journal_cb_id_ =
       journal_->RegisterOnChange([this, send_lsn](const JournalItem& item, bool allow_await) {
-        if (allow_await) {
-          ThrottleIfNeeded();
-          // No record to write, just await if data was written so consumer will read the data.
-          // TODO: shouldnt we trigger async write in noop??
-          if (item.opcode == Op::NOOP)
-            return;
-        }
-
         if (!ShouldWrite(item)) {
           return;
         }
 
+        DCHECK_GT(item.lsn, last_lsn_writen_);
         Write(item.data);
         time_t now = time(nullptr);
-
+        last_lsn_writen_ = item.lsn;
         // TODO: to chain it to the previous Write call.
         if (send_lsn && now - last_lsn_time_ > 3) {
           last_lsn_time_ = now;
@@ -80,6 +73,10 @@ void JournalStreamer::Start(util::FiberSocketBase* dest, bool send_lsn) {
           writer.Write(Entry{journal::Op::LSN, item.lsn});
           Write(std::move(sink).str());
         }
+
+        if (allow_await) {
+          ThrottleIfNeeded();
+        }
       });
 }
 

src/server/journal/streamer.h

@@ -48,7 +48,7 @@ class JournalStreamer {
   void ThrottleIfNeeded();
 
   virtual bool ShouldWrite(const journal::JournalItem& item) const {
-    return cntx_->IsRunning();
+    return cntx_->IsRunning() && item.opcode != journal::Op::NOOP;
   }
 
   void WaitForInflightToComplete();
@@ -68,6 +68,7 @@ class JournalStreamer {
 
   size_t in_flight_bytes_ = 0, total_sent_ = 0;
   time_t last_lsn_time_ = 0;
+  LSN last_lsn_writen_ = 0;
   util::fb2::EventCount waker_;
   uint32_t journal_cb_id_{0};
 };

src/server/snapshot.cc

@@ -131,6 +131,7 @@ void SliceSnapshot::FinalizeJournalStream(bool cancel) {
   journal->UnregisterOnChange(cb_id);
   if (!cancel) {
     // always succeeds because serializer_ flushes to string.
+    VLOG(1) << "FinalizeJournalStream lsn: " << journal->GetLsn();
     std::ignore = serializer_->SendJournalOffset(journal->GetLsn());
     PushSerialized(true);
   }

src/server/transaction.cc

@@ -149,11 +149,19 @@ cv_status Transaction::BatonBarrier::Wait(time_point tp) {
 
 Transaction::Guard::Guard(Transaction* tx) : tx(tx) {
   DCHECK(tx->cid_->opt_mask() & CO::GLOBAL_TRANS);
-  tx->Execute([](auto*, auto*) { return OpStatus::OK; }, false);
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    namespaces->GetDefaultNamespace().GetDbSlice(shard->shard_id()).SetExpireAllowed(false);
+    return OpStatus::OK;
+  };
+  tx->Execute(cb, false);
 }
 
 Transaction::Guard::~Guard() {
-  tx->Conclude();
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    namespaces->GetDefaultNamespace().GetDbSlice(shard->shard_id()).SetExpireAllowed(true);
+    return OpStatus::OK;
+  };
+  tx->Execute(cb, true);
   tx->Refurbish();
 }
 

tests/dragonfly/replication_test.py

@@ -2959,12 +2959,15 @@ async def test_preempt_in_atomic_section_of_heartbeat(df_factory: DflyInstanceFa
     await fill_task
 
 
-@pytest.mark.skip("temporarily skipped")
 async def test_bug_in_json_memory_tracking(df_factory: DflyInstanceFactory):
     """
     This test reproduces a bug in the JSON memory tracking.
     """
-    master = df_factory.create(proactor_threads=2, serialization_max_chunk_size=1)
+    master = df_factory.create(
+        proactor_threads=2,
+        serialization_max_chunk_size=1,
+        vmodule="replica=2,dflycmd=2,snapshot=1,rdb_save=1,rdb_load=1,journal_slice=2",
+    )
     replicas = [df_factory.create(proactor_threads=2) for i in range(2)]
 
     # Start instances and connect clients
@@ -2982,6 +2985,7 @@ async def test_bug_in_json_memory_tracking(df_factory: DflyInstanceFactory):
 
     seeder = SeederV2(key_target=50_000)
     fill_task = asyncio.create_task(seeder.run(master.client()))
+    await asyncio.sleep(0.2)
 
     for replica in c_replicas:
         await replica.execute_command(f"REPLICAOF LOCALHOST {master.port}")