ClickHouse/src/Common/CancelToken.cpp

244 lines
6.3 KiB
C++
Raw Normal View History

2023-01-09 14:58:44 +00:00
#include <Common/CancelToken.h>
namespace DB
{
namespace ErrorCodes
{
extern const int THREAD_WAS_CANCELED;
}
}
#ifdef OS_LINUX /// Because of futex
#include <base/getThreadId.h>
#include <linux/futex.h>
#include <sys/types.h>
#include <sys/syscall.h>
#include <unistd.h>
namespace DB
{
namespace
{
inline Int64 futexWait(void * address, UInt32 value)
{
return syscall(SYS_futex, address, FUTEX_WAIT_PRIVATE, value, nullptr, nullptr, 0);
}
inline Int64 futexWake(void * address, int count)
{
return syscall(SYS_futex, address, FUTEX_WAKE_PRIVATE, count, nullptr, nullptr, 0);
}
}
void CancelToken::Registry::insert(CancelToken * token)
{
2023-03-02 13:36:47 +00:00
std::lock_guard lock(mutex);
2023-01-09 14:58:44 +00:00
threads[token->thread_id] = token;
}
void CancelToken::Registry::remove(CancelToken * token)
{
2023-03-02 13:36:47 +00:00
std::lock_guard lock(mutex);
2023-01-09 14:58:44 +00:00
threads.erase(token->thread_id);
}
void CancelToken::Registry::signal(UInt64 tid)
{
2023-03-02 13:36:47 +00:00
std::lock_guard lock(mutex);
2023-01-09 14:58:44 +00:00
if (auto it = threads.find(tid); it != threads.end())
it->second->signalImpl();
}
void CancelToken::Registry::signal(UInt64 tid, int code, const String & message)
{
2023-03-02 13:36:47 +00:00
std::lock_guard lock(mutex);
2023-01-09 14:58:44 +00:00
if (auto it = threads.find(tid); it != threads.end())
it->second->signalImpl(code, message);
}
const std::shared_ptr<CancelToken::Registry> & CancelToken::Registry::instance()
{
static std::shared_ptr<Registry> registry{new Registry()}; // shared_ptr is used to enforce correct destruction order of tokens and registry
return registry;
}
CancelToken::CancelToken()
: state(disabled)
, thread_id(getThreadId())
, registry(Registry::instance())
{
registry->insert(this);
}
CancelToken::~CancelToken()
{
registry->remove(this);
}
void CancelToken::signal(UInt64 tid)
{
Registry::instance()->signal(tid);
}
void CancelToken::signal(UInt64 tid, int code, const String & message)
{
Registry::instance()->signal(tid, code, message);
}
bool CancelToken::wait(UInt32 * address, UInt32 value)
{
chassert((reinterpret_cast<UInt64>(address) & canceled) == 0); // An `address` must be 2-byte aligned
if (value & signaled) // Can happen after spurious wake-up due to cancel of other thread
return true; // Spin-wait unless signal is handled
UInt64 s = state.load();
while (true)
{
if (s & disabled)
{
// Start non-cancelable wait on futex. Spurious wake-up is possible.
futexWait(address, value);
return true; // Disabled - true is forced
}
if (s & canceled)
return false; // Has already been canceled
if (state.compare_exchange_strong(s, reinterpret_cast<UInt64>(address)))
break; // This futex has been "acquired" by this token
}
// Start cancelable wait. Spurious wake-up is possible.
futexWait(address, value);
2023-01-10 01:26:59 +00:00
// "Release" futex and check for cancellation
2023-01-09 14:58:44 +00:00
s = state.load();
while (true)
{
chassert((s & disabled) != disabled); // `disable()` must not be called from another thread
if (s & canceled)
{
if (s == canceled)
break; // Signaled; futex "release" has been done by the signaling thread
else
{
s = state.load();
continue; // To avoid race (may lead to futex destruction) we have to wait for signaling thread to finish
}
}
if (state.compare_exchange_strong(s, 0))
2023-01-10 01:26:59 +00:00
return true; // There was no cancellation; futex "released"
2023-01-09 14:58:44 +00:00
}
// Reset signaled bit
reinterpret_cast<std::atomic<UInt32> *>(address)->fetch_and(~signaled);
return false;
}
void CancelToken::raise()
{
2023-01-09 15:05:41 +00:00
std::unique_lock lock(signal_mutex);
2023-01-09 14:58:44 +00:00
if (exception_code != 0)
2023-01-26 09:52:47 +00:00
throw DB::Exception::createRuntime(
2023-01-09 14:58:44 +00:00
std::exchange(exception_code, 0),
std::exchange(exception_message, {}));
else
throw DB::Exception(ErrorCodes::THREAD_WAS_CANCELED, "Thread was canceled");
}
void CancelToken::notifyOne(UInt32 * address)
{
futexWake(address, 1);
}
void CancelToken::notifyAll(UInt32 * address)
{
futexWake(address, INT_MAX);
}
void CancelToken::signalImpl()
{
signalImpl(0, {});
}
std::mutex CancelToken::signal_mutex;
void CancelToken::signalImpl(int code, const String & message)
{
// Serialize all signaling threads to avoid races due to concurrent signal()/raise() calls
2023-01-09 15:05:41 +00:00
std::unique_lock lock(signal_mutex);
2023-01-09 14:58:44 +00:00
UInt64 s = state.load();
while (true)
{
if (s & canceled)
return; // Already canceled - don't signal twice
if (state.compare_exchange_strong(s, s | canceled))
break; // It is the canceling thread - should deliver signal if necessary
}
exception_code = code;
exception_message = message;
if ((s & disabled) == disabled)
2023-01-10 01:26:59 +00:00
return; // cancellation is disabled - just signal token for later, but don't wake
2023-01-09 14:58:44 +00:00
std::atomic<UInt32> * address = reinterpret_cast<std::atomic<UInt32> *>(s & disabled);
if (address == nullptr)
return; // Thread is currently not waiting on futex - wake-up not required
// Set signaled bit
UInt32 value = address->load();
while (true)
{
if (value & signaled) // Already signaled, just spin-wait until previous signal is handled by waiter
value = address->load();
else if (address->compare_exchange_strong(value, value | signaled))
break;
}
// Wake all threads waiting on `address`, one of them will be canceled and others will get spurious wake-ups
// Woken canceled thread will reset signaled bit
futexWake(address, INT_MAX);
// Signaling thread must remove address from state to notify canceled thread that `futexWake()` is done, thus `wake()` can return.
// Otherwise we may have race condition: signaling thread may try to wake futex that has been already destructed.
state.store(canceled);
}
Cancelable::Cancelable()
{
CancelToken::local().reset();
}
Cancelable::~Cancelable()
{
CancelToken::local().disable();
}
NonCancelable::NonCancelable()
{
CancelToken::local().disable();
}
NonCancelable::~NonCancelable()
{
CancelToken::local().enable();
}
}
#else
namespace DB
{
void CancelToken::raise()
{
throw DB::Exception(ErrorCodes::THREAD_WAS_CANCELED, "Thread was canceled");
}
}
#endif