return 0;
 }
 
-static void sched_core_enqueue(struct rq *rq, struct task_struct *p)
+void sched_core_enqueue(struct rq *rq, struct task_struct *p)
 {
        rq->core->core_task_seq++;
 
        rb_add(&p->core_node, &rq->core_tree, rb_sched_core_less);
 }
 
-static void sched_core_dequeue(struct rq *rq, struct task_struct *p)
+void sched_core_dequeue(struct rq *rq, struct task_struct *p)
 {
        rq->core->core_task_seq++;
 
-       if (!p->core_cookie)
+       if (!sched_core_enqueued(p))
                return;
 
        rb_erase(&p->core_node, &rq->core_tree);
+       RB_CLEAR_NODE(&p->core_node);
 }
 
 /*
 
--- /dev/null
+// SPDX-License-Identifier: GPL-2.0-only
+
+#include "sched.h"
+
+/*
+ * A simple wrapper around refcount. An allocated sched_core_cookie's
+ * address is used to compute the cookie of the task.
+ */
+struct sched_core_cookie {
+       refcount_t refcnt;
+};
+
+unsigned long sched_core_alloc_cookie(void)
+{
+       struct sched_core_cookie *ck = kmalloc(sizeof(*ck), GFP_KERNEL);
+       if (!ck)
+               return 0;
+
+       refcount_set(&ck->refcnt, 1);
+       sched_core_get();
+
+       return (unsigned long)ck;
+}
+
+void sched_core_put_cookie(unsigned long cookie)
+{
+       struct sched_core_cookie *ptr = (void *)cookie;
+
+       if (ptr && refcount_dec_and_test(&ptr->refcnt)) {
+               kfree(ptr);
+               sched_core_put();
+       }
+}
+
+unsigned long sched_core_get_cookie(unsigned long cookie)
+{
+       struct sched_core_cookie *ptr = (void *)cookie;
+
+       if (ptr)
+               refcount_inc(&ptr->refcnt);
+
+       return cookie;
+}
+
+/*
+ * sched_core_update_cookie - replace the cookie on a task
+ * @p: the task to update
+ * @cookie: the new cookie
+ *
+ * Effectively exchange the task cookie; caller is responsible for lifetimes on
+ * both ends.
+ *
+ * Returns: the old cookie
+ */
+unsigned long sched_core_update_cookie(struct task_struct *p, unsigned long cookie)
+{
+       unsigned long old_cookie;
+       struct rq_flags rf;
+       struct rq *rq;
+       bool enqueued;
+
+       rq = task_rq_lock(p, &rf);
+
+       /*
+        * Since creating a cookie implies sched_core_get(), and we cannot set
+        * a cookie until after we've created it, similarly, we cannot destroy
+        * a cookie until after we've removed it, we must have core scheduling
+        * enabled here.
+        */
+       SCHED_WARN_ON((p->core_cookie || cookie) && !sched_core_enabled(rq));
+
+       enqueued = sched_core_enqueued(p);
+       if (enqueued)
+               sched_core_dequeue(rq, p);
+
+       old_cookie = p->core_cookie;
+       p->core_cookie = cookie;
+
+       if (enqueued)
+               sched_core_enqueue(rq, p);
+
+       /*
+        * If task is currently running, it may not be compatible anymore after
+        * the cookie change, so enter the scheduler on its CPU to schedule it
+        * away.
+        */
+       if (task_running(rq, p))
+               resched_curr(rq);
+
+       task_rq_unlock(rq, p, &rf);
+
+       return old_cookie;
+}
+
+static unsigned long sched_core_clone_cookie(struct task_struct *p)
+{
+       unsigned long cookie, flags;
+
+       raw_spin_lock_irqsave(&p->pi_lock, flags);
+       cookie = sched_core_get_cookie(p->core_cookie);
+       raw_spin_unlock_irqrestore(&p->pi_lock, flags);
+
+       return cookie;
+}
+
+void sched_core_free(struct task_struct *p)
+{
+       sched_core_put_cookie(p->core_cookie);
+}
 
 
 extern void queue_core_balance(struct rq *rq);
 
+static inline bool sched_core_enqueued(struct task_struct *p)
+{
+       return !RB_EMPTY_NODE(&p->core_node);
+}
+
+extern void sched_core_enqueue(struct rq *rq, struct task_struct *p);
+extern void sched_core_dequeue(struct rq *rq, struct task_struct *p);
+
+extern void sched_core_get(void);
+extern void sched_core_put(void);
+
+extern unsigned long sched_core_alloc_cookie(void);
+extern void sched_core_put_cookie(unsigned long cookie);
+extern unsigned long sched_core_get_cookie(unsigned long cookie);
+extern unsigned long sched_core_update_cookie(struct task_struct *p, unsigned long cookie);
+
 #else /* !CONFIG_SCHED_CORE */
 
 static inline bool sched_core_enabled(struct rq *rq)