Commit 4d3ca89b authored by Tejun Heo's avatar Tejun Heo
Browse files

sched_ext: Refactor consume_remote_task() 



The tricky p->scx.holding_cpu handling was split across
consume_remote_task() body and move_task_to_local_dsq(). Refactor such that:

- All the tricky part is now in the new unlink_dsq_and_lock_src_rq() with
  consolidated documentation.

- move_task_to_local_dsq() now implements straightforward task migration
  making it easier to use in other places.

- dispatch_to_local_dsq() is another user move_task_to_local_dsq(). The
  usage is updated accordingly. This makes the local and remote cases more
  symmetric.

No functional changes intended.

v2: s/task_rq/src_rq/ for consistency.

Signed-off-by: default avatarTejun Heo <tj@kernel.org>
Acked-by: default avatarDavid Vernet <void@manifault.com>
parent fdaedba2

kernel/sched/ext.c

+76 −69
Original line number Diff line number Diff line
@@ -2178,49 +2178,13 @@ static bool yield_to_task_scx(struct rq *rq, struct task_struct *to) 
 * @src_rq: rq to move the task from, locked on entry, released on return

 * @dst_rq: rq to move the task into, locked on return

 *

 * Move @p which is currently on @src_rq to @dst_rq's local DSQ. The caller

 * must:

 *

 * 1. Start with exclusive access to @p either through its DSQ lock or

 *    %SCX_OPSS_DISPATCHING flag.

 *

 * 2. Set @p->scx.holding_cpu to raw_smp_processor_id().

 *

 * 3. Remember task_rq(@p) as @src_rq. Release the exclusive access so that we

 *    don't deadlock with dequeue.

 *

 * 4. Lock @src_rq from #3.

 *

 * 5. Call this function.

 *

 * Returns %true if @p was successfully moved. %false after racing dequeue and

 * losing. On return, @src_rq is unlocked and @dst_rq is locked.

 * Move @p which is currently on @src_rq to @dst_rq's local DSQ.

 */

static bool move_task_to_local_dsq(struct task_struct *p, u64 enq_flags,

static void move_task_to_local_dsq(struct task_struct *p, u64 enq_flags,

				   struct rq *src_rq, struct rq *dst_rq)

{

	lockdep_assert_rq_held(src_rq);



	/*

	 * If dequeue got to @p while we were trying to lock @src_rq, it'd have

	 * cleared @p->scx.holding_cpu to -1. While other cpus may have updated

	 * it to different values afterwards, as this operation can't be

	 * preempted or recurse, @p->scx.holding_cpu can never become

	 * raw_smp_processor_id() again before we're done. Thus, we can tell

	 * whether we lost to dequeue by testing whether @p->scx.holding_cpu is

	 * still raw_smp_processor_id().

	 *

	 * @p->rq couldn't have changed if we're still the holding cpu.

	 *

	 * See dispatch_dequeue() for the counterpart.

	 */

	if (unlikely(p->scx.holding_cpu != raw_smp_processor_id()) ||

	    WARN_ON_ONCE(src_rq != task_rq(p))) {

		raw_spin_rq_unlock(src_rq);

		raw_spin_rq_lock(dst_rq);

		return false;

	}



	/* the following marks @p MIGRATING which excludes dequeue */

	deactivate_task(src_rq, p, 0);

	set_task_cpu(p, cpu_of(dst_rq));
@@ -2239,8 +2203,6 @@ static bool move_task_to_local_dsq(struct task_struct *p, u64 enq_flags, 
	dst_rq->scx.extra_enq_flags = enq_flags;

	activate_task(dst_rq, p, 0);

	dst_rq->scx.extra_enq_flags = 0;



	return true;

}



#endif	/* CONFIG_SMP */
@@ -2305,28 +2267,69 @@ static bool task_can_run_on_remote_rq(struct task_struct *p, struct rq *rq, 
	return true;

}



static bool consume_remote_task(struct rq *rq, struct scx_dispatch_q *dsq,

				struct task_struct *p, struct rq *task_rq)

{

	lockdep_assert_held(&dsq->lock);	/* released on return */

/**

 * unlink_dsq_and_lock_src_rq() - Unlink task from its DSQ and lock its task_rq

 * @p: target task

 * @dsq: locked DSQ @p is currently on

 * @src_rq: rq @p is currently on, stable with @dsq locked

 *

 * Called with @dsq locked but no rq's locked. We want to move @p to a different

 * DSQ, including any local DSQ, but are not locking @src_rq. Locking @src_rq is

 * required when transferring into a local DSQ. Even when transferring into a

 * non-local DSQ, it's better to use the same mechanism to protect against

 * dequeues and maintain the invariant that @p->scx.dsq can only change while

 * @src_rq is locked, which e.g. scx_dump_task() depends on.

 *

 * We want to grab @src_rq but that can deadlock if we try while locking @dsq,

 * so we want to unlink @p from @dsq, drop its lock and then lock @src_rq. As

 * this may race with dequeue, which can't drop the rq lock or fail, do a little

 * dancing from our side.

 *

 * @p->scx.holding_cpu is set to this CPU before @dsq is unlocked. If @p gets

 * dequeued after we unlock @dsq but before locking @src_rq, the holding_cpu

 * would be cleared to -1. While other cpus may have updated it to different

 * values afterwards, as this operation can't be preempted or recurse, the

 * holding_cpu can never become this CPU again before we're done. Thus, we can

 * tell whether we lost to dequeue by testing whether the holding_cpu still

 * points to this CPU. See dispatch_dequeue() for the counterpart.

 *

 * On return, @dsq is unlocked and @src_rq is locked. Returns %true if @p is

 * still valid. %false if lost to dequeue.

 */

static bool unlink_dsq_and_lock_src_rq(struct task_struct *p,

				       struct scx_dispatch_q *dsq,

				       struct rq *src_rq)

{

	s32 cpu = raw_smp_processor_id();



	lockdep_assert_held(&dsq->lock);



	/*

	 * @dsq is locked and @p is on a remote rq. @p is currently protected by

	 * @dsq->lock. We want to pull @p to @rq but may deadlock if we grab

	 * @task_rq while holding @dsq and @rq locks. As dequeue can't drop the

	 * rq lock or fail, do a little dancing from our side. See

	 * move_task_to_local_dsq().

	 */

	WARN_ON_ONCE(p->scx.holding_cpu >= 0);

	task_unlink_from_dsq(p, dsq);

	dsq_mod_nr(dsq, -1);

	p->scx.holding_cpu = raw_smp_processor_id();

	p->scx.holding_cpu = cpu;



	raw_spin_unlock(&dsq->lock);

	raw_spin_rq_lock(src_rq);



	raw_spin_rq_unlock(rq);

	raw_spin_rq_lock(task_rq);

	/* task_rq couldn't have changed if we're still the holding cpu */

	return likely(p->scx.holding_cpu == cpu) &&

		!WARN_ON_ONCE(src_rq != task_rq(p));

}



static bool consume_remote_task(struct rq *this_rq, struct scx_dispatch_q *dsq,

				struct task_struct *p, struct rq *src_rq)

{

	raw_spin_rq_unlock(this_rq);



	return move_task_to_local_dsq(p, 0, task_rq, rq);

	if (unlink_dsq_and_lock_src_rq(p, dsq, src_rq)) {

		move_task_to_local_dsq(p, 0, src_rq, this_rq);

		return true;

	} else {

		raw_spin_rq_unlock(src_rq);

		raw_spin_rq_lock(this_rq);

		return false;

	}

}

#else	/* CONFIG_SMP */

static inline bool task_can_run_on_remote_rq(struct task_struct *p, struct rq *rq, bool trigger_error) { return false; }
@@ -2430,7 +2433,8 @@ dispatch_to_local_dsq(struct rq *rq, u64 dsq_id, struct task_struct *p, 
		 * As DISPATCHING guarantees that @p is wholly ours, we can

		 * pretend that we're moving from a DSQ and use the same

		 * mechanism - mark the task under transfer with holding_cpu,

		 * release DISPATCHING and then follow the same protocol.

		 * release DISPATCHING and then follow the same protocol. See

		 * unlink_dsq_and_lock_src_rq().

		 */

		p->scx.holding_cpu = raw_smp_processor_id();
@@ -2443,28 +2447,31 @@ dispatch_to_local_dsq(struct rq *rq, u64 dsq_id, struct task_struct *p, 
			raw_spin_rq_lock(src_rq);

		}



		if (src_rq == dst_rq) {

		/* task_rq couldn't have changed if we're still the holding cpu */

		dsp = p->scx.holding_cpu == raw_smp_processor_id() &&

			!WARN_ON_ONCE(src_rq != task_rq(p));



		if (likely(dsp)) {

			/*

			 * As @p is staying on the same rq, there's no need to

			 * If @p is staying on the same rq, there's no need to

			 * go through the full deactivate/activate cycle.

			 * Optimize by abbreviating the operations in

			 * move_task_to_local_dsq().

			 */

			dsp = p->scx.holding_cpu == raw_smp_processor_id();

			if (likely(dsp)) {

			if (src_rq == dst_rq) {

				p->scx.holding_cpu = -1;

				dispatch_enqueue(&dst_rq->scx.local_dsq, p,

						 enq_flags);

			}

				dispatch_enqueue(&dst_rq->scx.local_dsq,

						 p, enq_flags);

			} else {

			dsp = move_task_to_local_dsq(p, enq_flags,

				move_task_to_local_dsq(p, enq_flags,

						       src_rq, dst_rq);

			}



			/* if the destination CPU is idle, wake it up */

		if (dsp && sched_class_above(p->sched_class,

			if (sched_class_above(p->sched_class,

					      dst_rq->curr->sched_class))

				resched_curr(dst_rq);

		}



		/* switch back to @rq lock */

		if (rq != dst_rq) {