#include "Python.h"

#include "pycore_object.h" // _Py_ExplicitMergeRefcount

#include "pycore_brc.h" // struct _brc_thread_state

#include "pycore_ceval.h" // _Py_set_eval_breaker_bit

#include "pycore_llist.h" // struct llist_node

#include "pycore_pystate.h" // _PyThreadStateImpl

#ifdef Py_GIL_DISABLED

static struct _brc_bucket *

get_bucket(PyInterpreterState *interp, uintptr_t tid)

{

return &interp->brc.table[tid % _Py_BRC_NUM_BUCKETS];

}

static _PyThreadStateImpl *

find_thread_state(struct _brc_bucket *bucket, uintptr_t thread_id)

{

struct llist_node *node;

llist_for_each(node, &bucket->root) {

// Get the containing _PyThreadStateImpl from the linked-list node.

_PyThreadStateImpl *ts = llist_data(node, _PyThreadStateImpl,

brc.bucket_node);

if (ts->brc.tid == thread_id) {

return ts;

}

}

return NULL;

}

static void

merge_queued_refcounts(_PyObjectStack *stack)

{

for (_PyObjectStackChunk *buf = stack->head; buf != NULL; buf = buf->prev) {

for (Py_ssize_t i = 0; i < buf->n; i++) {

_Py_ExplicitMergeRefcount(buf->objs[i], 0);

}

}

}

static void

decref_merged_objects(_PyObjectStack *stack)

{

PyObject *ob;

while ((ob = _PyObjectStack_Pop(stack)) != NULL) {

Py_DECREF(ob);

}

}

_Py_brc_queue_object(PyObject *ob)

{

PyInterpreterState *interp = _PyInterpreterState_GET();

uintptr_t ob_tid = _Py_atomic_load_uintptr(&ob->ob_tid);

if (ob_tid == 0) {

// The owning thread may have concurrently decided to merge the

// refcount fields.

Py_DECREF(ob);

return;

}

struct _brc_bucket *bucket = get_bucket(interp, ob_tid);

PyMutex_Lock(&bucket->mutex);

_PyThreadStateImpl *tstate = find_thread_state(bucket, ob_tid);

if (tstate == NULL) {

// If we didn't find the owning thread then it must have already exited.

// It's safe (and necessary) to merge the refcount. Subtract one when

// merging because we've stolen a reference.

Py_ssize_t refcount = _Py_ExplicitMergeRefcount(ob, -1);

PyMutex_Unlock(&bucket->mutex);

if (refcount == 0) {

_Py_Dealloc(ob);

}

return;

}

if (_PyObjectStack_Push(&tstate->brc.objects_to_merge, ob) < 0) {

PyMutex_Unlock(&bucket->mutex);

// Fall back to stopping all threads and manually merging the refcount

// if we can't enqueue the object to be merged.

_PyEval_StopTheWorld(interp);

Py_ssize_t refcount = _Py_ExplicitMergeRefcount(ob, -1);

_PyEval_StartTheWorld(interp);

if (refcount == 0) {

_Py_Dealloc(ob);

}

return;

}

if (_PyThreadState_TrySuspendDetached(&tstate->base)) {

// The owning thread is detached (e.g. blocked on a lock or in a

// system call) and may not run Python code again for a long time,

// so merge its queue on its behalf instead of waiting for it. While

// it is held in the "suspended" state it cannot attach and therefore

// cannot touch ob_ref_local or ob_tid.

_PyObjectStack merged = {0};

_PyObjectStack_Merge(&merged, &tstate->brc.objects_to_merge);

merge_queued_refcounts(&merged);

_PyThreadState_ResumeDetached(&tstate->base);

PyMutex_Unlock(&bucket->mutex);

decref_merged_objects(&merged);

return;

}

// Notify owning thread

_Py_set_eval_breaker_bit(&tstate->base, _PY_EVAL_EXPLICIT_MERGE_BIT);

PyMutex_Unlock(&bucket->mutex);

}

static void

merge_queued_objects(_PyObjectStack *to_merge)

{

PyObject *ob;

while ((ob = _PyObjectStack_Pop(to_merge)) != NULL) {

// Subtract one when merging because the queue had a reference.

Py_ssize_t refcount = _Py_ExplicitMergeRefcount(ob, -1);

if (refcount == 0) {

_Py_Dealloc(ob);

}

}

}

_Py_brc_merge_refcounts(PyThreadState *tstate)

{

struct _brc_thread_state *brc = &((_PyThreadStateImpl *)tstate)->brc;

struct _brc_bucket *bucket = get_bucket(tstate->interp, brc->tid);

assert(brc->tid == _Py_ThreadId());

// Append all objects into a local stack. We don't want to hold the lock

// while calling destructors.

PyMutex_Lock(&bucket->mutex);

_PyObjectStack_Merge(&brc->local_objects_to_merge, &brc->objects_to_merge);

PyMutex_Unlock(&bucket->mutex);

// Process the local stack until it's empty

merge_queued_objects(&brc->local_objects_to_merge);

}

_Py_brc_init_state(PyInterpreterState *interp)

{

struct _brc_state *brc = &interp->brc;

for (Py_ssize_t i = 0; i < _Py_BRC_NUM_BUCKETS; i++) {

llist_init(&brc->table[i].root);

}

}

_Py_brc_init_thread(PyThreadState *tstate)

{

struct _brc_thread_state *brc = &((_PyThreadStateImpl *)tstate)->brc;

uintptr_t tid = _Py_ThreadId();

// Add ourself to the hashtable

struct _brc_bucket *bucket = get_bucket(tstate->interp, tid);

PyMutex_Lock(&bucket->mutex);

brc->tid = tid;

llist_insert_tail(&bucket->root, &brc->bucket_node);

PyMutex_Unlock(&bucket->mutex);

}

_Py_brc_remove_thread(PyThreadState *tstate)

{

struct _brc_thread_state *brc = &((_PyThreadStateImpl *)tstate)->brc;

if (brc->tid == 0) {

// The thread state may have been created, but never bound to a native

// thread and therefore never added to the hashtable.

assert(tstate->_status.bound == 0);

return;

}

struct _brc_bucket *bucket = get_bucket(tstate->interp, brc->tid);

// We need to fully process any objects to merge before removing ourself

// from the hashtable. It is not safe to perform any refcount operations

// after we are removed. After that point, other threads treat our objects

// as abandoned and may merge the objects' refcounts directly.

bool empty = false;

while (!empty) {

// Process the local stack until it's empty

merge_queued_objects(&brc->local_objects_to_merge);

PyMutex_Lock(&bucket->mutex);

empty = (brc->objects_to_merge.head == NULL);

if (empty) {

llist_remove(&brc->bucket_node);

}

else {

_PyObjectStack_Merge(&brc->local_objects_to_merge,

&brc->objects_to_merge);

}

PyMutex_Unlock(&bucket->mutex);

}

assert(brc->local_objects_to_merge.head == NULL);

assert(brc->objects_to_merge.head == NULL);

}

_Py_brc_after_fork(PyInterpreterState *interp)

{

// Unlock all bucket mutexes. Some of the buckets may be locked because

// locks can be handed off to a parked thread (see lock.c). We don't have

// to worry about consistency here, because no thread can be actively

// modifying a bucket, but it might be paused (not yet woken up) on a

// PyMutex_Lock while holding that lock.

for (Py_ssize_t i = 0; i < _Py_BRC_NUM_BUCKETS; i++) {

_PyMutex_at_fork_reinit(&interp->brc.table[i].mutex);

}

}

#endif /* Py_GIL_DISABLED */