Kernel wait-queue preemption

preempt.txt

Resources: {
        atomic_t done = 0;
        struct task *wtask;
}

Lets first have a look at the code that is executed on each CPU without
specifying any timings:

-----------------------------------+-----------------------------------
               CPU 1               |                CPU 2
-----------------------------------+-----------------------------------
  wtask = current;                 |
                             <some barrier>
                                   |
  if (atomic_read(&done))          |  atomic_inc(&done);
          goto done;               |  wake_up_process(wtask);
                                   |
  for(;;) {                        |
          set_current_task(        |
             TASK_INTERRUPTIBLE);  |
                                   |
          if (atomic_read(&done))  |
                  break;           |
                                   |
          schedule();              |
  }                                |
                                   |
  set_current_task(TASK_RUNNING);  |
done:                              |
                                   |
-----------------------------------+-----------------------------------

Now lets take this example, add some unfortunate timing and a kernel-preemption:

-----------------------------------+-----------------------------------
               CPU 1               |                CPU 2
-----------------------------------+-----------------------------------
  wtask = current;                 |
                             <some barrier>
                                   |
  if (atomic_read(&done))          |
          goto done;               |
                                   |
  for(;;) {                        |
                                   |
                                   |  atomic_inc(&done);
                                   |  wake_up_process(wtask);
                                   |
          set_current_task(        |
             TASK_INTERRUPTIBLE);  |
                                   |
  <interrupt>                      |
  <kernel preemption>              |
  <switch to more important task>  |
                                   |
  Who will wake us up now? We      |
  will never get to the code below |
  as we are marked as              |
  TASK_INTERRUPTIBLE and so the    |
  scheduler will not run us on     |
  any CPU...                       |
                                   |
          if (atomic_read(&done))  |
                  break;           |
                                   |
          preempt_enable();        |
                                   |
          schedule();              |
                                   |
-----------------------------------+-----------------------------------

What's the correct fix? Disable preemption like below?

-----------------------------------+-----------------------------------
               CPU 1               |                CPU 2
-----------------------------------+-----------------------------------
  wtask = current;                 |
                             <some barrier>
                                   |
  if (atomic_read(&done))          |  atomic_inc(&done);
          goto done;               |  wake_up_process(wtask);
                                   |
  for(;;) {                        |
          preempt_disable();       |
                                   |
          set_current_task(        |
             TASK_INTERRUPTIBLE);  |
                                   |
          if (atomic_read(&done))  |
                  break;           |
                                   |
          preempt_enable();        |
                                   |
          schedule();              |
  }                                |
                                   |
  set_current_task(TASK_RUNNING);  |
  preempt_enable();                |
done:                              |
                                   |
-----------------------------------+-----------------------------------