From 6f9d8260f09edc42cd53a7680b9b67601bf1df2e Mon Sep 17 00:00:00 2001
From: George Reynya <greynya@fb.com>
Date: Wed, 27 Nov 2024 16:45:59 -0800
Subject: [PATCH] Add abort check to yield hook (#13164)

Summary:
Adding ability to kill mysql queries traversing long lists of tombstones. Outside of mysql where RocksDbThreadYieldAndCheckAbort is not implemented all of this should still be optimized out by the compiler.

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

Reviewed By: cbi42

Differential Revision: D66556004

Pulled By: george-reynya

fbshipit-source-id: 727875569209cd6d2f29c07f89ecfa641d5ee36f
---
 db/db_impl/db_impl.cc |  7 ++++++-
 db/db_iter.cc         |  8 +++++++-
 port/port.h           | 14 +++++---------
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc
index 7076f5ff4d0..125ef3515b3 100644
--- a/db/db_impl/db_impl.cc
+++ b/db/db_impl/db_impl.cc
@@ -3247,8 +3247,13 @@ Status DBImpl::MultiGetImpl(
       s = Status::Aborted();
       break;
     }
+
     // This could be a long-running operation
-    ROCKSDB_THREAD_YIELD_HOOK();
+    bool aborted = ROCKSDB_THREAD_YIELD_CHECK_ABORT();
+    if (aborted) {
+      s = Status::Aborted("Query abort.");
+      break;
+    }
   }
 
   // Post processing (decrement reference counts and record statistics)
diff --git a/db/db_iter.cc b/db/db_iter.cc
index 49537f70112..c5a09910365 100644
--- a/db/db_iter.cc
+++ b/db/db_iter.cc
@@ -581,8 +581,14 @@ bool DBIter::FindNextUserEntryInternal(bool skipping_saved_key,
     } else {
       iter_.Next();
     }
+
     // This could be a long-running operation due to tombstones, etc.
-    ROCKSDB_THREAD_YIELD_HOOK();
+    bool aborted = ROCKSDB_THREAD_YIELD_CHECK_ABORT();
+    if (aborted) {
+      valid_ = false;
+      status_ = Status::Aborted("Query abort.");
+      return false;
+    }
   } while (iter_.Valid());
 
   valid_ = false;
diff --git a/port/port.h b/port/port.h
index 141716e5b9f..5afb7929d7b 100644
--- a/port/port.h
+++ b/port/port.h
@@ -24,14 +24,10 @@
 // A temporary hook into long-running RocksDB threads to support modifying their
 // priority etc. This should become a public API hook once the requirements
 // are better understood.
-extern "C" void RocksDbThreadYield() __attribute__((__weak__));
-#define ROCKSDB_THREAD_YIELD_HOOK() \
-  {                                 \
-    if (RocksDbThreadYield) {       \
-      RocksDbThreadYield();         \
-    }                               \
-  }
+// Returns true if query is aborted.
+extern "C" bool RocksDbThreadYieldAndCheckAbort() __attribute__((__weak__));
+#define ROCKSDB_THREAD_YIELD_CHECK_ABORT() \
+  (RocksDbThreadYieldAndCheckAbort ? RocksDbThreadYieldAndCheckAbort() : false)
 #else
-#define ROCKSDB_THREAD_YIELD_HOOK() \
-  {}
+#define ROCKSDB_THREAD_YIELD_CHECK_ABORT() (false)
 #endif