Re: [PATCH] sched_ext: Add cookie API for early qseq capture

[Kuba Piecuch <jpiecuch@google.com>]

Hi Cheng-Yang,

On Wed May 6, 2026 at 7:59 AM UTC, Cheng-Yang Chou wrote:
> @@ -8505,13 +8507,67 @@ __bpf_kfunc void scx_bpf_dsq_insert_vtime(struct task_struct *p, u64 dsq_id,
>  	scx_dsq_insert_vtime(sch, p, dsq_id, slice, vtime, enq_flags);
>  }
>  
> +/**
> + * scx_bpf_task_get_cookie - Get an opaque dispatch cookie for a task
> + * @p: task_struct to read cookie from
> + *
> + * Returns an opaque u64 cookie encoding @p's current qseq. Call this
> + * before pre-dispatch validity checks and pass the result to
> + * scx_bpf_dsq_insert_with_cookie() to extend the qseq protection window.

This is just my opinion, but I thought the point of introducing cookies is to
not have to expose qseq as a concept to the BPF schedulers (or people writing
them). This includes comments documenting the kfuncs.

Wouldn't it be better if we just focus on describing the semantics without
mentioning qseq? Perhaps something like:

  Returns an opaque u64 dispatch cookie. Pass the cookie to
  scx_bpf_dsq_insert_with_cookie() to extend the time window during which
  sched_ext will detect racing dequeues/enqueues of the task being dispatched.
  The extended time window begins with the call to scx_bpf_task_get_cookie()
  and ends at the same point as for dispatches without cookies, i.e. at the
  point where sched_ext attempts to finish the dispatch.

> + *
> + * For schedulers that do not implement properly synchronized dequeue only.

"only" seems quite strong here. How about "This API is intended for schedulers
that do not implement properly synchronized dequeue"?

> + */
> +__bpf_kfunc u64 scx_bpf_task_get_cookie(struct task_struct *p)
> +{
> +	return atomic_long_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK;
> +}
> +
> +/**
> + * scx_bpf_dsq_insert_with_cookie - Insert a task using an early-captured cookie
> + * @p: task_struct to insert
> + * @dsq_id: DSQ to insert into
> + * @enq_flags: SCX_ENQ_*
> + * @cookie: cookie from scx_bpf_task_get_cookie()
> + * @aux: implicit BPF argument
> + *
> + * Like scx_bpf_dsq_insert() with slice=0, but uses @cookie's qseq instead
> + * of re-reading ops_state at insert time. A stale cookie causes
> + * finish_dispatch() to silently discard the dispatch. Use
> + * scx_bpf_task_set_slice() to set a non-default slice.
> + *
> + * Returns %true on success, %false on failure.
> + */
> +__bpf_kfunc bool scx_bpf_dsq_insert_with_cookie(struct task_struct *p,
> +						 u64 dsq_id, u64 enq_flags,
> +						 u64 cookie,
> +						 const struct bpf_prog_aux *aux)
> +{
> +	struct scx_sched *sch;
> +
> +	guard(rcu)();
> +	sch = scx_prog_sched(aux);
> +	if (unlikely(!sch))
> +		return false;
> +
> +	if (!scx_dsq_insert_preamble(sch, p, dsq_id, &enq_flags))
> +		return false;
> +
> +	p->scx.slice = p->scx.slice ?: 1;
> +
> +	scx_dsq_insert_commit(sch, p, dsq_id, enq_flags, (unsigned long)cookie);
> +
> +	return true;
> +}
> +
>  __bpf_kfunc_end_defs();
>  
>  BTF_KFUNCS_START(scx_kfunc_ids_enqueue_dispatch)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_insert, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_insert___v2, KF_IMPLICIT_ARGS | KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_dsq_insert_with_cookie, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, __scx_bpf_dsq_insert_vtime, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_insert_vtime, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_task_get_cookie, KF_RCU)
>  BTF_KFUNCS_END(scx_kfunc_ids_enqueue_dispatch)
>  
>  static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
> @@ -10181,6 +10237,7 @@ BTF_ID_FLAGS(func, scx_bpf_put_cpumask, KF_RELEASE)
>  BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_task_cid, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_task_get_cookie, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_cpu_rq, KF_IMPLICIT_ARGS)
>  BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_IMPLICIT_ARGS | KF_RET_NULL)
>  BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)

Thank you for the patch!

Could you add a simple selftest for the two new kfuncs?

Thanks,
Kuba