/*
 * Implementation of safe memory reclamation scheme using
 * quiescent states.  See InternalDocs/qsbr.md.
 *
 * This is derived from the "GUS" safe memory reclamation technique
 * in FreeBSD written by Jeffrey Roberson. It is heavily modified. Any bugs
 * in this code are likely due to the modifications.
 *
 * The original copyright is preserved below.
 *
 * Copyright (c) 2019,2020 Jeffrey Roberson <jeff@FreeBSD.org>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice unmodified, this list of conditions, and the following
 *    disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include "Python.h"
#include "pycore_interp.h"          // PyInterpreterState
#include "pycore_pystate.h"         // _PyThreadState_GET()
#include "pycore_qsbr.h"
#include "pycore_tstate.h"          // _PyThreadStateImpl
#include "pycore_stats.h"           // FT_STAT_QSBR_POLL_INC()


// Starting size of the array of qsbr thread states
#define MIN_ARRAY_SIZE 8

// Allocate a QSBR thread state from the freelist
static struct _qsbr_thread_state *
qsbr_allocate(struct _qsbr_shared *shared)
{
    struct _qsbr_thread_state *qsbr = shared->freelist;
    if (qsbr == NULL) {
        return NULL;
    }
    shared->freelist = qsbr->freelist_next;
    qsbr->freelist_next = NULL;
    qsbr->shared = shared;
    qsbr->allocated = true;
    return qsbr;
}

// Initialize (or reinitialize) the freelist of QSBR thread states
static void
initialize_new_array(struct _qsbr_shared *shared)
{
    for (Py_ssize_t i = 0; i != shared->size; i++) {
        struct _qsbr_thread_state *qsbr = &shared->array[i].qsbr;
        if (qsbr->tstate != NULL) {
            // Update the thread state pointer to its QSBR state
            _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)qsbr->tstate;
            tstate->qsbr = qsbr;
        }
        if (!qsbr->allocated) {
            // Push to freelist
            qsbr->freelist_next = shared->freelist;
            shared->freelist = qsbr;
        }
    }
}

// Grow the array of QSBR thread states. Returns 0 on success, -1 on failure.
static int
grow_thread_array(struct _qsbr_shared *shared)
{
    Py_ssize_t new_size = shared->size * 2;
    if (new_size < MIN_ARRAY_SIZE) {
        new_size = MIN_ARRAY_SIZE;
    }

    // Overallocate by 63 bytes so we can align to a 64-byte boundary.
    // This avoids potential false sharing between the first entry and other
    // allocations.
    size_t alignment = 64;
    size_t alloc_size = (size_t)new_size * sizeof(struct _qsbr_pad) + alignment - 1;
    void *raw = PyMem_RawCalloc(1, alloc_size);
    if (raw == NULL) {
        return -1;
    }
    struct _qsbr_pad *array = _Py_ALIGN_UP(raw, alignment);

    void *old_raw = shared->array_raw;
    if (shared->array != NULL) {
        memcpy(array, shared->array, shared->size * sizeof(*array));
    }

    shared->array = array;
    shared->array_raw = raw;
    shared->size = new_size;
    shared->freelist = NULL;
    initialize_new_array(shared);

    PyMem_RawFree(old_raw);
    return 0;
}

uint64_t
_Py_qsbr_advance(struct _qsbr_shared *shared)
{
    // NOTE: with 64-bit sequence numbers, we don't have to worry too much
    // about the wr_seq getting too far ahead of rd_seq, but if we ever use
    // 32-bit sequence numbers, we'll need to be more careful.
    return _Py_atomic_add_uint64(&shared->wr_seq, QSBR_INCR) + QSBR_INCR;
}

uint64_t
_Py_qsbr_shared_next(struct _qsbr_shared *shared)
{
    return _Py_qsbr_shared_current(shared) + QSBR_INCR;
}

static uint64_t
qsbr_poll_scan(struct _qsbr_shared *shared)
{
    // Synchronize with store in _Py_qsbr_attach(). We need to ensure that
    // the reads from each thread's sequence number are not reordered to see
    // earlier "offline" states.
    _Py_atomic_fence_seq_cst();

    // Compute the minimum sequence number of all attached threads
    uint64_t min_seq = _Py_atomic_load_uint64(&shared->wr_seq);
    struct _qsbr_pad *array = shared->array;
    for (Py_ssize_t i = 0, size = shared->size; i != size; i++) {
        struct _qsbr_thread_state *qsbr = &array[i].qsbr;

        uint64_t seq = _Py_atomic_load_uint64(&qsbr->seq);
        if (seq != QSBR_OFFLINE && QSBR_LT(seq, min_seq)) {
            min_seq = seq;
        }
    }

    // Update the shared read sequence
    uint64_t rd_seq = _Py_atomic_load_uint64(&shared->rd_seq);
    if (QSBR_LT(rd_seq, min_seq)) {
        // It's okay if the compare-exchange failed: another thread updated it
        (void)_Py_atomic_compare_exchange_uint64(&shared->rd_seq, &rd_seq, min_seq);
        rd_seq = min_seq;
    }

    return rd_seq;
}

bool
_Py_qsbr_poll(struct _qsbr_thread_state *qsbr, uint64_t goal)
{
    assert(_Py_atomic_load_int_relaxed(&_PyThreadState_GET()->state) == _Py_THREAD_ATTACHED);
    assert(((_PyThreadStateImpl *)_PyThreadState_GET())->qsbr == qsbr);

    if (_Py_qbsr_goal_reached(qsbr, goal)) {
        return true;
    }
    FT_STAT_QSBR_POLL_INC();
    uint64_t rd_seq = qsbr_poll_scan(qsbr->shared);
    return QSBR_LEQ(goal, rd_seq);
}

void
_Py_qsbr_attach(struct _qsbr_thread_state *qsbr)
{
    assert(qsbr->seq == 0 && "already attached");

    uint64_t seq = _Py_qsbr_shared_current(qsbr->shared);
    _Py_atomic_store_uint64(&qsbr->seq, seq);  // needs seq_cst
}

void
_Py_qsbr_detach(struct _qsbr_thread_state *qsbr)
{
    assert(qsbr->seq != 0 && "already detached");

    _Py_atomic_store_uint64_release(&qsbr->seq, QSBR_OFFLINE);
}

Py_ssize_t
_Py_qsbr_reserve(PyInterpreterState *interp)
{
    struct _qsbr_shared *shared = &interp->qsbr;

    PyMutex_Lock(&shared->mutex);
    // Try allocating from our internal freelist
    struct _qsbr_thread_state *qsbr = qsbr_allocate(shared);

    // If there are no free entries, we pause all threads, grow the array,
    // and update the pointers in PyThreadState to entries in the new array.
    if (qsbr == NULL) {
        _PyEval_StopTheWorld(interp);
        if (grow_thread_array(shared) == 0) {
            qsbr = qsbr_allocate(shared);
        }
        _PyEval_StartTheWorld(interp);
    }

    // Return an index rather than the pointer because the array may be
    // resized and the pointer invalidated.
    Py_ssize_t index = -1;
    if (qsbr != NULL) {
        index = (struct _qsbr_pad *)qsbr - shared->array;
    }
    PyMutex_Unlock(&shared->mutex);
    return index;
}

void
_Py_qsbr_register(_PyThreadStateImpl *tstate, PyInterpreterState *interp,
                  Py_ssize_t index)
{
    // Associate the QSBR state with the thread state
    struct _qsbr_shared *shared = &interp->qsbr;

    PyMutex_Lock(&shared->mutex);
    struct _qsbr_thread_state *qsbr = &interp->qsbr.array[index].qsbr;
    assert(qsbr->allocated && qsbr->tstate == NULL);
    qsbr->tstate = (PyThreadState *)tstate;
    tstate->qsbr = qsbr;
    PyMutex_Unlock(&shared->mutex);
}

void
_Py_qsbr_unregister(PyThreadState *tstate)
{
    struct _qsbr_shared *shared = &tstate->interp->qsbr;
    struct _PyThreadStateImpl *tstate_imp = (_PyThreadStateImpl*) tstate;

    // gh-119369: GIL must be released (if held) to prevent deadlocks, because
    // we might not have an active tstate, which means that blocking on PyMutex
    // locks will not implicitly release the GIL.
    assert(!tstate->holds_gil);

    PyMutex_Lock(&shared->mutex);
    // NOTE: we must load (or reload) the thread state's qbsr inside the mutex
    // because the array may have been resized (changing tstate->qsbr) while
    // we waited to acquire the mutex.
    struct _qsbr_thread_state *qsbr = tstate_imp->qsbr;

    assert(qsbr->seq == 0 && "thread state must be detached");
    assert(qsbr->allocated && qsbr->tstate == tstate);

    tstate_imp->qsbr = NULL;
    qsbr->tstate = NULL;
    qsbr->allocated = false;
    qsbr->freelist_next = shared->freelist;
    shared->freelist = qsbr;
    PyMutex_Unlock(&shared->mutex);
}

void
_Py_qsbr_fini(PyInterpreterState *interp)
{
    struct _qsbr_shared *shared = &interp->qsbr;
    PyMem_RawFree(shared->array_raw);
    shared->array = NULL;
    shared->array_raw = NULL;
    shared->size = 0;
    shared->freelist = NULL;
}

void
_Py_qsbr_after_fork(_PyThreadStateImpl *tstate)
{
    struct _qsbr_thread_state *this_qsbr = tstate->qsbr;
    struct _qsbr_shared *shared = this_qsbr->shared;

    _PyMutex_at_fork_reinit(&shared->mutex);

    for (Py_ssize_t i = 0; i != shared->size; i++) {
        struct _qsbr_thread_state *qsbr = &shared->array[i].qsbr;
        if (qsbr != this_qsbr && qsbr->allocated) {
            qsbr->tstate = NULL;
            qsbr->allocated = false;
            qsbr->freelist_next = shared->freelist;
            shared->freelist = qsbr;
        }
    }
}