ccc03ce0fc
If a thread forks, while another thread is holding an snmalloc lock, then the allocator could stop working. This patch attempts to protect against the cases of this. There is one case that is not covered. If a fork occurs during the very first allocation. This can result in the installation of the fork handler racing with the fork, and all bets are off.
306 lines
8.7 KiB
C++
306 lines
8.7 KiB
C++
#pragma once
|
|
|
|
#include "../aal/aal.h"
|
|
#include "../pal/pal.h"
|
|
#include "snmalloc/stl/atomic.h"
|
|
|
|
namespace snmalloc
|
|
{
|
|
class CombiningLockNode;
|
|
|
|
struct CombiningLock
|
|
{
|
|
// Fast path lock incase there is no contention.
|
|
stl::Atomic<bool> flag{false};
|
|
|
|
// MCS queue of work items
|
|
stl::Atomic<CombiningLockNode*> last{nullptr};
|
|
|
|
void release()
|
|
{
|
|
flag.store(false, stl::memory_order_release);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* @brief Combinations of MCS queue lock with Flat Combining
|
|
*
|
|
* Each element in the queue has a pointer to a work item.
|
|
* This means when under contention the thread holding the lock
|
|
* can perform the work.
|
|
*
|
|
* As the work items are arbitrary lambdas there are no simplifications
|
|
* for combining related work items. I.e. original Flat Combining paper
|
|
* might sort a collection of inserts, and perform them in a single traversal.
|
|
*
|
|
* Note that, we should perhaps add a Futex/WakeOnAddress mode to improve
|
|
* performance in the contended case, rather than spinning.
|
|
*/
|
|
class CombiningLockNode
|
|
{
|
|
template<typename Pal>
|
|
static constexpr bool use_wait_on_address =
|
|
pal_supports<PalFeatures::WaitOnAddress, Pal> &&
|
|
SNMALLOC_USE_WAIT_ON_ADDRESS;
|
|
|
|
template<bool HasWaitOnAddress, typename Pal>
|
|
struct WaitWordTypeSelect;
|
|
|
|
template<typename Pal>
|
|
struct WaitWordTypeSelect<true, Pal>
|
|
{
|
|
using type = typename Pal::WaitingWord;
|
|
};
|
|
|
|
template<typename Pal>
|
|
struct WaitWordTypeSelect<false, Pal>
|
|
{
|
|
using type = int;
|
|
};
|
|
|
|
using WaitingWordType =
|
|
typename WaitWordTypeSelect<use_wait_on_address<DefaultPal>, DefaultPal>::
|
|
type;
|
|
|
|
template<typename F>
|
|
friend class CombiningLockNodeTempl;
|
|
|
|
enum class LockStatus : WaitingWordType
|
|
{
|
|
// The work for this node has not been completed.
|
|
WAITING,
|
|
|
|
// The work for this thread has been completed, and it is not the
|
|
// last element in the queue.
|
|
DONE,
|
|
|
|
// The work for this thread has not been completed, and it is the
|
|
// head of the queue.
|
|
HEAD,
|
|
|
|
// The waiter is currently sleeping.
|
|
SLEEPING
|
|
};
|
|
|
|
// Status of the queue, set by the thread at the head of the queue,
|
|
// When it makes the thread for this node either the head of the queue
|
|
// or completes its work.
|
|
stl::Atomic<LockStatus> status{LockStatus::WAITING};
|
|
|
|
// Used to store the queue
|
|
stl::Atomic<CombiningLockNode*> next{nullptr};
|
|
|
|
// Stores the C++ lambda associated with this node in the queue.
|
|
void (*f_raw)(CombiningLockNode*);
|
|
|
|
constexpr CombiningLockNode(void (*f)(CombiningLockNode*)) : f_raw(f) {}
|
|
|
|
void set_status(LockStatus s)
|
|
{
|
|
status.store(s, stl::memory_order_release);
|
|
}
|
|
|
|
template<typename Pal = DefaultPal>
|
|
static void wake(CombiningLockNode* node, LockStatus message)
|
|
{
|
|
if constexpr (!use_wait_on_address<Pal>)
|
|
{
|
|
node->set_status(message);
|
|
}
|
|
else
|
|
{
|
|
if (
|
|
node->status.exchange(message, stl::memory_order_acq_rel) ==
|
|
LockStatus::SLEEPING)
|
|
{
|
|
Pal::notify_one_on_address(node->status);
|
|
}
|
|
}
|
|
}
|
|
|
|
template<typename Pal = DefaultPal>
|
|
void wait()
|
|
{
|
|
if constexpr (!use_wait_on_address<Pal>)
|
|
{
|
|
while (status.load(stl::memory_order_acquire) == LockStatus::WAITING)
|
|
Aal::pause();
|
|
}
|
|
else
|
|
{
|
|
int remaining = 100;
|
|
while (remaining > 0)
|
|
{
|
|
if (status.load(stl::memory_order_acquire) != LockStatus::WAITING)
|
|
return;
|
|
Aal::pause();
|
|
remaining--;
|
|
}
|
|
LockStatus expected = LockStatus::WAITING;
|
|
if (status.compare_exchange_strong(
|
|
expected, LockStatus::SLEEPING, stl::memory_order_acq_rel))
|
|
{
|
|
Pal::wait_on_address(status, LockStatus::SLEEPING);
|
|
if (status.load(stl::memory_order_acquire) == LockStatus::SLEEPING)
|
|
{
|
|
error("Corruption in core locking primitive. Aborting execution.");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
SNMALLOC_SLOW_PATH void attach_slow(CombiningLock& lock)
|
|
{
|
|
// There is contention for the lock, we need to add our work to the
|
|
// queue of pending work
|
|
auto prev = lock.last.exchange(this, stl::memory_order_acq_rel);
|
|
|
|
if (prev != nullptr)
|
|
{
|
|
// If we aren't the head, link into predecessor
|
|
prev->next.store(this, stl::memory_order_release);
|
|
|
|
// Wait to for predecessor to complete
|
|
wait();
|
|
|
|
// Determine if another thread completed our work.
|
|
if (status.load(stl::memory_order_acquire) == LockStatus::DONE)
|
|
return;
|
|
}
|
|
else
|
|
{
|
|
// We are the head of the queue. Spin until we acquire the fast path
|
|
// lock. As we are in the queue future requests shouldn't try to
|
|
// acquire the fast path lock, but stale views of the queue being empty
|
|
// could still be concurrent with this thread.
|
|
while (lock.flag.exchange(true, stl::memory_order_acquire))
|
|
{
|
|
while (lock.flag.load(stl::memory_order_relaxed))
|
|
{
|
|
Aal::pause();
|
|
}
|
|
}
|
|
|
|
// We could set
|
|
// status = LockStatus::HEAD
|
|
// However, the subsequent state assumes it is HEAD, and
|
|
// nothing would read it.
|
|
}
|
|
|
|
// We are the head of the queue, and responsible for
|
|
// waking/performing our and subsequent work.
|
|
auto curr = this;
|
|
while (true)
|
|
{
|
|
// Start pulling in the next element of the queue
|
|
auto n = curr->next.load(stl::memory_order_acquire);
|
|
Aal::prefetch(n);
|
|
|
|
// Perform work for head of the queue
|
|
curr->f_raw(curr);
|
|
|
|
// Determine if there are more elements.
|
|
n = curr->next.load(stl::memory_order_acquire);
|
|
if (n == nullptr)
|
|
break;
|
|
// Signal this work was completed and move on to
|
|
// next item.
|
|
wake(curr, LockStatus::DONE);
|
|
curr = n;
|
|
}
|
|
|
|
// This could be the end of the queue, attempt to close the
|
|
// queue.
|
|
auto curr_c = curr;
|
|
if (lock.last.compare_exchange_strong(
|
|
curr_c,
|
|
nullptr,
|
|
stl::memory_order_release,
|
|
stl::memory_order_relaxed))
|
|
{
|
|
// Queue was successfully closed.
|
|
// Notify last element the work was completed.
|
|
wake(curr, LockStatus::DONE);
|
|
lock.release();
|
|
return;
|
|
}
|
|
|
|
// Failed to close the queue wait for next thread to be
|
|
// added.
|
|
while (curr->next.load(stl::memory_order_relaxed) == nullptr)
|
|
Aal::pause();
|
|
|
|
auto n = curr->next.load(stl::memory_order_acquire);
|
|
|
|
// As we had to wait, give the job to the next thread
|
|
// to carry on performing the work.
|
|
wake(n, LockStatus::HEAD);
|
|
|
|
// Notify the thread that we completed its work.
|
|
// Note that this needs to be before setting curr->status,
|
|
// as after the status is set the thread may deallocate the
|
|
// queue node.
|
|
wake(curr, LockStatus::DONE);
|
|
return;
|
|
}
|
|
};
|
|
|
|
template<typename F>
|
|
class CombiningLockNodeTempl : CombiningLockNode
|
|
{
|
|
template<typename FF>
|
|
friend void with(CombiningLock&, FF&&);
|
|
|
|
// This holds the closure for the lambda
|
|
F f;
|
|
|
|
CombiningLockNodeTempl(CombiningLock& lock, F&& f_)
|
|
: CombiningLockNode([](CombiningLockNode* self) {
|
|
CombiningLockNodeTempl* self_templ =
|
|
reinterpret_cast<CombiningLockNodeTempl*>(self);
|
|
self_templ->f();
|
|
}),
|
|
f(stl::forward<F>(f_))
|
|
{
|
|
attach_slow(lock);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* Lock primitive. This takes a reference to a Lock, and a thunk to
|
|
* call when the lock is available. The thunk should be independent of
|
|
* the current thread as the thunk may be executed by a different thread.
|
|
*/
|
|
template<typename F>
|
|
inline void with(CombiningLock& lock, F&& f)
|
|
{
|
|
// A unix fork while holding a lock can lead to deadlock. Protect against
|
|
// this by not allowing a fork while holding a lock.
|
|
PreventFork pf;
|
|
snmalloc::UNUSED(pf);
|
|
{
|
|
// Test if no one is waiting
|
|
if (SNMALLOC_LIKELY(lock.last.load(stl::memory_order_relaxed) == nullptr))
|
|
{
|
|
// No one was waiting so low contention. Attempt to acquire the flag
|
|
// lock.
|
|
if (SNMALLOC_LIKELY(
|
|
lock.flag.exchange(true, stl::memory_order_acquire) == false))
|
|
{
|
|
// We grabbed the lock.
|
|
// Execute the thunk.
|
|
f();
|
|
|
|
// Release the lock
|
|
lock.release();
|
|
return;
|
|
}
|
|
}
|
|
|
|
// There is contention for the lock, we need to take the slow path
|
|
// with the queue.
|
|
CombiningLockNodeTempl<F> node(lock, stl::forward<F>(f));
|
|
}
|
|
}
|
|
} // namespace snmalloc
|