Make concurrent Memtable stats more robust (#14506)

Summary:
Fix thread-safety issues in memtable stats tracking and `MaybeUpdateNewestUDT` to prepare for PR https://github.com/facebook/rocksdb/issues/14448, which introduces concurrent range tombstone insertion from the read path into the mutable memtable. Currently the non-concurrent write path updates memtable counters (`num_entries_`, `num_deletes_`, `num_range_deletes_`, `data_size_`) using non-atomic load/store pairs. While this is safe today PR https://github.com/facebook/rocksdb/issues/14448 will make range_tombstone memtable concurrent, but the main write path can still be non-concurrent. This fix ensures stats are tracked correctly.

Similarly, `MaybeUpdateNewestUDT` was not thread-safe and was not called in the concurrent write path at all. This PR also fixes the `post_process_info` delete counter which missed `kTypeSingleDeletion` and `kTypeDeletionWithTimestamp`.

## Key changes
- Switch non-concurrent counter updates from `LoadRelaxed`/`StoreRelaxed` pairs to `FetchAddRelaxed` to avoid races with concurrent range tombstone inserts (e.g. `AddLogicallyRedundantRangeTombstone` calling `BatchPostProcess`).
- Make `MaybeUpdateNewestUDT` thread-safe by replacing `Slice newest_udt_` with `RelaxedAtomic<const char*> newest_udt_data_` and using a CAS loop. The pointed-to memory lives in the arena and remains valid for the memtable's lifetime.
- Call `MaybeUpdateNewestUDT` in the concurrent write path (was previously skipped with a TODO).
- Fix `post_process_info->num_deletes++` to also count `kTypeSingleDeletion` and `kTypeDeletionWithTimestamp`, matching the non-concurrent path.
- Change `GetNewestUDT()` return type from `const Slice&` to `Slice` (returning by value) to avoid dangling reference issues with the new atomic pointer storage. Updated across `MemTable`, `ReadOnlyMemTable`, `MemTableListVersion`, and `WBWIMemTable`.
- Remove unused `Slice newest_udt_` member from `WBWIMemTable`.

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

Test Plan:
- Added `ConcurrentWriteMemTableProperties` test in `db_properties_test.cc`: 4 threads writing puts, deletes, and single deletes concurrently, then verifying `rocksdb.num-entries-active-mem-table` and `rocksdb.num-deletes-active-mem-table` properties are correct. Flush with `flush_verify_memtable_count=true` validates integrity.
- Added `ConcurrentGetTableNewestUDT` test in `memtable_list_test.cc`: 4 threads concurrently inserting entries with UDTs via `allow_concurrent=true`, verifying the newest UDT is correctly tracked.
- Extended `DuplicateSeq` and `ConcurrentMergeWrite` tests in `db_memtable_test.cc` to verify stat counters after `BatchPostProcess`.
- Benchmark results show no performance regression (3 runs each, averaged):

**Workload 1: fillrandom** (pure Put workload)
```
./db_bench -benchmarks=fillrandom -seed=1 -compression_type=none -threads=8 -db=<DB>
```

**Workload 2: readrandomwriterandom** (mixed read/write/delete)
```
# Setup: fillrandom,compact to populate DB (single-threaded)
./db_bench -benchmarks=readrandomwriterandom -seed=1 -compression_type=none \
  -use_existing_db=1 -readwritepercent=50 -deletepercent=20 -threads=8 -db=<DB>
```

**Workload 3: fillrandom with range tombstones** (puts interleaved with range deletes)
```
./db_bench -benchmarks=fillrandom -seed=1 -compression_type=none \
  -writes_per_range_tombstone=100 -range_tombstone_width=10 \
  -max_num_range_tombstones=1000 -threads=8 -db=<DB>
```

```
| Benchmark                          | avg ops/s (main) | avg ops/s (feature) | % change |
|------------------------------------|-----------------|--------------------:|----------|
| fillrandom                         |          531,986 |             532,099 |   +0.02% |
| readrandomwriterandom (50r/30w/20d)|          786,400 |             815,143 |   +3.65% |
| fillrandom (range tombstones)      |          541,656 |             531,008 |   -1.97% |
```

Reviewed By: xingbowang

Differential Revision: D97974456

Pulled By: joshkang97

fbshipit-source-id: ea95f23953fe37771a599aed61ece1a504b12c2e

db/db_impl/db_impl_write.cc

@@ -2540,7 +2540,7 @@ Status DBImpl::SwitchMemtable(ColumnFamilyData* cfd, WriteContext* context,
   memtable_info.num_deletes = cfd->mem()->NumDeletion();
   if (!cfd->ioptions().persist_user_defined_timestamps &&
       cfd->user_comparator()->timestamp_size() > 0) {
-    const Slice& newest_udt = cfd->mem()->GetNewestUDT();
+    Slice newest_udt = cfd->mem()->GetNewestUDT();
     memtable_info.newest_udt.assign(newest_udt.data(), newest_udt.size());
   }
   // Log this later after lock release. It may be outdated, e.g., if background

db/db_memtable_test.cc

@@ -163,6 +163,10 @@ TEST_F(DBMemTableTest, DuplicateSeq) {
       mem->Add(seq, kTypeSingleDeletion, "key", "", nullptr /* kv_prot_info */)
           .IsTryAgain());
 
+  ASSERT_EQ(mem->NumEntries(), 2u);
+  ASSERT_EQ(mem->NumDeletion(), 0u);
+  ASSERT_EQ(mem->NumRangeDeletion(), 0u);
+
   // Test the duplicate keys under stress
   for (int i = 0; i < 10000; i++) {
     bool insert_dup = i % 10 == 1;
@@ -204,6 +208,11 @@ TEST_F(DBMemTableTest, DuplicateSeq) {
   ASSERT_TRUE(mem->Add(seq, kTypeValue, "key", "value",
                        nullptr /* kv_prot_info */, true, &post_process_info)
                   .IsTryAgain());
+
+  mem->BatchPostProcess(post_process_info);
+  ASSERT_EQ(mem->NumEntries(), 1u);
+  ASSERT_EQ(mem->NumDeletion(), 0u);
+  ASSERT_EQ(mem->NumRangeDeletion(), 0u);
   delete mem;
 }
 
@@ -242,6 +251,7 @@ TEST_F(DBMemTableTest, ConcurrentMergeWrite) {
                          true, &post_process_info1));
       v1.clear();
     }
+    mem->BatchPostProcess(post_process_info1);
   });
   ROCKSDB_NAMESPACE::port::Thread write_thread2([&]() {
     MemTablePostProcessInfo post_process_info2;
@@ -252,9 +262,14 @@ TEST_F(DBMemTableTest, ConcurrentMergeWrite) {
                          true, &post_process_info2));
       v2.clear();
     }
+    mem->BatchPostProcess(post_process_info2);
   });
   write_thread1.join();
   write_thread2.join();
+  // 1 non-concurrent Put + (num_ops - 1) concurrent Merges
+  ASSERT_EQ(mem->NumEntries(), static_cast<uint64_t>(num_ops));
+  ASSERT_EQ(mem->NumDeletion(), 0u);
+  ASSERT_EQ(mem->NumRangeDeletion(), 0u);
 
   Status status;
   ReadOptions roptions;

db/db_properties_test.cc

@@ -772,6 +772,72 @@ TEST_F(DBPropertiesTest, NumImmutableMemTable) {
   } while (ChangeCompactOptions());
 }
 
+TEST_F(DBPropertiesTest, ConcurrentWriteMemTableProperties) {
+  Options options = CurrentOptions();
+  options.allow_concurrent_memtable_write = true;
+  options.memtable_factory = std::make_shared<SkipListFactory>();
+  options.flush_verify_memtable_count = true;
+  DestroyAndReopen(options);
+
+  // Multi-threaded writes that go through the concurrent Add() +
+  // BatchPostProcess path.
+  const int kNumThreads = 4;
+  const int kEntriesPerThread = 100;
+  std::vector<port::Thread> threads;
+  threads.reserve(kNumThreads);
+  for (int t = 0; t < kNumThreads; t++) {
+    threads.emplace_back([&, t]() {
+      for (int i = 0; i < kEntriesPerThread; i++) {
+        std::string key = "k_" + std::to_string(t) + "_" + std::to_string(i);
+        ASSERT_OK(Put(key, "val"));
+      }
+    });
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+
+  uint64_t num_entries;
+  ASSERT_TRUE(dbfull()->GetIntProperty("rocksdb.num-entries-active-mem-table",
+                                       &num_entries));
+  ASSERT_EQ(num_entries,
+            static_cast<uint64_t>(kNumThreads * kEntriesPerThread));
+
+  // Write deletes and single deletes via concurrent path
+  threads.clear();
+  for (int t = 0; t < kNumThreads; t++) {
+    threads.emplace_back([&, t]() {
+      for (int i = 0; i < kEntriesPerThread; i++) {
+        std::string key = "d_" + std::to_string(t) + "_" + std::to_string(i);
+        if (i % 2 == 0) {
+          ASSERT_OK(Delete(key));
+        } else {
+          ASSERT_OK(SingleDelete(key));
+        }
+      }
+    });
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+
+  uint64_t total_entries;
+  ASSERT_TRUE(dbfull()->GetIntProperty("rocksdb.num-entries-active-mem-table",
+                                       &total_entries));
+  ASSERT_EQ(total_entries,
+            static_cast<uint64_t>(kNumThreads * kEntriesPerThread * 2));
+
+  uint64_t num_deletes;
+  ASSERT_TRUE(dbfull()->GetIntProperty("rocksdb.num-deletes-active-mem-table",
+                                       &num_deletes));
+  ASSERT_EQ(num_deletes,
+            static_cast<uint64_t>(kNumThreads * kEntriesPerThread));
+
+  // Flush exercises the integrity check (flush_verify_memtable_count = true).
+  // If stat counters are wrong, this would return Corruption.
+  ASSERT_OK(Flush());
+}
+
 // TODO(techdept) : Disabled flaky test #12863555
 TEST_F(DBPropertiesTest, DISABLED_GetProperty) {
   // Set sizes to both background thread pool to be 1 and block them.

db/memtable.cc

@@ -1012,16 +1012,13 @@ Status MemTable::Add(SequenceNumber s, ValueType type,
       }
     }
 
-    // this is a bit ugly, but is the way to avoid locked instructions
-    // when incrementing an atomic
-    num_entries_.StoreRelaxed(num_entries_.LoadRelaxed() + 1);
-    data_size_.StoreRelaxed(data_size_.LoadRelaxed() + encoded_len);
+    num_entries_.FetchAddRelaxed(1);
+    data_size_.FetchAddRelaxed(encoded_len);
     if (type == kTypeDeletion || type == kTypeSingleDeletion ||
         type == kTypeDeletionWithTimestamp) {
-      num_deletes_.StoreRelaxed(num_deletes_.LoadRelaxed() + 1);
+      num_deletes_.FetchAddRelaxed(1);
     } else if (type == kTypeRangeDeletion) {
-      uint64_t val = num_range_deletes_.LoadRelaxed() + 1;
-      num_range_deletes_.StoreRelaxed(val);
+      num_range_deletes_.FetchAddRelaxed(1);
     }
 
     if (bloom_filter_ && prefix_extractor_ &&
@@ -1044,8 +1041,6 @@ Status MemTable::Add(SequenceNumber s, ValueType type,
       assert(first_seqno_.load() >= earliest_seqno_.load());
     }
     assert(post_process_info == nullptr);
-    // TODO(yuzhangyu): support updating newest UDT for when `allow_concurrent`
-    // is true.
     MaybeUpdateNewestUDT(key_slice);
     UpdateFlushState();
   } else {
@@ -1059,7 +1054,8 @@ Status MemTable::Add(SequenceNumber s, ValueType type,
     assert(post_process_info != nullptr);
     post_process_info->num_entries++;
     post_process_info->data_size += encoded_len;
-    if (type == kTypeDeletion) {
+    if (type == kTypeDeletion || type == kTypeSingleDeletion ||
+        type == kTypeDeletionWithTimestamp) {
       post_process_info->num_deletes++;
     }
 
@@ -1083,6 +1079,7 @@ Status MemTable::Add(SequenceNumber s, ValueType type,
         (cur_earliest_seqno == kMaxSequenceNumber || s < cur_earliest_seqno) &&
         !earliest_seqno_.compare_exchange_weak(cur_earliest_seqno, s)) {
     }
+    MaybeUpdateNewestUDT(key_slice);
   }
   if (type == kTypeRangeDeletion) {
     auto new_cache = std::make_shared<FragmentedRangeTombstoneListCache>();
@@ -2006,14 +2003,22 @@ void MemTable::MaybeUpdateNewestUDT(const Slice& user_key) {
   }
   const Comparator* ucmp = GetInternalKeyComparator().user_comparator();
   Slice udt = ExtractTimestampFromUserKey(user_key, ts_sz_);
-  if (newest_udt_.empty() || ucmp->CompareTimestamp(udt, newest_udt_) > 0) {
-    newest_udt_ = udt;
+  const char* cur = newest_udt_data_.Load();
+  while (cur == nullptr ||
+         ucmp->CompareTimestamp(udt, Slice(cur, ts_sz_)) > 0) {
+    if (newest_udt_data_.CasWeak(cur, udt.data())) {
+      break;
+    }
   }
 }
 
-const Slice& MemTable::GetNewestUDT() const {
+Slice MemTable::GetNewestUDT() const {
   assert(ts_sz_ > 0);
-  return newest_udt_;
+  const char* data = newest_udt_data_.Load();
+  if (data == nullptr) {
+    return Slice();
+  }
+  return Slice(data, ts_sz_);
 }
 
 }  // namespace ROCKSDB_NAMESPACE

db/memtable.h

@@ -329,7 +329,7 @@ class ReadOnlyMemTable {
   // `persist_user_defined_timestamps` to false. The tracked newest UDT will be
   // used by flush job in the background to help check the MemTable's
   // eligibility for Flush.
-  virtual const Slice& GetNewestUDT() const = 0;
+  virtual Slice GetNewestUDT() const = 0;
 
   // Increase reference count.
   // REQUIRES: external synchronization to prevent simultaneous
@@ -826,7 +826,7 @@ class MemTable final : public ReadOnlyMemTable {
 
   //  Gets the newest user defined timestamps in the memtable. This should only
   //  be called when user defined timestamp is enabled.
-  const Slice& GetNewestUDT() const override;
+  Slice GetNewestUDT() const override;
 
   // Returns Corruption status if verification fails.
   static Status VerifyEntryChecksum(const char* entry,
@@ -910,11 +910,11 @@ class MemTable final : public ReadOnlyMemTable {
   // Size in bytes for the user-defined timestamps.
   size_t ts_sz_;
 
-  // Newest user-defined timestamp contained in this MemTable. For ts1, and ts2
-  // if Comparator::CompareTimestamp(ts1, ts2) > 0, ts1 is considered newer than
-  // ts2. We track this field for a MemTable if its column family has UDT
-  // feature enabled.
-  Slice newest_udt_;
+  // Pointer to the newest user-defined timestamp data in this MemTable. The
+  // pointed-to memory lives in the arena and remains valid for the lifetime of
+  // the memtable. Stored as an atomic pointer so that concurrent range
+  // tombstone inserts from the read path can safely update it via CAS.
+  Atomic<const char*> newest_udt_data_{nullptr};
 
   // Updates flush_state_ using ShouldFlushNow()
   void UpdateFlushState();

db/memtable_list.cc

@@ -374,17 +374,16 @@ bool MemTableListVersion::TrimHistory(autovector<ReadOnlyMemTable*>* to_delete,
   return ret;
 }
 
-const Slice& MemTableListVersion::GetNewestUDT() const {
-  static Slice kEmptySlice;
+Slice MemTableListVersion::GetNewestUDT() const {
   for (auto it = memlist_.begin(); it != memlist_.end(); ++it) {
     ReadOnlyMemTable* m = *it;
     Slice timestamp = m->GetNewestUDT();
     assert(!timestamp.empty() || m->IsEmpty());
     if (!timestamp.empty()) {
-      return m->GetNewestUDT();
+      return timestamp;
     }
   }
-  return kEmptySlice;
+  return Slice();
 }
 
 // Returns true if there is at least one memtable on which flush has

db/memtable_list.h

@@ -153,7 +153,7 @@ class MemTableListVersion {
   // This returns the newest user defined timestamp found in the most recent
   // immutable memtable. This should only be called when user defined timestamp
   // is enabled.
-  const Slice& GetNewestUDT() const;
+  Slice GetNewestUDT() const;
 
  private:
   friend class MemTableList;

db/memtable_list_test.cc

@@ -1153,6 +1153,63 @@ TEST_F(MemTableListWithTimestampTest, GetTableNewestUDT) {
   to_delete.clear();
 }
 
+TEST_F(MemTableListWithTimestampTest, ConcurrentGetTableNewestUDT) {
+  const int num_threads = 4;
+  const int num_entries_per_thread = 100;
+
+  auto factory = std::make_shared<SkipListFactory>();
+  options.memtable_factory = factory;
+  options.persist_user_defined_timestamps = false;
+  ImmutableOptions ioptions(options);
+  const Comparator* ucmp = test::BytewiseComparatorWithU64TsWrapper();
+  InternalKeyComparator cmp(ucmp);
+  WriteBufferManager wb(options.db_write_buffer_size);
+  MutableCFOptions mutable_cf_options(options);
+
+  MemTable* mem = new MemTable(cmp, ioptions, mutable_cf_options, &wb,
+                               kMaxSequenceNumber, 0 /* column_family_id */);
+  mem->Ref();
+
+  std::atomic<uint64_t> next_seq{1};
+  uint64_t max_ts = num_threads * num_entries_per_thread - 1;
+
+  std::vector<port::Thread> threads;
+  threads.reserve(num_threads);
+
+  for (int t = 0; t < num_threads; t++) {
+    threads.emplace_back([&, t]() {
+      MemTablePostProcessInfo post_process_info;
+      for (int j = 0; j < num_entries_per_thread; j++) {
+        uint64_t ts = t * num_entries_per_thread + j;
+        std::string key = "key" + std::to_string(ts);
+        std::string write_ts;
+        PutFixed64(&write_ts, ts);
+        key.append(write_ts);
+
+        SequenceNumber seq = next_seq.fetch_add(1);
+        ASSERT_OK(mem->Add(seq, kTypeValue, key, "val",
+                           nullptr /* kv_prot_info */,
+                           true /* allow_concurrent */, &post_process_info));
+      }
+      mem->BatchPostProcess(post_process_info);
+    });
+  }
+
+  for (auto& thread : threads) {
+    thread.join();
+  }
+
+  ASSERT_EQ(num_threads * num_entries_per_thread, mem->NumEntries());
+
+  Slice newest_udt = mem->GetNewestUDT();
+  ASSERT_EQ(sizeof(uint64_t), newest_udt.size());
+  uint64_t got_ts = DecodeFixed64(newest_udt.data());
+  ASSERT_EQ(max_ts, got_ts);
+
+  ReadOnlyMemTable* m = mem->Unref();
+  delete m;
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

memtable/wbwi_memtable.h

@@ -199,10 +199,10 @@ class WBWIMemTable final : public ReadOnlyMemTable {
     return true;
   }
 
-  const Slice& GetNewestUDT() const override {
+  Slice GetNewestUDT() const override {
     // FIXME: support UDT
     assert(false);
-    return newest_udt_;
+    return Slice();
   }
 
   // Assign a sequence number to the entries in this memtable.
@@ -224,7 +224,6 @@ class WBWIMemTable final : public ReadOnlyMemTable {
  private:
   inline InternalIterator* NewIterator() const;
 
-  Slice newest_udt_;
   std::shared_ptr<WriteBatchWithIndex> wbwi_;
   const Comparator* comparator_;
   InternalKeyComparator ikey_comparator_;