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
This commit is contained in:
George Reynya 2024-11-27 16:45:59 -08:00 committed by Changyu Bi
parent b0c40216ff
commit 0d9fc96c44
3 changed files with 18 additions and 11 deletions

View file

@ -3247,8 +3247,13 @@ Status DBImpl::MultiGetImpl(
s = Status::Aborted(); s = Status::Aborted();
break; break;
} }
// This could be a long-running operation // 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) // Post processing (decrement reference counts and record statistics)

View file

@ -581,8 +581,14 @@ bool DBIter::FindNextUserEntryInternal(bool skipping_saved_key,
} else { } else {
iter_.Next(); iter_.Next();
} }
// This could be a long-running operation due to tombstones, etc. // 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()); } while (iter_.Valid());
valid_ = false; valid_ = false;

View file

@ -24,14 +24,10 @@
// A temporary hook into long-running RocksDB threads to support modifying their // A temporary hook into long-running RocksDB threads to support modifying their
// priority etc. This should become a public API hook once the requirements // priority etc. This should become a public API hook once the requirements
// are better understood. // are better understood.
extern "C" void RocksDbThreadYield() __attribute__((__weak__)); // Returns true if query is aborted.
#define ROCKSDB_THREAD_YIELD_HOOK() \ extern "C" bool RocksDbThreadYieldAndCheckAbort() __attribute__((__weak__));
{ \ #define ROCKSDB_THREAD_YIELD_CHECK_ABORT() \
if (RocksDbThreadYield) { \ (RocksDbThreadYieldAndCheckAbort ? RocksDbThreadYieldAndCheckAbort() : false)
RocksDbThreadYield(); \
} \
}
#else #else
#define ROCKSDB_THREAD_YIELD_HOOK() \ #define ROCKSDB_THREAD_YIELD_CHECK_ABORT() (false)
{}
#endif #endif