sched_ext: Make explicit scx_task_iter_relock() calls unnecessary

During tasks iteration, the locks can be dropped using
scx_task_iter_unlock() to perform e.g. sleepable allocations. Afterwards,
scx_task_iter_relock() has to be called prior to other iteration operations,
which is error-prone. This can be easily automated by tracking whether
scx_tasks_lock is held in scx_task_iter and re-acquiring when necessary. It
already tracks whether the task's rq is locked after all.

- Add scx_task_iter->list_locked which remembers whether scx_tasks_lock is
  held.

- Rename scx_task_iter->locked to scx_task_iter->locked_task to better
  distinguish it from ->list_locked.

- Replace scx_task_iter_relock() with __scx_task_iter_maybe_relock() which
  is automatically called by scx_task_iter_next() and scx_task_iter_stop().

- Drop explicit scx_task_iter_relock() calls.

The resulting behavior should be equivalent.

Signed-off-by: Tejun Heo <tj@kernel.org>
Acked-by: Andrea Righi <arighi@nvidia.com>

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 7dedc9a16281b3d5fbae416fe23c51638ae2fa96..7f799345c899442939202c7721ba18ae133e9b7c 100644 (file)
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1488,10 +1488,11 @@ struct bpf_iter_scx_dsq {
  */
 struct scx_task_iter {
        struct sched_ext_entity         cursor;
-       struct task_struct              *locked;
+       struct task_struct              *locked_task;
        struct rq                       *rq;
        struct rq_flags                 rf;
        u32                             cnt;
+       bool                            list_locked;
 };
 
 /**
@@ -1519,15 +1520,16 @@ static void scx_task_iter_start(struct scx_task_iter *iter)
 
        iter->cursor = (struct sched_ext_entity){ .flags = SCX_TASK_CURSOR };
        list_add(&iter->cursor.tasks_node, &scx_tasks);
-       iter->locked = NULL;
+       iter->locked_task = NULL;
        iter->cnt = 0;
+       iter->list_locked = true;
 }
 
 static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
 {
-       if (iter->locked) {
-               task_rq_unlock(iter->rq, iter->locked, &iter->rf);
-               iter->locked = NULL;
+       if (iter->locked_task) {
+               task_rq_unlock(iter->rq, iter->locked_task, &iter->rf);
+               iter->locked_task = NULL;
        }
 }
 
@@ -1537,24 +1539,24 @@ static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
  *
  * If @iter is in the middle of a locked iteration, it may be locking the rq of
  * the task currently being visited in addition to scx_tasks_lock. Unlock both.
- * This function can be safely called anytime during an iteration.
+ * This function can be safely called anytime during an iteration. The next
+ * iterator operation will automatically restore the necessary locking.
  */
 static void scx_task_iter_unlock(struct scx_task_iter *iter)
 {
        __scx_task_iter_rq_unlock(iter);
-       spin_unlock_irq(&scx_tasks_lock);
+       if (iter->list_locked) {
+               iter->list_locked = false;
+               spin_unlock_irq(&scx_tasks_lock);
+       }
 }
 
-/**
- * scx_task_iter_relock - Lock scx_tasks_lock released by scx_task_iter_unlock()
- * @iter: iterator to re-lock
- *
- * Re-lock scx_tasks_lock unlocked by scx_task_iter_unlock(). Note that it
- * doesn't re-lock the rq lock. Must be called before other iterator operations.
- */
-static void scx_task_iter_relock(struct scx_task_iter *iter)
+static void __scx_task_iter_maybe_relock(struct scx_task_iter *iter)
 {
-       spin_lock_irq(&scx_tasks_lock);
+       if (!iter->list_locked) {
+               spin_lock_irq(&scx_tasks_lock);
+               iter->list_locked = true;
+       }
 }
 
 /**
@@ -1567,6 +1569,7 @@ static void scx_task_iter_relock(struct scx_task_iter *iter)
  */
 static void scx_task_iter_stop(struct scx_task_iter *iter)
 {
+       __scx_task_iter_maybe_relock(iter);
        list_del_init(&iter->cursor.tasks_node);
        scx_task_iter_unlock(iter);
 }
@@ -1584,10 +1587,12 @@ static struct task_struct *scx_task_iter_next(struct scx_task_iter *iter)
        struct list_head *cursor = &iter->cursor.tasks_node;
        struct sched_ext_entity *pos;
 
+       __scx_task_iter_maybe_relock(iter);
+
        if (!(++iter->cnt % SCX_TASK_ITER_BATCH)) {
                scx_task_iter_unlock(iter);
                cond_resched();
-               scx_task_iter_relock(iter);
+               __scx_task_iter_maybe_relock(iter);
        }
 
        list_for_each_entry(pos, cursor, tasks_node) {
@@ -1650,7 +1655,7 @@ static struct task_struct *scx_task_iter_next_locked(struct scx_task_iter *iter)
                return NULL;
 
        iter->rq = task_rq_lock(p, &iter->rf);
-       iter->locked = p;
+       iter->locked_task = p;
 
        return p;
 }
@@ -5713,7 +5718,6 @@ static int scx_enable(struct sched_ext_ops *ops, struct bpf_link *link)
                ret = scx_init_task(p, task_group(p), false);
                if (ret) {
                        put_task_struct(p);
-                       scx_task_iter_relock(&sti);
                        scx_task_iter_stop(&sti);
                        scx_error(sch, "ops.init_task() failed (%d) for %s[%d]",
                                  ret, p->comm, p->pid);
@@ -5723,7 +5727,6 @@ static int scx_enable(struct sched_ext_ops *ops, struct bpf_link *link)
                scx_set_task_state(p, SCX_TASK_READY);
 
                put_task_struct(p);
-               scx_task_iter_relock(&sti);
        }
        scx_task_iter_stop(&sti);
        scx_cgroup_unlock();