From f8c601b91cc45842b77cb71fda0c552adda9eddc Mon Sep 17 00:00:00 2001
From: Xingbo Wang <xbw@meta.com>
Date: Tue, 11 Aug 2026 08:12:46 -0700
Subject: [PATCH] Fix io_activity assertion in stress test after
 MultiGetEntityLazy (#15076)

Summary:
Pull Request resolved: https://github.com/facebook/rocksdb/pull/15076

Fix the `io_activity == Env::IOActivity::kUnknown || io_activity == options.io_activity`
assertion failure in the crash test. The crash occurs because `TestMultiGetEntity` in
`no_batched_ops_stress.cc` changes the thread operation to `OP_GETENTITY` during its
consistency-check phase (the `check_results` lambda calls `SetThreadOperation(OP_GETENTITY)`
before calling `GetEntity` for comparison). When `MaybeTestMultiGetEntityLazy` is called
afterward, it invokes `MultiGetEntityLazy` which correctly sets
`read_options.io_activity = kMultiGetEntity`, but the thread operation is still stale at
`OP_GETENTITY`. The debug assertion in `DbStressRandomAccessFileWrapper::Read` then sees
`TEST_GetExpectedIOActivity(OP_GETENTITY) = kGetEntity != kMultiGetEntity` and aborts.

The fix restores `SetThreadOperation(OP_MULTIGETENTITY)` before each call to
`MaybeTestMultiGetEntityLazy` in `TestMultiGetEntity`.

Triggered by commit cb73fc845fe0 which added the `MaybeTestMultiGetEntityLazy` calls.
The `--lazy_entity_read_one_in=100` flag (non-default) enables this code path.

Reviewed By: pdillinger

Differential Revision: D115021723

fbshipit-source-id: e3eed600db4986cab27d7b1924282b1dfb8b0a7a
---
 db/wide/db_lazy_entity_test.cc          | 45 +++++++++++++++++++++++++
 db_stress_tool/no_batched_ops_stress.cc |  4 +++
 2 files changed, 49 insertions(+)

diff --git a/db/wide/db_lazy_entity_test.cc b/db/wide/db_lazy_entity_test.cc
index fae88db0b2..1adc4fb42c 100644
--- a/db/wide/db_lazy_entity_test.cc
+++ b/db/wide/db_lazy_entity_test.cc
@@ -16,6 +16,7 @@
 
 #include "db/db_test_util.h"
 #include "db/wide/wide_column_test_util.h"
+#include "monitoring/thread_status_util.h"
 #include "port/stack_trace.h"
 #include "rocksdb/cache.h"
 #include "rocksdb/comparator.h"
@@ -1554,6 +1555,50 @@ TEST_F(DBLazyEntityTest, SecondaryInstance) {
   ASSERT_EQ(Slice(value), big);
 }
 
+// Verify that MultiGetEntityLazy sets IOOptions::io_activity to
+// kMultiGetEntity (not kGetEntity or kUnknown) so that the stress test's
+// CheckIOActivity assertion in db_stress_env_wrapper.h passes. This
+// reproduces the scenario from T283693234 where the thread operation was left
+// stale at OP_GETENTITY after a consistency check, causing a mismatch with the
+// kMultiGetEntity activity that MultiGetEntityLazy correctly propagates.
+TEST_F(DBLazyEntityTest, MultiGetEntityLazyIOActivity) {
+  Options options = GetLazyTestOptions();
+  options.enable_thread_tracking = true;
+  DestroyAndReopen(options);
+
+  constexpr char key1[] = "k1";
+  constexpr char key2[] = "k2";
+  const std::string val(200, 'x');
+  const WideColumns cols{{kDefaultWideColumnName, val}};
+
+  ASSERT_OK(
+      db_->PutEntity(WriteOptions(), db_->DefaultColumnFamily(), key1, cols));
+  ASSERT_OK(
+      db_->PutEntity(WriteOptions(), db_->DefaultColumnFamily(), key2, cols));
+  ASSERT_OK(Flush());
+
+  // Simulate the stress test scenario: thread operation is set to
+  // OP_MULTIGETENTITY (as in the stress test OperateDb loop).
+  ThreadStatusUtil::SetThreadOperation(
+      ThreadStatus::OperationType::OP_MULTIGETENTITY);
+
+  std::array<Slice, 2> keys{Slice(key1), Slice(key2)};
+  LazyWideColumnsBatch batch;
+  std::array<Status, 2> statuses;
+
+  // MultiGetEntityLazy should set io_activity = kMultiGetEntity internally,
+  // which matches OP_MULTIGETENTITY. If the thread op were stale at
+  // OP_GETENTITY, this would trigger the assertion in debug builds.
+  db_->MultiGetEntityLazy(ReadOptions(), db_->DefaultColumnFamily(),
+                          keys.size(), keys.data(), &batch, statuses.data());
+  for (const auto& s : statuses) {
+    ASSERT_OK(s);
+  }
+  ASSERT_EQ(batch.size(), 2U);
+
+  ThreadStatusUtil::SetThreadOperation(ThreadStatus::OperationType::OP_UNKNOWN);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {
diff --git a/db_stress_tool/no_batched_ops_stress.cc b/db_stress_tool/no_batched_ops_stress.cc
index 6fe45ab4e7..f5f59bd1eb 100644
--- a/db_stress_tool/no_batched_ops_stress.cc
+++ b/db_stress_tool/no_batched_ops_stress.cc
@@ -1685,6 +1685,8 @@ class NonBatchedOpsStressTest : public StressTest {
           eager_refs[i].columns = &results[i][0].columns();
         }
       }
+      ThreadStatusUtil::SetThreadOperation(
+          ThreadStatus::OperationType::OP_MULTIGETENTITY);
       MaybeTestMultiGetEntityLazy(thread, read_opts_copy, cfh, num_keys,
                                   key_slices.data(), &eager_refs);
     } else {
@@ -1723,6 +1725,8 @@ class NonBatchedOpsStressTest : public StressTest {
           eager_refs[i].columns = &results[i].columns();
         }
       }
+      ThreadStatusUtil::SetThreadOperation(
+          ThreadStatus::OperationType::OP_MULTIGETENTITY);
       MaybeTestMultiGetEntityLazy(thread, read_opts_copy, cfh, num_keys,
                                   key_slices.data(), &eager_refs);
     }