server: Carry db index in item

Signed-off-by: Abhijat Malviya <abhijat@dragonflydb.io>

Author: abhijat

src/server/rdb_load.cc

@@ -2803,14 +2803,14 @@ void RdbLoader::FlushShardAsync(ShardId sid) {
   if (out_buf.empty())
     return;
 
-  auto cb = [indx = this->cur_db_index_, this, ib = std::move(out_buf)] {
+  auto cb = [this, ib = std::move(out_buf)] {
     auto& db_slice = GetCurrentDbSlice();
 
     // Before we start loading, increment LoadInProgress.
     // This is required because FlushShardAsync dispatches to multiple shards, and those shards
     // might have not yet have their state (load in progress) incremented.
     db_slice.IncrLoadInProgress();
-    this->LoadItemsBuffer(indx, ib);
+    this->LoadItemsBuffer(ib);
     db_slice.DecrLoadInProgress();
   };
 
@@ -2960,17 +2960,19 @@ void RdbLoader::CreateObjectOnShard(const DbContext& db_cntx, const Item* item,
   }
 }
 
-void RdbLoader::LoadItemsBuffer(DbIndex db_ind, const ItemsBuf& ib) {
+void RdbLoader::LoadItemsBuffer(const ItemsBuf& ib) {
   EngineShard* es = EngineShard::tlocal();
-  DbContext db_cntx{&namespaces->GetDefaultNamespace(), db_ind, GetCurrentTimeMs()};
-  DbSlice& db_slice = db_cntx.GetDbSlice(es->shard_id());
-
-  DCHECK(!db_slice.IsCacheMode());
+  const uint64_t now_ms = GetCurrentTimeMs();
+  Namespace* ns = &namespaces->GetDefaultNamespace();
 
   for (const auto* item : ib) {
+    DbContext db_cntx{ns, item->db_index, now_ms};
+    DbSlice& db_slice = db_cntx.GetDbSlice(es->shard_id());
+    DCHECK(!db_slice.IsCacheMode());
     CreateObjectOnShard(db_cntx, item, &db_slice);
     if (stop_early_) {
-      return;
+      // force all items in ib to move into item_queue_ so they can be cleaned up later.
+      break;
     }
   }
 
@@ -3121,6 +3123,7 @@ io::Result<bool> RdbLoader::ReadAndDispatchObject(int object_type, std::string&
   item->has_mc_flags = obj_settings.has_mc_flags;
   item->mc_flags = obj_settings.mc_flags;
   item->expire_ms = obj_settings.expiretime;
+  item->db_index = db_index;
 
   std::move(cleanup).Cancel();
 

src/server/rdb_load.h

@@ -353,6 +353,7 @@ class RdbLoader : protected RdbLoaderBase {
     OpaqueObj val;
     uint64_t expire_ms;
     std::atomic<Item*> next;
+    DbIndex db_index = 0;
     bool is_sticky = false;
     bool has_mc_flags = false;
     uint32_t mc_flags = 0;
@@ -395,7 +396,7 @@ class RdbLoader : protected RdbLoaderBase {
   void FlushShardAsync(ShardId sid);
   void FlushAllShards();
 
-  void LoadItemsBuffer(DbIndex db_ind, const ItemsBuf& ib);
+  void LoadItemsBuffer(const ItemsBuf& ib);
 
   void CreateObjectOnShard(const DbContext& db_cntx, const Item* item, DbSlice* db_slice);
 

src/server/rdb_test.cc

@@ -1470,32 +1470,56 @@ TEST_F(RdbTest, LoadTwoChunks) {
 }
 
 TEST_F(RdbTest, InterleavedLoad) {
+  // must have >1 shards for non inlined path check. find a key that lands in shard 1 by hashing, to
+  // test non inlined obj. creation
+  ASSERT_GT(shard_set->size(), 1u);
+  std::string key;
+  for (unsigned i = 0; i < 1000; ++i) {
+    key = StrCat("x", i);
+    if (Shard(key, shard_set->size()) == 1)
+      break;
+  }
+  ASSERT_EQ(Shard(key, shard_set->size()), 1u);
+
   std::string a1;
+  // hash chunk 1
   a1.push_back(RDB_TYPE_HASH);
-  AppendString(&a1, "a");
+  AppendString(&a1, key);
   AppendLen(&a1, 2);
   AddKV(&a1, "f1", "v1");
 
+  // string
   std::string b;
   b.push_back(RDB_TYPE_STRING);
   AppendString(&b, "b");
   AppendString(&b, "plain");
 
+  // hash chunk 2
   std::string a2;
   AddKV(&a2, "f2", "v2");
 
   std::string body;
-  // interleave - one tag, one plain, then one tag for the same id=1
+  // chunk for db 0
   body += MakeTaggedChunk(1, a1);
+  // simple string b=plain
   body += b;
+  body.push_back(static_cast<char>(RDB_OPCODE_SELECTDB));
+  // switch to db 1
+  AppendLen(&body, 1);
+  // back to chunk for db 0
   body += MakeTaggedChunk(1, a2);
 
   auto ec = pp_->at(0)->Await([&] { return LoadRdbData(service_.get(), WrapInRdb(body)); });
   ASSERT_FALSE(ec) << ec.message();
 
-  EXPECT_EQ(Run({"HGET", "a", "f1"}), "v1");
-  EXPECT_EQ(Run({"HGET", "a", "f2"}), "v2");
+  EXPECT_EQ(Run({"SELECT", "0"}), "OK");
+  EXPECT_EQ(Run({"HGET", key, "f1"}), "v1");
+  EXPECT_EQ(Run({"HGET", key, "f2"}), "v2");
   EXPECT_EQ(Run({"GET", "b"}), "plain");
+
+  EXPECT_EQ(Run({"SELECT", "1"}), "OK");
+  EXPECT_THAT(Run({"EXISTS", key}), IntArg(0));
+  EXPECT_EQ(Run({"SELECT", "0"}), "OK");
 }
 
 TEST_F(RdbTest, ChunksAroundJournalOffset) {