fix(generic): stop RM from freeing a container under a yielding iteration

Author: backport test

src/server/container_utils.cc

@@ -161,6 +161,7 @@ using namespace std;
 bool IterateList(const PrimeValue& pv, const IterateFunc& func, size_t start, size_t end,
                  bool allow_yield) {
   DCHECK_LE(start, end);
+  ScopedValueIteration iteration_guard{pv};
   bool success = true;
   size_t len = pv.Size();
   if (len == 0) {
@@ -211,6 +212,7 @@ bool IterateList(const PrimeValue& pv, const IterateFunc& func, size_t start, si
 }
 
 bool IterateSet(const PrimeValue& pv, const IterateFunc& func, bool allow_yield) {
+  ScopedValueIteration iteration_guard{pv};
   bool success = true;
   if (pv.Encoding() == kEncodingIntSet) {
     intset* is = static_cast<intset*>(pv.RObjPtr());
@@ -242,6 +244,7 @@ bool IterateSet(const PrimeValue& pv, const IterateFunc& func, bool allow_yield)
 
 bool IterateSortedSet(const PrimeValue& pv, const IterateSortedFunc& func, size_t start, size_t end,
                       bool reverse, bool use_score, bool allow_yield) {
+  ScopedValueIteration iteration_guard{pv};
   size_t llen = pv.Size();
   if (llen == 0)
     return true;
@@ -305,6 +308,7 @@ bool IterateSortedSet(const PrimeValue& pv, const IterateSortedFunc& func, size_
 }
 
 bool IterateMap(const PrimeValue& pv, const IterateKVFunc& func, bool allow_yield) {
+  ScopedValueIteration iteration_guard{pv};
   bool finished = true;
 
   if (pv.Encoding() == kEncodingListPack) {

src/server/db_slice.cc

@@ -1462,19 +1462,26 @@ DbSlice::Iterator DbSlice::ExpireIfNeeded(const Context& cntx, Iterator it) cons
   return Iterator::FromPrime(ExpireIfNeeded(cntx, it.GetInnerIt()));
 }
 
+bool DbSlice::WillExpire(const Context& cntx, PrimeIterator it) const {
+  if (!it->first.HasExpire())
+    return false;
+
+  int64_t expire_time = it->first.GetExpireTime();
+
+  // Never do expiration if expiration is disabled, or on replicas unless replica_delete_expired
+  // is enabled (which allows replicas to proactively delete expired keys on the read path).
+  return int64_t(cntx.time_now_ms) >= expire_time && expire_allowed_ &&
+         (!owner_->IsReplica() || absl::GetFlag(FLAGS_replica_delete_expired));
+}
+
 PrimeIterator DbSlice::ExpireIfNeeded(const Context& cntx, PrimeIterator it,
                                       vector<string>* events) const {
   if (!it->first.HasExpire()) {
     LOG(DFATAL) << "Invalid call to ExpireIfNeeded";
     return it;
   }
 
-  int64_t expire_time = it->first.GetExpireTime();
-
-  // Never do expiration if expiration is disabled, or on replicas unless replica_delete_expired
-  // is enabled (which allows replicas to proactively delete expired keys on the read path).
-  if (int64_t(cntx.time_now_ms) < expire_time || !expire_allowed_ ||
-      (owner_->IsReplica() && !absl::GetFlag(FLAGS_replica_delete_expired))) {
+  if (!WillExpire(cntx, it)) {
     return it;
   }
 

src/server/db_slice.h

@@ -421,6 +421,9 @@ class DbSlice {
   // and returns Iterator{}.
   Iterator ExpireIfNeeded(const Context& cntx, Iterator it) const;
 
+  // Whether ExpireIfNeeded() would erase 'it' right now.
+  bool WillExpire(const Context& cntx, PrimeIterator it) const;
+
   // Iterate over all expire table entries and delete expired.
   void ExpireAllIfNeeded();
 

src/server/generic_family.cc

@@ -649,6 +649,11 @@ bool ScanCb(const OpArgs& op_args, PrimeIterator prime_it, const ScanOpts& opts,
 
   DbSlice::Iterator it = DbSlice::Iterator::FromPrime(prime_it);
   if (prime_it->first.HasExpire()) {
+    // Expiring frees the value and this scan holds no key lock, so a parked walk would be left
+    // pointing into it. Only a key that is expiring anyway may be skipped - it is gone either
+    // way, and the heartbeat will free it once the walk ends.
+    if (db_slice.WillExpire(op_args.db_cntx, prime_it) && IsValueIterated(prime_it->second))
+      return false;
     it = db_slice.ExpireIfNeeded(op_args.db_cntx, it);
     if (!IsValid(it))
       return false;
@@ -775,28 +780,82 @@ uint64_t ScanGeneric(uint64_t cursor, const ScanOpts& scan_opts, StringVec* keys
   return cursor;
 }
 
-void OpScanAndDelete(const OpArgs& op_args, const ScanOpts& scan_opts, uint64_t* cursor,
-                     uint32_t* deleted) {
+// A fiber parked inside container_utils::Iterate* holds raw pointers into the value and RM holds
+// no key lock, so freeing it here is a use-after-free. kIterated means the key is still there and
+// the caller must not let the cursor move past it.
+enum class DeleteScannedResult { kDeleted, kMissing, kIterated };
+
+DeleteScannedResult DeleteScannedKey(const OpArgs& op_args, string_view key, uint64_t deadline_ms) {
+  auto& db_slice = op_args.GetDbSlice();
+
+  while (true) {
+    // Raw peek: DbSlice::Find* would lazily expire - that is, free - the very value we protect,
+    // and would re-run the change callbacks on every retry.
+    PrimeIterator peek = db_slice.GetTables(op_args.db_cntx.db_index)->Find(key);
+    if (!IsValid(peek))
+      return DeleteScannedResult::kMissing;
+
+    if (!IsValueIterated(peek->second)) {
+      auto res = db_slice.FindMutable(op_args.db_cntx, key);
+      if (!IsValid(res.it))
+        return DeleteScannedResult::kMissing;
+
+      if (!IsValueIterated(res.it->second)) {  // FindMutable preempts, so re-check before freeing
+        db_slice.DelMutable(op_args.db_cntx, std::move(res));
+        if (op_args.shard->journal()) {
+          RecordDelete(op_args.db_cntx.db_index, key);
+        }
+        return DeleteScannedResult::kDeleted;
+      }
+      res.post_updater.Cancel();  // nothing changed, do not wake its watchers
+    }
+
+    if (GetCurrentTimeMs() > deadline_ms)
+      return DeleteScannedResult::kIterated;
+    // Sleep rather than yield: a yield loop keeps this proactor's ready queue non-empty, so it
+    // stops polling and nothing timer-driven on this shard fires until the wait is over.
+    util::ThisFiber::SleepFor(std::chrono::microseconds(100));
+  }
+}
+
+// Returns false if a key could not be deleted before the deadline. *cursor is then rewound to
+// where this scan started, so the next pass sees that key again instead of losing it behind a
+// cursor OpScan already advanced. Rescanning is idempotent - what we deleted no longer matches.
+bool OpScanAndDelete(const OpArgs& op_args, const ScanOpts& scan_opts, uint64_t* cursor,
+                     uint32_t* deleted, uint64_t deadline_ms) {
+  const uint64_t scan_start = *cursor;
   StringVec keys;
   OpScan(op_args, scan_opts, cursor, &keys);
 
-  auto& db_slice = op_args.GetDbSlice();
   uint32_t count = 0;
+  bool complete = true;
   for (const auto& key : keys) {
-    auto it = db_slice.FindMutable(op_args.db_cntx, key).it;
-    if (!IsValid(it))
-      continue;
-    db_slice.Del(op_args.db_cntx, it);
-    if (op_args.shard->journal()) {
-      RecordDelete(op_args.db_cntx.db_index, key);
+    switch (DeleteScannedKey(op_args, key, deadline_ms)) {
+      case DeleteScannedResult::kDeleted:
+        ++count;
+        break;
+      case DeleteScannedResult::kMissing:
+        break;
+      case DeleteScannedResult::kIterated:
+        complete = false;  // the rest of the batch is still worth deleting
+        break;
     }
-    ++count;
   }
   *deleted += count;
+
+  if (!complete)
+    *cursor = scan_start;
+  return complete;
 }
 
 uint64_t RmGeneric(uint64_t cursor, const ScanOpts& scan_opts, uint32_t* deleted,
                    ConnectionContext* cntx) {
+  // A returned cursor of 0 means the pass is over, so it can not also mean "resume at shard 0,
+  // bucket 0" - which is what we must say when we give up on a key in shard 0's first batch.
+  // A dash token stays below 2^40 and we shift it by 10, so the high bits are ours.
+  constexpr uint64_t kResumeBit = 1ULL << 62;
+  cursor &= ~kResumeBit;
+
   ShardId sid = cursor % 1024;
 
   EngineShardSet* ess = shard_set;
@@ -811,13 +870,15 @@ uint64_t RmGeneric(uint64_t cursor, const ScanOpts& scan_opts, uint32_t* deleted
 
   cursor >>= 10;
   DbContext db_cntx{cntx->ns, cntx->conn_state.db_index, GetCurrentTimeMs()};
+  const uint64_t deadline_ms = db_cntx.time_now_ms + kMaxRmTimeMs;
 
   *deleted = 0;
+  bool complete = true;
 
   do {
     auto cb = [&] {
       OpArgs op_args{EngineShard::tlocal(), nullptr, db_cntx};
-      OpScanAndDelete(op_args, scan_opts, &cursor, deleted);
+      complete = OpScanAndDelete(op_args, scan_opts, &cursor, deleted, deadline_ms);
     };
 
     if (EngineShard::tlocal() && EngineShard::tlocal()->shard_id() == sid) {
@@ -827,6 +888,9 @@ uint64_t RmGeneric(uint64_t cursor, const ScanOpts& scan_opts, uint32_t* deleted
       ess->Await(sid, cb);
     }
 
+    if (!complete)  // cursor was rewound; do not step past a key we failed to delete
+      break;
+
     if (cursor == 0) {
       ++sid;
       if (unsigned(sid) == shard_count)
@@ -841,6 +905,8 @@ uint64_t RmGeneric(uint64_t cursor, const ScanOpts& scan_opts, uint32_t* deleted
 
   if (sid < shard_count) {
     cursor = (cursor << 10) | sid;
+    if (cursor == 0)  // gave up in shard 0's first batch; a bare 0 would read as "done"
+      cursor = kResumeBit;
   } else {
     DCHECK_EQ(0u, cursor);
   }