Add abort check to yield hook (facebook#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: facebook#13164

Reviewed By: cbi42

Differential Revision: D66556004

Pulled By: george-reynya

fbshipit-source-id: 727875569209cd6d2f29c07f89ecfa641d5ee36f

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)

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;

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