MB-56917: Check memory when fetching item from disk in Magma

mohammadzaeem · mohammadzaeem · commit 2b05aaaaad74 · 2024-01-10T14:18:55.000Z
Allow MagmaKVStore::makeItem to use createItemCb to determine if sufficient memory is available to create a new item. If there is available memory, then a GetValue with a pointer to a new item and an engine_errc::success is returned. Otherwise, a GetValue with nullptr is returned as an item and engine_errc::temp_faillure / no_memory. Change-Id: I0dd236fb7889d7ccc8858a215c102377f38a5005 Reviewed-on: https://review.couchbase.org/c/kv_engine/+/203384 Reviewed-by: Jim Walker <jim@couchbase.com> Tested-by: Mohammad Zaeem <mohammad.zaeem@couchbase.com>
diff --git a/engines/ep/src/kvstore/magma-kvstore/magma-kvstore.cc b/engines/ep/src/kvstore/magma-kvstore/magma-kvstore.cc
@@ -1156,7 +1156,7 @@ void MagmaKVStore::getMulti(Vbid vbid,
                                              metaSlice,
                                              valueSlice,
                                              bg_itm_ctx->getValueFilter(),
-                                             std::move(createItemCb));
+                                             createItemCb);
             GetValue* rv = &bg_itm_ctx->value;
 
             for (auto& fetch : bg_itm_ctx->getRequests()) {
@@ -1165,8 +1165,7 @@ void MagmaKVStore::getMulti(Vbid vbid,
                         std::chrono::duration_cast<std::chrono::microseconds>(
                                 std::chrono::steady_clock::now() -
                                 fetch->initTime));
-                st.readSizeHisto.add(bg_itm_ctx->value.item->getKey().size() +
-                                     bg_itm_ctx->value.item->getNBytes());
+                st.readSizeHisto.add(op.Key.Len() + valueSlice.Len());
             }
             ++st.io_bg_fetch_docs_read;
             st.io_bgfetch_doc_bytes +=
@@ -1375,12 +1374,12 @@ bool MagmaKVStore::snapshotVBucket(Vbid vbid, const VB::Commit& meta) {
     return true;
 }
 
-std::unique_ptr<Item> MagmaKVStore::makeItem(Vbid vb,
-                                             const Slice& keySlice,
-                                             const Slice& metaSlice,
-                                             const Slice& valueSlice,
-                                             ValueFilter filter,
-                                             CreateItemCB createItemcb) const {
+GetValue MagmaKVStore::makeItem(Vbid vb,
+                                const Slice& keySlice,
+                                const Slice& metaSlice,
+                                const Slice& valueSlice,
+                                ValueFilter filter,
+                                CreateItemCB createItemCb) const {
     auto key = makeDiskDocKey(keySlice);
     const auto meta = magmakv::getDocMeta(metaSlice);
 
@@ -1390,25 +1389,33 @@ std::unique_ptr<Item> MagmaKVStore::makeItem(Vbid vb,
     const bool includeValue = filter != ValueFilter::KEYS_ONLY ||
                               key.getDocKey().isInSystemCollection();
 
-    value_t body;
+    size_t nbytes = 0;
     // Only create the body for the value filter and when a valueSlice is given
     if ((includeValue || forceValueFetch) && valueSlice.Data()) {
-        body.reset(TaggedPtr<Blob>(
+        nbytes = valueSlice.Len();
+    }
+
+    auto [status, item] = createItemCb(key.getDocKey(),
+                                       nbytes,
+                                       meta.getFlags(),
+                                       meta.getExptime(),
+                                       value_t(),
+                                       meta.getDatatype(),
+                                       meta.getCas(),
+                                       meta.getBySeqno(),
+                                       vb,
+                                       meta.getRevSeqno());
+    if (status != cb::engine_errc::success) {
+        return GetValue(nullptr, status);
+    }
+
+    // Blob creation is deferred to avoid exceeding quota.
+    if ((includeValue || forceValueFetch) && valueSlice.Data()) {
+        item->replaceValue(TaggedPtr<Blob>(
                 Blob::New(valueSlice.Data(), meta.getValueSize()),
                 TaggedPtrBase::NoTagValue));
     }
 
-    auto item =
-            std::make_unique<Item>(key.getDocKey(),
-                                   meta.getFlags(),
-                                   meta.getExptime(),
-                                   body,
-                                   meta.getDatatype(),
-                                   meta.getCas(),
-                                   meta.getBySeqno(),
-                                   vb,
-                                   meta.getRevSeqno());
-
     if (filter != ValueFilter::KEYS_ONLY) {
         checkAndFixKVStoreCreatedItem(*item);
     }
@@ -1445,13 +1452,13 @@ std::unique_ptr<Item> MagmaKVStore::makeItem(Vbid vb,
         item->setPrepareSeqno(meta.getPrepareSeqno());
     }
 
-    if (body) {
+    if (item->getValue()) {
         // function shared with KEY_ONLY creation paths, so only sanitize when
         // a body was created
         checkAndFixKVStoreCreatedItem(*item);
     }
 
-    return item;
+    return GetValue(std::move(item), status);
 }
 
 GetValue MagmaKVStore::makeGetValue(Vbid vb,
@@ -1460,11 +1467,12 @@ GetValue MagmaKVStore::makeGetValue(Vbid vb,
                                     const Slice& valueSlice,
                                     ValueFilter filter,
                                     CreateItemCB createItemCb) const {
-    return GetValue(
-            makeItem(vb, keySlice, metaSlice, valueSlice, filter, createItemCb),
-            cb::engine_errc::success,
-            -1,
-            false);
+    return makeItem(vb,
+                    keySlice,
+                    metaSlice,
+                    valueSlice,
+                    filter,
+                    std::move(createItemCb));
 }
 
 int MagmaKVStore::saveDocs(MagmaKVStoreTransactionContext& txnCtx,
@@ -2380,18 +2388,19 @@ ScanStatus MagmaKVStore::scanOne(
                 magmakv::isCompressed(docMeta));
     }
 
-    auto itm = makeItem(ctx.vbid,
-                        keySlice,
-                        metaSlice,
-                        value,
-                        ctx.valFilter,
-                        getDefaultCreateItemCallback());
+    auto gv = makeItem(ctx.vbid,
+                       keySlice,
+                       metaSlice,
+                       value,
+                       ctx.valFilter,
+                       getDefaultCreateItemCallback());
+    Expects(gv.getStatus() == cb::engine_errc::success);
 
     // When we are requested to return the values as compressed AND
     // the value isn't compressed, attempt to compress the value.
     if (ctx.valFilter == ValueFilter::VALUES_COMPRESSED &&
         !magmakv::isCompressed(docMeta)) {
-        if (!itm->compressValue()) {
+        if (!gv.item->compressValue()) {
             logger->warn(
                     "MagmaKVStore::scanOne failed to compress value - {} "
                     "key:{} "
@@ -2404,11 +2413,10 @@ ScanStatus MagmaKVStore::scanOne(
         }
     }
 
-    GetValue rv(std::move(itm),
-                cb::engine_errc::success,
-                -1,
-                ctx.valFilter == ValueFilter::KEYS_ONLY);
-    ctx.getValueCallback().callback(rv);
+    if (ctx.valFilter == ValueFilter::KEYS_ONLY) {
+        gv.setPartial();
+    }
+    ctx.getValueCallback().callback(gv);
     auto callbackStatus = ctx.getValueCallback().getStatus();
     if (callbackStatus == cb::engine_errc::success) {
         return ScanStatus::Success;
@@ -4120,13 +4128,13 @@ GetValue MagmaKVStore::getBySeqno(KVFileHandle& handle,
     }
 
     if (status.IsOkDocFound()) {
-        rv.item = makeItem(vbid,
-                           *key,
-                           *meta,
-                           *value,
-                           filter,
-                           getDefaultCreateItemCallback());
-        rv.setStatus(cb::engine_errc::success);
+        rv = makeItem(vbid,
+                      *key,
+                      *meta,
+                      *value,
+                      filter,
+                      getDefaultCreateItemCallback());
+        Expects(rv.getStatus() == cb::engine_errc::success);
         return rv;
     }
     return rv;
diff --git a/engines/ep/src/kvstore/magma-kvstore/magma-kvstore.h b/engines/ep/src/kvstore/magma-kvstore/magma-kvstore.h
@@ -616,12 +616,12 @@ class MagmaKVStore : public KVStore {
     CompactDBStatus compactDBInternal(std::unique_lock<std::mutex>& vbLock,
                                       std::shared_ptr<CompactionContext> ctx);
 
-    std::unique_ptr<Item> makeItem(Vbid vb,
-                                   const magma::Slice& keySlice,
-                                   const magma::Slice& metaSlice,
-                                   const magma::Slice& valueSlice,
-                                   ValueFilter filter,
-                                   CreateItemCB createItemCb) const;
+    GetValue makeItem(Vbid vb,
+                      const magma::Slice& keySlice,
+                      const magma::Slice& metaSlice,
+                      const magma::Slice& valueSlice,
+                      ValueFilter filter,
+                      CreateItemCB createItemCb) const;
 
     GetValue makeGetValue(Vbid vb,
                           const magma::Slice& keySlice,
diff --git a/engines/ep/tests/module_tests/evp_store_test.cc b/engines/ep/tests/module_tests/evp_store_test.cc
@@ -629,12 +629,14 @@ TEST_P(EPBucketFullEvictionTest, xattrExpiryOnFullyEvictedItem) {
 }
 
 TEST_P(EPBucketFullEvictionTest, GetMultiShouldNotExceedMutationWatermark) {
-    if (!isCouchstore()) {
+    // This test relies on precise memory tracking
+    const auto& stats = engine->getEpStats();
+    if (!stats.isMemoryTrackingEnabled()) {
         GTEST_SKIP();
     }
+
     ASSERT_EQ(cb::engine_errc::success,
               store->setVBucketState(vbid, vbucket_state_active, {}));
-    const auto& stats = engine->getEpStats();
     EXPECT_LT(stats.getPreciseTotalMemoryUsed(), stats.getMaxDataSize());
 
     // Load a document of size 1MiB
@@ -662,18 +664,18 @@ TEST_P(EPBucketFullEvictionTest, GetMultiShouldNotExceedMutationWatermark) {
     // run bgFetch
     runBGFetcherTask();
 
-    // BGFetch should fail to allocate memory above bucekt quota (now works for
-    // couch-kvstore), however this still fails for magma
     auto vb = store->getVBucket(vbid);
     EXPECT_FALSE(vb->hasPendingBGFetchItems());
     EXPECT_LE(stats.getPreciseTotalMemoryUsed(), stats.getMaxDataSize());
 }
 
 TEST_P(EPBucketFullEvictionTest, BgfetchSucceedsUntilMutationWatermark) {
-    if (!isCouchstore()) {
+    // This test involves adjusting the bucket quota for the backend. Since
+    // Nexus utilises both Magma and Couchstore, manipulating the bucket quota
+    // to accommodate one backend will result in failure for the other.
+    if (isNexus()) {
         GTEST_SKIP();
     }
-
     // This test relies on precise memory tracking
     const auto& stats = engine->getEpStats();
     if (!stats.isMemoryTrackingEnabled()) {
@@ -688,17 +690,17 @@ TEST_P(EPBucketFullEvictionTest, BgfetchSucceedsUntilMutationWatermark) {
 
     // Load documents of size 1MiB
     const std::string value(valueSize, 'x');
-    auto key1 = makeStoredDocKey("key_0");
-    auto key2 = makeStoredDocKey("key_1");
+    auto key0 = makeStoredDocKey("key_0");
+    auto key1 = makeStoredDocKey("key_1");
 
-    // store an items
+    // store items
+    store_item(vbid, key0, value);
     store_item(vbid, key1, value);
-    store_item(vbid, key2, value);
     flush_vbucket_to_disk(vbid, 2);
 
     // evict items
+    evict_key(vbid, key0);
     evict_key(vbid, key1);
-    evict_key(vbid, key2);
 
     // check item has been removed
     auto options = static_cast<get_options_t>(
@@ -707,37 +709,64 @@ TEST_P(EPBucketFullEvictionTest, BgfetchSucceedsUntilMutationWatermark) {
     auto cookie1 = create_mock_cookie();
     auto cookie2 = create_mock_cookie();
 
-    auto gv = store->get(key1, vbid, cookie1, options);
+    // Couchstore sorts key pointers and retrieves the first item in the queue.
+    // Initially, the queue is ["Key_1", "Key_0"], which is then sorted to
+    // ["Key_0", "Key_1"]. In Magma, there is no sorting; instead, the first
+    // item from the queue is popped and requested. When a new item is
+    // requested, it is pushed to the front of the queue. To ensure Key_0 always
+    // succeeds and Key_1 fails, we first request Key_1 ["Key_1"] and then
+    // request Key_0 ["Key_0", "Key_1"].
+    auto gv = store->get(key1, vbid, cookie2, options);
     EXPECT_EQ(cb::engine_errc::would_block, gv.getStatus());
-    gv = store->get(key2, vbid, cookie2, options);
+    gv = store->get(key0, vbid, cookie1, options);
     EXPECT_EQ(cb::engine_errc::would_block, gv.getStatus());
 
     // set bucket quota to current memory usage + 3MiB such that the first
     // bgfetch will pass and the second one will fail:
     // couchstore fetches value + addition overhead = ~1.3MiB
     // We want an additional 1MiB for fetching the first item successfully
-    engine->setMaxDataSize(stats.getPreciseTotalMemoryUsed() + (3 * valueSize));
+    // Magma requires an addition 1MiB overhead than couchstore.
+    if (engine->getConfiguration().getBackend() == "magma") {
+        engine->setMaxDataSize(stats.getPreciseTotalMemoryUsed() +
+                               (4 * valueSize));
+    } else {
+        engine->setMaxDataSize(stats.getPreciseTotalMemoryUsed() +
+                               (3 * valueSize));
+    }
 
     int callbackCounter = 0;
-    cookie1->setUserNotifyIoComplete(
-            [&callbackCounter](cb::engine_errc status) {
-                EXPECT_EQ(cb::engine_errc::success, status);
-                callbackCounter++;
-            });
-    cookie2->setUserNotifyIoComplete(
-            [&callbackCounter](cb::engine_errc status) {
-                EXPECT_EQ(cb::engine_errc::temporary_failure, status);
-                callbackCounter++;
-            });
+    int numSucess = 0;
+    int numFailure = 0;
+    cookie1->setUserNotifyIoComplete([&callbackCounter,
+                                      &numSucess,
+                                      &numFailure](cb::engine_errc status) {
+        if (status == cb::engine_errc::success) {
+            numSucess++;
+        } else {
+            numFailure++;
+        }
+        callbackCounter++;
+    });
+    cookie2->setUserNotifyIoComplete([&callbackCounter,
+                                      &numSucess,
+                                      &numFailure](cb::engine_errc status) {
+        if (status == cb::engine_errc::success) {
+            numSucess++;
+        } else {
+            numFailure++;
+        }
+        callbackCounter++;
+    });
 
     // run bgFetch
     runBGFetcherTask();
 
-    // BGFetch should fail to allocate memory above bucekt quota (now works for
-    // couch-kvstore), however this still fails for magma
+    // BGFetch should fail to allocate memory above bucekt quota
     auto vb = store->getVBucket(vbid);
     EXPECT_FALSE(vb->hasPendingBGFetchItems());
     EXPECT_LE(stats.getPreciseTotalMemoryUsed(), stats.getMaxDataSize());
+    EXPECT_EQ(1, numSucess);
+    EXPECT_EQ(1, numFailure);
     EXPECT_EQ(2, callbackCounter);
 
     destroy_mock_cookie(cookie1);
@@ -1075,7 +1104,7 @@ TEST_P(EPBucketFullEvictionTest, ExpiryFindsPrepareWithSameCas) {
     vb.reset();
     // Before destroying the engine, prevent createItemCallback from becoming
     // invalid
-    createItemCallback = nullptr;
+    createItemCallback = KVStoreIface::getDefaultCreateItemCallback();
     resetEngineAndWarmup();
     createItemCallback = this->engine->getCreateItemCallback();
     vb = store->getVBucket(vbid);
