// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025 Ctrl IQ, Inc.
 * Author: Sultan Alsawaf <sultan@ciq.com>
 */
#ifndef _LKRG_ED_TASK_TREE_H_
#define _LKRG_ED_TASK_TREE_H_

struct p_ed_process {
	struct rb_node node;
	raw_spinlock_t lock;
	struct p_ed_process_task p_ed_task;
	struct rcu_head rcu;
	unsigned long kill_task_on_unlock;
	struct p_ed_process *defer_kill_next;
};

extern unsigned long cookie_kill_task_on_unlock;
int init_ed_task_cache(void);
void destroy_ed_task_cache(void);
struct p_ed_process *alloc_ed_task(struct task_struct *tsk);
void free_ed_task(struct p_ed_process *edp);
int ed_task_add(struct p_ed_process *edp);
int ed_task_del_current(void);

/*
 * The underscored functions shouldn't be used directly since they don't take
 * struct p_ed_process::lock, which can easily lead to forgetting to take the
 * lock in new code.
 *
 * When using __ed_task_find_rcu() for a task other than `current`, the RCU read
 * lock must stay held for the entire duration the task's ed struct is accessed.
 */
struct p_ed_process *__ed_task_find_rcu(const struct task_struct *tsk);
struct p_ed_process *__ed_task_current(void);

/* Find and lock the current task's ed struct if found */
static inline struct p_ed_process *ed_task_lock_current(void)
{
	struct p_ed_process *edp = __ed_task_current();

	if (edp)
		raw_spin_lock(&edp->lock);

	return edp;
}

/* Find and trylock the current task's ed struct if found */
static inline struct p_ed_process *ed_task_trylock_current(void)
{
	struct p_ed_process *edp = __ed_task_current();

	return edp && raw_spin_trylock(&edp->lock) ? edp : NULL;
}

/*
 * Find and lock the specified task's ed struct if found. Must be called under
 * RCU read lock. When the specified task isn't the current task, the RCU read
 * lock must stay held for the entire duration the task's ed struct is accessed.
 */
static inline struct p_ed_process *
ed_task_find_lock_rcu(const struct task_struct *tsk)
{
	struct p_ed_process *edp = __ed_task_find_rcu(tsk);

	if (edp)
		raw_spin_lock(&edp->lock);

	return edp;
}

/*
 * Unlock the ed task. Kill the task if requested, unless deferral of the task's
 * killing is requested by providing a non-NULL @kill_list. Deferred kills are
 * instead placed onto the provided singly-linked list. Deferral is needed when
 * it's possible for an ed task to be killed underneath the sighand lock, which
 * would otherwise deadlock because killing a task requires taking the sighand
 * lock. Deferred kills must be processed by calling ed_task_kill_deferred().
 *
 * The RCU read lock held when looking up the kill-deferred ed task must remain
 * held until after the complementary call to ed_task_kill_deferred().
 */
static inline void ed_task_unlock_defer_kill(struct p_ed_process *edp,
					     struct p_ed_process **kill_list)
{
	struct task_struct *task_to_kill = NULL;

	/* Kill if kill_task_on_unlock differs from the original cookie value */
	if (edp->kill_task_on_unlock != cookie_kill_task_on_unlock)
		task_to_kill = edp->p_ed_task.p_task;
	raw_spin_unlock(&edp->lock);

	if (!task_to_kill)
		return;

	/* Append the task to a deferred kill list if deferral is requested */
	if (kill_list) {
		edp->defer_kill_next = *kill_list;
		*kill_list = edp;
		return;
	}

	/*
	 * Kill the task outside of the ed task lock to avoid acquiring the
	 * sighand lock underneath the ed task lock. This is needed to make it
	 * safe to nest the ed task lock under the sighand lock.
	 */
	p_kill_task_by_task(task_to_kill);
}

/* Unlock the ed task struct and kill the task if requested */
static inline void ed_task_unlock(struct p_ed_process *edp)
{
	ed_task_unlock_defer_kill(edp, NULL);
}

/*
 * Kill the list of ed tasks whose killing was deferred. The RCU read lock held
 * when looking up the kill-deferred ed task must remain held until after this
 * function finishes killing every ed task on the list.
 */
static inline void ed_task_kill_deferred(struct p_ed_process *kill_list)
{
	struct p_ed_process *edp;

	for (edp = kill_list; edp; edp = edp->defer_kill_next)
		p_kill_task_by_task(edp->p_ed_task.p_task);
}

#endif /* _LKRG_ED_TASK_TREE_H_ */
