/*-
 * Copyright (c) 2016 Matthew Macy (mmacy@mattmacy.io)
 * Copyright (c) 2017-2021 Hans Petter Selasky (hselasky@freebsd.org)
 * All rights reserved.
 * Copyright (c) 2024 The FreeBSD Foundation
 *
 * Portions of this software were developed by Björn Zeeb
 * under sponsorship from the FreeBSD Foundation.
 *
 * 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 <sys/types.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/kernel.h>
#include <sys/lock.h>
#include <sys/mutex.h>
#include <sys/proc.h>
#include <sys/sched.h>
#include <sys/smp.h>
#include <sys/queue.h>
#include <sys/taskqueue.h>
#include <sys/kdb.h>

#include <ck_epoch.h>

#include <linux/rcupdate.h>
#include <linux/sched.h>
#include <linux/srcu.h>
#include <linux/slab.h>
#include <linux/kernel.h>
#include <linux/compat.h>
#include <linux/llist.h>
#include <linux/irq_work.h>

/*
 * By defining CONFIG_NO_RCU_SKIP LinuxKPI RCU locks and asserts will
 * not be skipped during panic().
 */
#ifdef CONFIG_NO_RCU_SKIP
#define	RCU_SKIP(void) 0
#else
#define	RCU_SKIP(void)	unlikely(SCHEDULER_STOPPED() || kdb_active)
#endif

struct callback_head {
	union {
		STAILQ_ENTRY(callback_head) entry;
		struct llist_node node;
	};
	rcu_callback_t func;
};

struct linux_epoch_head {
	struct llist_head cb_head;
	struct task task;
} __aligned(CACHE_LINE_SIZE);

struct linux_epoch_record {
	ck_epoch_record_t epoch_record;
	TAILQ_HEAD(, task_struct) ts_head;
	int cpuid;
	int type;
} __aligned(CACHE_LINE_SIZE);

/*
 * Verify that "struct rcu_head" is big enough to hold "struct
 * callback_head". This has been done to avoid having to add special
 * compile flags for including ck_epoch.h to all clients of the
 * LinuxKPI.
 */
CTASSERT(sizeof(struct rcu_head) == sizeof(struct callback_head));

/*
 * Verify that "rcu_section[0]" has the same size as
 * "ck_epoch_section_t". This has been done to avoid having to add
 * special compile flags for including ck_epoch.h to all clients of
 * the LinuxKPI.
 */
CTASSERT(sizeof(((struct task_struct *)0)->rcu_section[0] ==
    sizeof(ck_epoch_section_t)));

/*
 * Verify that "epoch_record" is at beginning of "struct
 * linux_epoch_record":
 */
CTASSERT(offsetof(struct linux_epoch_record, epoch_record) == 0);

CTASSERT(TS_RCU_TYPE_MAX == RCU_TYPE_MAX);

static ck_epoch_t linux_epoch[RCU_TYPE_MAX];
static struct linux_epoch_head linux_epoch_head[RCU_TYPE_MAX];
DPCPU_DEFINE_STATIC(struct linux_epoch_record, linux_epoch_record[RCU_TYPE_MAX]);

static void linux_rcu_cleaner_func(void *, int);

static void
linux_rcu_runtime_init(void *arg __unused)
{
	struct linux_epoch_head *head;
	int i;
	int j;

	for (j = 0; j != RCU_TYPE_MAX; j++) {
		ck_epoch_init(&linux_epoch[j]);

		head = &linux_epoch_head[j];

		TASK_INIT(&head->task, 0, linux_rcu_cleaner_func, head);
		init_llist_head(&head->cb_head);

		CPU_FOREACH(i) {
			struct linux_epoch_record *record;

			record = &DPCPU_ID_GET(i, linux_epoch_record[j]);

			record->cpuid = i;
			record->type = j;
			ck_epoch_register(&linux_epoch[j],
			    &record->epoch_record, NULL);
			TAILQ_INIT(&record->ts_head);
		}
	}
}
SYSINIT(linux_rcu_runtime, SI_SUB_CPU, SI_ORDER_ANY, linux_rcu_runtime_init, NULL);

static void
linux_rcu_cleaner_func(void *context, int pending __unused)
{
	struct linux_epoch_head *head = context;
	struct callback_head *rcu;
	STAILQ_HEAD(, callback_head) tmp_head;
	struct llist_node *node, *next;
	uintptr_t offset;

	/* move current callbacks into own queue */
	STAILQ_INIT(&tmp_head);
	llist_for_each_safe(node, next, llist_del_all(&head->cb_head)) {
		rcu = container_of(node, struct callback_head, node);
		/* re-reverse list to restore chronological order */
		STAILQ_INSERT_HEAD(&tmp_head, rcu, entry);
	}

	/* synchronize */
	linux_synchronize_rcu(head - linux_epoch_head);

	/* dispatch all callbacks, if any */
	while ((rcu = STAILQ_FIRST(&tmp_head)) != NULL) {
		STAILQ_REMOVE_HEAD(&tmp_head, entry);

		offset = (uintptr_t)rcu->func;

		if (offset < LINUX_KFREE_RCU_OFFSET_MAX)
			kfree((char *)rcu - offset);
		else
			rcu->func((struct rcu_head *)rcu);
	}
}

void
linux_rcu_read_lock(unsigned type)
{
	struct linux_epoch_record *record;
	struct task_struct *ts;

	MPASS(type < RCU_TYPE_MAX);

	if (RCU_SKIP())
		return;

	ts = current;

	/* assert valid refcount */
	MPASS(ts->rcu_recurse[type] != INT_MAX);

	if (++(ts->rcu_recurse[type]) != 1)
		return;

	/*
	 * Pin thread to current CPU so that the unlock code gets the
	 * same per-CPU epoch record:
	 */
	sched_pin();

	record = &DPCPU_GET(linux_epoch_record[type]);

	/*
	 * Use a critical section to prevent recursion inside
	 * ck_epoch_begin(). Else this function supports recursion.
	 */
	critical_enter();
	ck_epoch_begin(&record->epoch_record,
	    (ck_epoch_section_t *)&ts->rcu_section[type]);
	TAILQ_INSERT_TAIL(&record->ts_head, ts, rcu_entry[type]);
	critical_exit();
}

void
linux_rcu_read_unlock(unsigned type)
{
	struct linux_epoch_record *record;
	struct task_struct *ts;

	MPASS(type < RCU_TYPE_MAX);

	if (RCU_SKIP())
		return;

	ts = current;

	/* assert valid refcount */
	MPASS(ts->rcu_recurse[type] > 0);
	
	if (--(ts->rcu_recurse[type]) != 0)
		return;

	record = &DPCPU_GET(linux_epoch_record[type]);

	/*
	 * Use a critical section to prevent recursion inside
	 * ck_epoch_end(). Else this function supports recursion.
	 */
	critical_enter();
	ck_epoch_end(&record->epoch_record,
	    (ck_epoch_section_t *)&ts->rcu_section[type]);
	TAILQ_REMOVE(&record->ts_head, ts, rcu_entry[type]);
	critical_exit();

	sched_unpin();
}

bool
linux_rcu_read_lock_held(unsigned type)
{
#ifdef INVARINATS
	struct linux_epoch_record *record __diagused;
	struct task_struct *ts;

	MPASS(type < RCU_TYPE_MAX);

	if (RCU_SKIP())
		return (false);

	if (__current_unallocated(curthread))
		return (false);

	ts = current;
	if (ts->rcu_recurse[type] == 0)
		return (false);

	MPASS(curthread->td_pinned != 0);
	MPASS((record = &DPCPU_GET(linux_epoch_record[type])) &&
	    record->epoch_record.active != 0);
#endif

	return (true);
}

static void
linux_synchronize_rcu_cb(ck_epoch_t *epoch __unused, ck_epoch_record_t *epoch_record, void *arg __unused)
{
	struct linux_epoch_record *record =
	    container_of(epoch_record, struct linux_epoch_record, epoch_record);
	struct thread *td = curthread;
	struct task_struct *ts;

	/* check if blocked on the current CPU */
	if (record->cpuid == PCPU_GET(cpuid)) {
		bool is_sleeping = 0;
		u_char prio = 0;

		/*
		 * Find the lowest priority or sleeping thread which
		 * is blocking synchronization on this CPU core. All
		 * the threads in the queue are CPU-pinned and cannot
		 * go anywhere while the current thread is locked.
		 */
		TAILQ_FOREACH(ts, &record->ts_head, rcu_entry[record->type]) {
			if (ts->task_thread->td_priority > prio)
				prio = ts->task_thread->td_priority;
			is_sleeping |= (ts->task_thread->td_inhibitors != 0);
		}

		if (is_sleeping) {
			thread_unlock(td);
			pause("W", 1);
			thread_lock(td);
		} else {
			/* set new thread priority */
			sched_prio(td, prio);
			/* task switch */
			mi_switch(SW_VOL | SWT_RELINQUISH);
			/*
			 * It is important the thread lock is dropped
			 * while yielding to allow other threads to
			 * acquire the lock pointed to by
			 * TDQ_LOCKPTR(td). Currently mi_switch() will
			 * unlock the thread lock before
			 * returning. Else a deadlock like situation
			 * might happen.
			 */
			thread_lock(td);
		}
	} else {
		/*
		 * To avoid spinning move execution to the other CPU
		 * which is blocking synchronization. Set highest
		 * thread priority so that code gets run. The thread
		 * priority will be restored later.
		 */
		sched_prio(td, 0);
		sched_bind(td, record->cpuid);
	}
}

void
linux_synchronize_rcu(unsigned type)
{
	struct thread *td;
	int was_bound;
	int old_cpu;
	int old_pinned;
	u_char old_prio;

	MPASS(type < RCU_TYPE_MAX);

	if (RCU_SKIP())
		return;

	WITNESS_WARN(WARN_GIANTOK | WARN_SLEEPOK, NULL,
	    "linux_synchronize_rcu() can sleep");

	td = curthread;
	DROP_GIANT();

	/*
	 * Synchronizing RCU might change the CPU core this function
	 * is running on. Save current values:
	 */
	thread_lock(td);

	old_cpu = PCPU_GET(cpuid);
	old_pinned = td->td_pinned;
	old_prio = td->td_priority;
	was_bound = sched_is_bound(td);
	sched_unbind(td);
	td->td_pinned = 0;
	sched_bind(td, old_cpu);

	ck_epoch_synchronize_wait(&linux_epoch[type],
	    &linux_synchronize_rcu_cb, NULL);

	/* restore CPU binding, if any */
	if (was_bound != 0) {
		sched_bind(td, old_cpu);
	} else {
		/* get thread back to initial CPU, if any */
		if (old_pinned != 0)
			sched_bind(td, old_cpu);
		sched_unbind(td);
	}
	/* restore pinned after bind */
	td->td_pinned = old_pinned;

	/* restore thread priority */
	sched_prio(td, old_prio);
	thread_unlock(td);

	PICKUP_GIANT();
}

void
linux_rcu_barrier(unsigned type)
{
	struct linux_epoch_head *head;

	MPASS(type < RCU_TYPE_MAX);

	/*
	 * This function is not obligated to wait for a grace period.
	 * It only waits for RCU callbacks that have already been posted.
	 * If there are no RCU callbacks posted, rcu_barrier() can return
	 * immediately.
	 */
	head = &linux_epoch_head[type];

	/* wait for callbacks to complete */
	taskqueue_drain(linux_irq_work_tq, &head->task);
}

void
linux_call_rcu(unsigned type, struct rcu_head *context, rcu_callback_t func)
{
	struct callback_head *rcu;
	struct linux_epoch_head *head;

	MPASS(type < RCU_TYPE_MAX);

	rcu = (struct callback_head *)context;
	head = &linux_epoch_head[type];

	rcu->func = func;
	llist_add(&rcu->node, &head->cb_head);
	taskqueue_enqueue(linux_irq_work_tq, &head->task);
}

int
init_srcu_struct(struct srcu_struct *srcu)
{
	return (0);
}

void
cleanup_srcu_struct(struct srcu_struct *srcu)
{
}

int
srcu_read_lock(struct srcu_struct *srcu)
{
	linux_rcu_read_lock(RCU_TYPE_SLEEPABLE);
	return (0);
}

void
srcu_read_unlock(struct srcu_struct *srcu, int key __unused)
{
	linux_rcu_read_unlock(RCU_TYPE_SLEEPABLE);
}

void
synchronize_srcu(struct srcu_struct *srcu)
{
	linux_synchronize_rcu(RCU_TYPE_SLEEPABLE);
}

void
srcu_barrier(struct srcu_struct *srcu)
{
	linux_rcu_barrier(RCU_TYPE_SLEEPABLE);
}