2 files changed, 24 insertions, 1 deletions

diff --git a/share/man/man9/taskqueue.9 b/share/man/man9/taskqueue.9
index 5ee7fc8b36a4..312cffe9d9f2 100644
--- a/share/man/man9/taskqueue.9
+++ b/share/man/man9/taskqueue.9
@@ -223,6 +223,8 @@ Otherwise, the task is scheduled for enqueueing in the future,
 after the absolute value of
 .Va ticks
 is passed.
+This function returns -1 if the task is being drained.
+Otherwise, the number of pending calls is returned.
 .Pp
 The
 .Fn taskqueue_cancel
diff --git a/sys/kern/subr_taskqueue.c b/sys/kern/subr_taskqueue.c
index 5a20148f573c..d380b9cbd3a7 100644
--- a/sys/kern/subr_taskqueue.c
+++ b/sys/kern/subr_taskqueue.c
@@ -81,6 +81,7 @@ struct taskqueue {
 #define	TQ_FLAGS_UNLOCKED_ENQUEUE	(1 << 2)
 
 #define	DT_CALLOUT_ARMED	(1 << 0)
+#define	DT_DRAIN_IN_PROGRESS	(1 << 1)
 
 #define	TQ_LOCK(tq)							\
 	do {								\
@@ -299,7 +300,11 @@ taskqueue_enqueue_timeout(struct taskqueue *queue,
 	KASSERT(!queue->tq_spin, ("Timeout for spin-queue"));
 	timeout_task->q = queue;
 	res = timeout_task->t.ta_pending;
-	if (ticks == 0) {
+	if (timeout_task->f & DT_DRAIN_IN_PROGRESS) {
+		/* Do nothing */
+		TQ_UNLOCK(queue);
+		res = -1;
+	} else if (ticks == 0) {
 		taskqueue_enqueue_locked(queue, &timeout_task->t);
 		/* The lock is released inside. */
 	} else {
@@ -559,8 +564,24 @@ taskqueue_drain_timeout(struct taskqueue *queue,
     struct timeout_task *timeout_task)
 {
 
+	/*
+	 * Set flag to prevent timer from re-starting during drain:
+	 */
+	TQ_LOCK(queue);
+	KASSERT((timeout_task->f & DT_DRAIN_IN_PROGRESS) == 0,
+	    ("Drain already in progress"));
+	timeout_task->f |= DT_DRAIN_IN_PROGRESS;
+	TQ_UNLOCK(queue);
+
 	callout_drain(&timeout_task->c);
 	taskqueue_drain(queue, &timeout_task->t);
+
+	/*
+	 * Clear flag to allow timer to re-start:
+	 */
+	TQ_LOCK(queue);
+	timeout_task->f &= ~DT_DRAIN_IN_PROGRESS;
+	TQ_UNLOCK(queue);
 }
 
 static void