@nxtedition/rocksdb 8.2.8 → 9.0.0

package/deps/rocksdb/rocksdb/utilities/transactions/optimistic_transaction_test.cc

@@ -3,8 +3,9 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
-
+#include <cstdint>
 #include <functional>
+#include <memory>
 #include <string>
 #include <thread>
 
@@ -27,49 +28,54 @@ class OptimisticTransactionTest
     : public testing::Test,
       public testing::WithParamInterface<OccValidationPolicy> {
  public:
-  OptimisticTransactionDB* txn_db;
+  std::unique_ptr<OptimisticTransactionDB> txn_db;
   std::string dbname;
   Options options;
+  OptimisticTransactionDBOptions occ_opts;
 
   OptimisticTransactionTest() {
     options.create_if_missing = true;
     options.max_write_buffer_number = 2;
     options.max_write_buffer_size_to_maintain = 2 * Arena::kInlineSize;
     options.merge_operator.reset(new TestPutOperator());
+    occ_opts.validate_policy = GetParam();
     dbname = test::PerThreadDBPath("optimistic_transaction_testdb");
 
     EXPECT_OK(DestroyDB(dbname, options));
     Open();
   }
   ~OptimisticTransactionTest() override {
-    delete txn_db;
+    EXPECT_OK(txn_db->Close());
+    txn_db.reset();
     EXPECT_OK(DestroyDB(dbname, options));
   }
 
   void Reopen() {
-    delete txn_db;
-    txn_db = nullptr;
+    txn_db.reset();
     Open();
   }
 
- private:
-  void Open() {
+  static void OpenImpl(const Options& options,
+                       const OptimisticTransactionDBOptions& occ_opts,
+                       const std::string& dbname,
+                       std::unique_ptr<OptimisticTransactionDB>* txn_db) {
     ColumnFamilyOptions cf_options(options);
-    OptimisticTransactionDBOptions occ_opts;
-    occ_opts.validate_policy = GetParam();
     std::vector<ColumnFamilyDescriptor> column_families;
     std::vector<ColumnFamilyHandle*> handles;
     column_families.push_back(
         ColumnFamilyDescriptor(kDefaultColumnFamilyName, cf_options));
-    Status s =
-        OptimisticTransactionDB::Open(DBOptions(options), occ_opts, dbname,
-                                      column_families, &handles, &txn_db);
-
+    OptimisticTransactionDB* raw_txn_db = nullptr;
+    Status s = OptimisticTransactionDB::Open(
+        options, occ_opts, dbname, column_families, &handles, &raw_txn_db);
     ASSERT_OK(s);
-    ASSERT_NE(txn_db, nullptr);
+    ASSERT_NE(raw_txn_db, nullptr);
+    txn_db->reset(raw_txn_db);
     ASSERT_EQ(handles.size(), 1);
     delete handles[0];
   }
+
+ private:
+  void Open() { OpenImpl(options, occ_opts, dbname, &txn_db); }
 };
 
 TEST_P(OptimisticTransactionTest, SuccessTest) {
@@ -316,17 +322,11 @@ TEST_P(OptimisticTransactionTest, FlushTest) {
   delete txn;
 }
 
-TEST_P(OptimisticTransactionTest, FlushTest2) {
-  WriteOptions write_options;
-  ReadOptions read_options, snapshot_read_options;
+namespace {
+void FlushTest2PopulateTxn(Transaction* txn) {
+  ReadOptions snapshot_read_options;
   std::string value;
 
-  ASSERT_OK(txn_db->Put(write_options, Slice("foo"), Slice("bar")));
-  ASSERT_OK(txn_db->Put(write_options, Slice("foo2"), Slice("bar")));
-
-  Transaction* txn = txn_db->BeginTransaction(write_options);
-  ASSERT_NE(txn, nullptr);
-
   snapshot_read_options.snapshot = txn->GetSnapshot();
 
   ASSERT_OK(txn->GetForUpdate(snapshot_read_options, "foo", &value));
@@ -336,6 +336,21 @@ TEST_P(OptimisticTransactionTest, FlushTest2) {
 
   ASSERT_OK(txn->GetForUpdate(snapshot_read_options, "foo", &value));
   ASSERT_EQ(value, "bar2");
+}
+}  // namespace
+
+TEST_P(OptimisticTransactionTest, FlushTest2) {
+  WriteOptions write_options;
+  ReadOptions read_options;
+  std::string value;
+
+  ASSERT_OK(txn_db->Put(write_options, Slice("foo"), Slice("bar")));
+  ASSERT_OK(txn_db->Put(write_options, Slice("foo2"), Slice("bar")));
+
+  Transaction* txn = txn_db->BeginTransaction(write_options);
+  ASSERT_NE(txn, nullptr);
+
+  FlushTest2PopulateTxn(txn);
 
   // Put a random key so we have a MemTable to flush
   ASSERT_OK(txn_db->Put(write_options, "dummy", "dummy"));
@@ -361,9 +376,23 @@ TEST_P(OptimisticTransactionTest, FlushTest2) {
   // txn should not commit since MemTableList History is not large enough
   ASSERT_TRUE(s.IsTryAgain());
 
+  // simply trying Commit again doesn't help
+  s = txn->Commit();
+  ASSERT_TRUE(s.IsTryAgain());
+
   ASSERT_OK(txn_db->Get(read_options, "foo", &value));
   ASSERT_EQ(value, "bar");
 
+  // But rolling back and redoing does
+  ASSERT_OK(txn->Rollback());
+
+  FlushTest2PopulateTxn(txn);
+
+  ASSERT_OK(txn->Commit());
+
+  ASSERT_OK(txn_db->Get(read_options, "foo", &value));
+  ASSERT_EQ(value, "bar2");
+
   delete txn;
 }
 
@@ -616,8 +645,11 @@ TEST_P(OptimisticTransactionTest, ColumnFamiliesTest) {
 
   delete cfa;
   delete cfb;
-  delete txn_db;
-  txn_db = nullptr;
+  txn_db.reset();
+
+  OptimisticTransactionDBOptions my_occ_opts = occ_opts;
+  const size_t bucket_count = 500;
+  my_occ_opts.shared_lock_buckets = MakeSharedOccLockBuckets(bucket_count);
 
   // open DB with three column families
   std::vector<ColumnFamilyDescriptor> column_families;
@@ -630,10 +662,11 @@ TEST_P(OptimisticTransactionTest, ColumnFamiliesTest) {
   column_families.push_back(
       ColumnFamilyDescriptor("CFB", ColumnFamilyOptions()));
   std::vector<ColumnFamilyHandle*> handles;
-  ASSERT_OK(OptimisticTransactionDB::Open(options, dbname, column_families,
-                                          &handles, &txn_db));
-  assert(txn_db != nullptr);
-  ASSERT_NE(txn_db, nullptr);
+  OptimisticTransactionDB* raw_txn_db = nullptr;
+  ASSERT_OK(OptimisticTransactionDB::Open(
+      options, my_occ_opts, dbname, column_families, &handles, &raw_txn_db));
+  ASSERT_NE(raw_txn_db, nullptr);
+  txn_db.reset(raw_txn_db);
 
   Transaction* txn = txn_db->BeginTransaction(write_options);
   ASSERT_NE(txn, nullptr);
@@ -671,6 +704,7 @@ TEST_P(OptimisticTransactionTest, ColumnFamiliesTest) {
   s = txn_db->Get(read_options, "AAA", &value);
   ASSERT_TRUE(s.IsNotFound());
   s = txn_db->Get(read_options, handles[2], "AAAZZZ", &value);
+  ASSERT_OK(s);
   ASSERT_EQ(value, "barbar");
 
   Slice key_slices[3] = {Slice("AAA"), Slice("ZZ"), Slice("Z")};
@@ -694,6 +728,7 @@ TEST_P(OptimisticTransactionTest, ColumnFamiliesTest) {
   delete txn;
   delete txn2;
 
+  // ** MultiGet **
   txn = txn_db->BeginTransaction(write_options, txn_options);
   snapshot_read_options.snapshot = txn->GetSnapshot();
 
@@ -745,11 +780,162 @@ TEST_P(OptimisticTransactionTest, ColumnFamiliesTest) {
   s = txn2->Commit();
   ASSERT_TRUE(s.IsBusy());
 
+  delete txn;
+  delete txn2;
+
+  // ** Test independence and/or sharing of lock buckets across CFs and DBs **
+  if (my_occ_opts.validate_policy == OccValidationPolicy::kValidateParallel) {
+    struct SeenStat {
+      uint64_t rolling_hash = 0;
+      uintptr_t min = 0;
+      uintptr_t max = 0;
+    };
+    SeenStat cur_seen;
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
+        "OptimisticTransaction::CommitWithParallelValidate::lock_bucket_ptr",
+        [&](void* arg) {
+          // Hash the pointer
+          cur_seen.rolling_hash = Hash64(reinterpret_cast<char*>(&arg),
+                                         sizeof(arg), cur_seen.rolling_hash);
+          uintptr_t val = reinterpret_cast<uintptr_t>(arg);
+          if (cur_seen.min == 0 || val < cur_seen.min) {
+            cur_seen.min = val;
+          }
+          if (cur_seen.max == 0 || val > cur_seen.max) {
+            cur_seen.max = val;
+          }
+        });
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
+
+    // Another db sharing lock buckets
+    auto shared_dbname =
+        test::PerThreadDBPath("optimistic_transaction_testdb_shared");
+    std::unique_ptr<OptimisticTransactionDB> shared_txn_db = nullptr;
+    OpenImpl(options, my_occ_opts, shared_dbname, &shared_txn_db);
+
+    // Another db not sharing lock buckets
+    auto nonshared_dbname =
+        test::PerThreadDBPath("optimistic_transaction_testdb_nonshared");
+    std::unique_ptr<OptimisticTransactionDB> nonshared_txn_db = nullptr;
+    my_occ_opts.occ_lock_buckets = bucket_count;
+    my_occ_opts.shared_lock_buckets = nullptr;
+    OpenImpl(options, my_occ_opts, nonshared_dbname, &nonshared_txn_db);
+
+    // Plenty of keys to avoid randomly hitting the same hash sequence
+    std::array<std::string, 30> keys;
+    for (size_t i = 0; i < keys.size(); ++i) {
+      keys[i] = std::to_string(i);
+    }
+
+    // Get a baseline pattern of bucket accesses
+    cur_seen = {};
+    txn = txn_db->BeginTransaction(write_options, txn_options);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(handles[0], key, "blah"));
+    }
+    ASSERT_OK(txn->Commit());
+    // Sufficiently large hash coverage of the space
+    const uintptr_t min_span_bytes = sizeof(port::Mutex) * bucket_count / 2;
+    ASSERT_GT(cur_seen.max - cur_seen.min, min_span_bytes);
+    // Save
+    SeenStat base_seen = cur_seen;
+
+    // Verify it is repeatable
+    cur_seen = {};
+    txn = txn_db->BeginTransaction(write_options, txn_options, txn);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(handles[0], key, "moo"));
+    }
+    ASSERT_OK(txn->Commit());
+    ASSERT_EQ(cur_seen.rolling_hash, base_seen.rolling_hash);
+    ASSERT_EQ(cur_seen.min, base_seen.min);
+    ASSERT_EQ(cur_seen.max, base_seen.max);
+
+    // Try another CF
+    cur_seen = {};
+    txn = txn_db->BeginTransaction(write_options, txn_options, txn);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(handles[1], key, "blah"));
+    }
+    ASSERT_OK(txn->Commit());
+    // Different access pattern (different hash seed)
+    ASSERT_NE(cur_seen.rolling_hash, base_seen.rolling_hash);
+    // Same pointer space
+    ASSERT_LT(cur_seen.min, base_seen.max);
+    ASSERT_GT(cur_seen.max, base_seen.min);
+    // Sufficiently large hash coverage of the space
+    ASSERT_GT(cur_seen.max - cur_seen.min, min_span_bytes);
+    // Save
+    SeenStat cf1_seen = cur_seen;
+
+    // And another CF
+    cur_seen = {};
+    txn = txn_db->BeginTransaction(write_options, txn_options, txn);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(handles[2], key, "blah"));
+    }
+    ASSERT_OK(txn->Commit());
+    // Different access pattern (different hash seed)
+    ASSERT_NE(cur_seen.rolling_hash, base_seen.rolling_hash);
+    ASSERT_NE(cur_seen.rolling_hash, cf1_seen.rolling_hash);
+    // Same pointer space
+    ASSERT_LT(cur_seen.min, base_seen.max);
+    ASSERT_GT(cur_seen.max, base_seen.min);
+    // Sufficiently large hash coverage of the space
+    ASSERT_GT(cur_seen.max - cur_seen.min, min_span_bytes);
+
+    // And DB with shared lock buckets
+    cur_seen = {};
+    delete txn;
+    txn = shared_txn_db->BeginTransaction(write_options, txn_options);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(key, "blah"));
+    }
+    ASSERT_OK(txn->Commit());
+    // Different access pattern (different hash seed)
+    ASSERT_NE(cur_seen.rolling_hash, base_seen.rolling_hash);
+    ASSERT_NE(cur_seen.rolling_hash, cf1_seen.rolling_hash);
+    // Same pointer space
+    ASSERT_LT(cur_seen.min, base_seen.max);
+    ASSERT_GT(cur_seen.max, base_seen.min);
+    // Sufficiently large hash coverage of the space
+    ASSERT_GT(cur_seen.max - cur_seen.min, min_span_bytes);
+
+    // And DB with distinct lock buckets
+    cur_seen = {};
+    delete txn;
+    txn = nonshared_txn_db->BeginTransaction(write_options, txn_options);
+    for (const auto& key : keys) {
+      ASSERT_OK(txn->Put(key, "blah"));
+    }
+    ASSERT_OK(txn->Commit());
+    // Different access pattern (different hash seed)
+    ASSERT_NE(cur_seen.rolling_hash, base_seen.rolling_hash);
+    ASSERT_NE(cur_seen.rolling_hash, cf1_seen.rolling_hash);
+    // Different pointer space
+    ASSERT_TRUE(cur_seen.min > base_seen.max || cur_seen.max < base_seen.min);
+    // Sufficiently large hash coverage of the space
+    ASSERT_GT(cur_seen.max - cur_seen.min, min_span_bytes);
+
+    delete txn;
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
+  }
+
+  // ** Test dropping column family before committing, or even creating txn **
+  txn = txn_db->BeginTransaction(write_options, txn_options);
+  ASSERT_OK(txn->Delete(handles[1], "AAA"));
+
   s = txn_db->DropColumnFamily(handles[1]);
   ASSERT_OK(s);
   s = txn_db->DropColumnFamily(handles[2]);
   ASSERT_OK(s);
 
+  ASSERT_NOK(txn->Commit());
+
+  txn2 = txn_db->BeginTransaction(write_options, txn_options);
+  ASSERT_OK(txn2->Delete(handles[2], "AAA"));
+  ASSERT_NOK(txn2->Commit());
+
   delete txn;
   delete txn2;
 
@@ -1237,7 +1423,7 @@ TEST_P(OptimisticTransactionTest, UndoGetForUpdateTest) {
   txn1->UndoGetForUpdate("A");
 
   Transaction* txn2 = txn_db->BeginTransaction(write_options);
-  txn2->Put("A", "x");
+  ASSERT_OK(txn2->Put("A", "x"));
   ASSERT_OK(txn2->Commit());
   delete txn2;
 
@@ -1402,7 +1588,7 @@ TEST_P(OptimisticTransactionTest, OptimisticTransactionStressTest) {
 
   std::function<void()> call_inserter = [&] {
     ASSERT_OK(OptimisticTransactionStressTestInserter(
-        txn_db, num_transactions_per_thread, num_sets, num_keys_per_set));
+        txn_db.get(), num_transactions_per_thread, num_sets, num_keys_per_set));
   };
 
   // Create N threads that use RandomTransactionInserter to write
@@ -1417,7 +1603,7 @@ TEST_P(OptimisticTransactionTest, OptimisticTransactionStressTest) {
   }
 
   // Verify that data is consistent
-  Status s = RandomTransactionInserter::Verify(txn_db, num_sets);
+  Status s = RandomTransactionInserter::Verify(txn_db.get(), num_sets);
   ASSERT_OK(s);
 }
 
@@ -1469,6 +1655,19 @@ INSTANTIATE_TEST_CASE_P(
     testing::Values(OccValidationPolicy::kValidateSerial,
                     OccValidationPolicy::kValidateParallel));
 
+TEST(OccLockBucketsTest, CacheAligned) {
+  // Typical x86_64 is 40 byte mutex, 64 byte cache line
+  if (sizeof(port::Mutex) >= sizeof(CacheAlignedWrapper<port::Mutex>)) {
+    ROCKSDB_GTEST_BYPASS("Test requires mutex smaller than cache line");
+    return;
+  }
+  auto buckets_unaligned = MakeSharedOccLockBuckets(100, false);
+  auto buckets_aligned = MakeSharedOccLockBuckets(100, true);
+  // Save at least one byte per bucket
+  ASSERT_LE(buckets_unaligned->ApproximateMemoryUsage() + 100,
+            buckets_aligned->ApproximateMemoryUsage());
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {
@@ -1476,4 +1675,3 @@ int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();
 }
-

package/deps/rocksdb/rocksdb/utilities/transactions/pessimistic_transaction.h

@@ -308,4 +308,3 @@ class WriteCommittedTxn : public PessimisticTransaction {
 };
 
 }  // namespace ROCKSDB_NAMESPACE
-

package/deps/rocksdb/rocksdb/utilities/transactions/transaction_base.cc

@@ -110,6 +110,8 @@ void TransactionBaseImpl::Reinitialize(DB* db,
   start_time_ = dbimpl_->GetSystemClock()->NowMicros();
   indexing_enabled_ = true;
   cmp_ = GetColumnFamilyUserComparator(db_->DefaultColumnFamily());
+  WriteBatchInternal::SetDefaultColumnFamilyTimestampSize(
+      write_batch_.GetWriteBatch(), cmp_->timestamp_size());
   WriteBatchInternal::UpdateProtectionInfo(
       write_batch_.GetWriteBatch(), write_options_.protection_bytes_per_key)
       .PermitUncheckedError();
@@ -232,27 +234,56 @@ Status TransactionBaseImpl::PopSavePoint() {
   return write_batch_.PopSavePoint();
 }
 
-Status TransactionBaseImpl::Get(const ReadOptions& read_options,
+Status TransactionBaseImpl::Get(const ReadOptions& _read_options,
                                 ColumnFamilyHandle* column_family,
                                 const Slice& key, std::string* value) {
-  if (read_options.io_activity != Env::IOActivity::kUnknown) {
+  if (_read_options.io_activity != Env::IOActivity::kUnknown &&
+      _read_options.io_activity != Env::IOActivity::kGet) {
     return Status::InvalidArgument(
-        "Cannot call Get with `ReadOptions::io_activity` != "
-        "`Env::IOActivity::kUnknown`");
+        "Can only call Get with `ReadOptions::io_activity` is "
+        "`Env::IOActivity::kUnknown` or `Env::IOActivity::kGet`");
   }
+  ReadOptions read_options(_read_options);
+  if (read_options.io_activity == Env::IOActivity::kUnknown) {
+    read_options.io_activity = Env::IOActivity::kGet;
+  }
+  auto s = GetImpl(read_options, column_family, key, value);
+  return s;
+}
+
+Status TransactionBaseImpl::GetImpl(const ReadOptions& read_options,
+                                    ColumnFamilyHandle* column_family,
+                                    const Slice& key, std::string* value) {
   assert(value != nullptr);
   PinnableSlice pinnable_val(value);
   assert(!pinnable_val.IsPinned());
-  auto s = Get(read_options, column_family, key, &pinnable_val);
+  auto s = GetImpl(read_options, column_family, key, &pinnable_val);
   if (s.ok() && pinnable_val.IsPinned()) {
     value->assign(pinnable_val.data(), pinnable_val.size());
   }  // else value is already assigned
   return s;
 }
 
-Status TransactionBaseImpl::Get(const ReadOptions& read_options,
+Status TransactionBaseImpl::Get(const ReadOptions& _read_options,
                                 ColumnFamilyHandle* column_family,
                                 const Slice& key, PinnableSlice* pinnable_val) {
+  if (_read_options.io_activity != Env::IOActivity::kUnknown &&
+      _read_options.io_activity != Env::IOActivity::kGet) {
+    return Status::InvalidArgument(
+        "Can only call Get with `ReadOptions::io_activity` is "
+        "`Env::IOActivity::kUnknown` or `Env::IOActivity::kGet`");
+  }
+  ReadOptions read_options(_read_options);
+  if (read_options.io_activity == Env::IOActivity::kUnknown) {
+    read_options.io_activity = Env::IOActivity::kGet;
+  }
+  return GetImpl(read_options, column_family, key, pinnable_val);
+}
+
+Status TransactionBaseImpl::GetImpl(const ReadOptions& read_options,
+                                    ColumnFamilyHandle* column_family,
+                                    const Slice& key,
+                                    PinnableSlice* pinnable_val) {
   return write_batch_.GetFromBatchAndDB(db_, read_options, column_family, key,
                                         pinnable_val);
 }
@@ -279,7 +310,7 @@ Status TransactionBaseImpl::GetForUpdate(const ReadOptions& read_options,
     assert(value != nullptr);
     PinnableSlice pinnable_val(value);
     assert(!pinnable_val.IsPinned());
-    s = Get(read_options, column_family, key, &pinnable_val);
+    s = GetImpl(read_options, column_family, key, &pinnable_val);
     if (s.ok() && pinnable_val.IsPinned()) {
       value->assign(pinnable_val.data(), pinnable_val.size());
     }  // else value is already assigned
@@ -307,39 +338,63 @@ Status TransactionBaseImpl::GetForUpdate(const ReadOptions& read_options,
       TryLock(column_family, key, true /* read_only */, exclusive, do_validate);
 
   if (s.ok() && pinnable_val != nullptr) {
-    s = Get(read_options, column_family, key, pinnable_val);
+    s = GetImpl(read_options, column_family, key, pinnable_val);
   }
   return s;
 }
 
 std::vector<Status> TransactionBaseImpl::MultiGet(
-    const ReadOptions& read_options,
+    const ReadOptions& _read_options,
     const std::vector<ColumnFamilyHandle*>& column_family,
     const std::vector<Slice>& keys, std::vector<std::string>* values) {
   size_t num_keys = keys.size();
-  if (read_options.io_activity != Env::IOActivity::kUnknown) {
+  std::vector<Status> stat_list(num_keys);
+  if (_read_options.io_activity != Env::IOActivity::kUnknown &&
+      _read_options.io_activity != Env::IOActivity::kMultiGet) {
     Status s = Status::InvalidArgument(
-        "Cannot call MultiGet with `ReadOptions::io_activity` != "
-        "`Env::IOActivity::kUnknown`");
-    return std::vector<Status>(num_keys, s);
+        "Can only call MultiGet with `ReadOptions::io_activity` is "
+        "`Env::IOActivity::kUnknown` or `Env::IOActivity::kMultiGet`");
+
+    for (size_t i = 0; i < num_keys; ++i) {
+      stat_list[i] = s;
+    }
+    return stat_list;
+  }
+  ReadOptions read_options(_read_options);
+  if (read_options.io_activity == Env::IOActivity::kUnknown) {
+    read_options.io_activity = Env::IOActivity::kMultiGet;
   }
 
   values->resize(num_keys);
-
-  std::vector<Status> stat_list(num_keys);
   for (size_t i = 0; i < num_keys; ++i) {
-    stat_list[i] = Get(read_options, column_family[i], keys[i], &(*values)[i]);
+    stat_list[i] =
+        GetImpl(read_options, column_family[i], keys[i], &(*values)[i]);
   }
 
   return stat_list;
 }
 
-void TransactionBaseImpl::MultiGet(const ReadOptions& read_options,
+void TransactionBaseImpl::MultiGet(const ReadOptions& _read_options,
                                    ColumnFamilyHandle* column_family,
                                    const size_t num_keys, const Slice* keys,
                                    PinnableSlice* values, Status* statuses,
                                    const bool sorted_input) {
-  assert(read_options.io_activity == Env::IOActivity::kUnknown);
+  if (_read_options.io_activity != Env::IOActivity::kUnknown &&
+      _read_options.io_activity != Env::IOActivity::kMultiGet) {
+    Status s = Status::InvalidArgument(
+        "Can only call MultiGet with `ReadOptions::io_activity` is "
+        "`Env::IOActivity::kUnknown` or `Env::IOActivity::kMultiGet`");
+    for (size_t i = 0; i < num_keys; ++i) {
+      if (statuses[i].ok()) {
+        statuses[i] = s;
+      }
+    }
+    return;
+  }
+  ReadOptions read_options(_read_options);
+  if (read_options.io_activity == Env::IOActivity::kUnknown) {
+    read_options.io_activity = Env::IOActivity::kMultiGet;
+  }
   write_batch_.MultiGetFromBatchAndDB(db_, read_options, column_family,
                                       num_keys, keys, values, statuses,
                                       sorted_input);
@@ -349,7 +404,6 @@ std::vector<Status> TransactionBaseImpl::MultiGetForUpdate(
     const ReadOptions& read_options,
     const std::vector<ColumnFamilyHandle*>& column_family,
     const std::vector<Slice>& keys, std::vector<std::string>* values) {
-  // Regardless of whether the MultiGet succeeded, track these keys.
   size_t num_keys = keys.size();
   if (read_options.io_activity != Env::IOActivity::kUnknown) {
     Status s = Status::InvalidArgument(
@@ -357,6 +411,7 @@ std::vector<Status> TransactionBaseImpl::MultiGetForUpdate(
         "`Env::IOActivity::kUnknown`");
     return std::vector<Status>(num_keys, s);
   }
+  // Regardless of whether the MultiGet succeeded, track these keys.
   values->resize(num_keys);
 
   // Lock all keys
@@ -372,7 +427,8 @@ std::vector<Status> TransactionBaseImpl::MultiGetForUpdate(
   // TODO(agiardullo): optimize multiget?
   std::vector<Status> stat_list(num_keys);
   for (size_t i = 0; i < num_keys; ++i) {
-    stat_list[i] = Get(read_options, column_family[i], keys[i], &(*values)[i]);
+    stat_list[i] =
+        GetImpl(read_options, column_family[i], keys[i], &(*values)[i]);
   }
 
   return stat_list;

package/deps/rocksdb/rocksdb/utilities/transactions/transaction_base.h

@@ -53,11 +53,13 @@ class TransactionBaseImpl : public Transaction {
   Status PopSavePoint() override;
 
   using Transaction::Get;
-  Status Get(const ReadOptions& options, ColumnFamilyHandle* column_family,
-             const Slice& key, std::string* value) override;
+  Status Get(const ReadOptions& _read_options,
+             ColumnFamilyHandle* column_family, const Slice& key,
+             std::string* value) override;
 
-  Status Get(const ReadOptions& options, ColumnFamilyHandle* column_family,
-             const Slice& key, PinnableSlice* value) override;
+  Status Get(const ReadOptions& _read_options,
+             ColumnFamilyHandle* column_family, const Slice& key,
+             PinnableSlice* value) override;
 
   Status Get(const ReadOptions& options, const Slice& key,
              std::string* value) override {
@@ -84,7 +86,7 @@ class TransactionBaseImpl : public Transaction {
 
   using Transaction::MultiGet;
   std::vector<Status> MultiGet(
-      const ReadOptions& options,
+      const ReadOptions& _read_options,
       const std::vector<ColumnFamilyHandle*>& column_family,
       const std::vector<Slice>& keys,
       std::vector<std::string>* values) override;
@@ -98,9 +100,10 @@ class TransactionBaseImpl : public Transaction {
                     keys, values);
   }
 
-  void MultiGet(const ReadOptions& options, ColumnFamilyHandle* column_family,
-                const size_t num_keys, const Slice* keys, PinnableSlice* values,
-                Status* statuses, const bool sorted_input = false) override;
+  void MultiGet(const ReadOptions& _read_options,
+                ColumnFamilyHandle* column_family, const size_t num_keys,
+                const Slice* keys, PinnableSlice* values, Status* statuses,
+                const bool sorted_input = false) override;
 
   using Transaction::MultiGetForUpdate;
   std::vector<Status> MultiGetForUpdate(
@@ -260,6 +263,13 @@ class TransactionBaseImpl : public Transaction {
   LockTracker& GetTrackedLocks() { return *tracked_locks_; }
 
  protected:
+  Status GetImpl(const ReadOptions& options, ColumnFamilyHandle* column_family,
+                 const Slice& key, std::string* value) override;
+
+  virtual Status GetImpl(const ReadOptions& options,
+                         ColumnFamilyHandle* column_family, const Slice& key,
+                         PinnableSlice* value) override;
+
   // Add a key to the list of tracked keys.
   //
   // seqno is the earliest seqno this key was involved with this transaction.
@@ -379,4 +389,3 @@ class TransactionBaseImpl : public Transaction {
 };
 
 }  // namespace ROCKSDB_NAMESPACE
-