bug/race in mtdchar device?

[Joakim Tjernlund]

Trying to find a falsh corretion bug I stumbled onto this in mtdchar.c:
			/*
			  FIXME: Allow INTERRUPTIBLE. Which means
			  not having the wait_queue head on the stack.

			  If the wq_head is on the stack, and we
			  leave because we got interrupted, then the
			  wq_head is no longer there when the
			  callback routine tries to wake us up.
			*/
			ret = mtd_erase(mtd, erase);
			if (!ret) {
				set_current_state(TASK_UNINTERRUPTIBLE);
				add_wait_queue(&waitq, &wait);
				if (erase->state != MTD_ERASE_DONE &&
				    erase->state != MTD_ERASE_FAILED)
					schedule();
				remove_wait_queue(&waitq, &wait);
				set_current_state(TASK_RUNNING);

				ret = (erase->state == MTD_ERASE_FAILED)?-EIO:0;
			}
			kfree(erase);

Is the FIXME comment correct? If so, should not init_waitqueue_head(&waitq) be
moved to after set_current_state(TASK_UNINTERRUPTIBLE) ?

Also, it feels a bit odd that mtd_erase is called before setting up the wait queue
as mtd_erase can call the callback before returning.

I am contemplating a changing the code into:
			set_current_state(TASK_UNINTERRUPTIBLE);
			init_waitqueue_head(&waitq);
			add_wait_queue(&waitq, &wait);
			ret = mtd_erase(mtd, erase);
			if (ret) {
				remove_wait_queue(&waitq, &wait);
				set_current_state(TASK_RUNNING);
				kfree(erase);
				return ret;
			}
			schedule(); /* Wait for erase to finish. */
			remove_wait_queue(&waitq, &wait);
			ret = (erase->state == MTD_ERASE_FAILED) ? -EIO : 0;
			kfree(erase);

Does that make sense?

 Jocke