From 14889b422910d683b06059f5b0eb340e4cc590d4 Mon Sep 17 00:00:00 2001
From: Warner Losh <imp@FreeBSD.org>
Date: Tue, 5 Oct 2004 04:16:01 +0000
Subject: [PATCH] Add taskqueue_drain.  This waits for the specified task to
 finish, if running, or returns.  The calling program is responsible for
 making sure that nothing new is enqueued.

# man page coming soon.
---
 sys/kern/subr_taskqueue.c | 14 ++++++++++++++
 sys/sys/_task.h           |  3 +++
 sys/sys/taskqueue.h       |  2 ++
 3 files changed, 19 insertions(+)

diff --git a/sys/kern/subr_taskqueue.c b/sys/kern/subr_taskqueue.c
index a10e7107bd89..7758fbdb1a27 100644
--- a/sys/kern/subr_taskqueue.c
+++ b/sys/kern/subr_taskqueue.c
@@ -186,11 +186,14 @@ taskqueue_run(struct taskqueue *queue)
 		STAILQ_REMOVE_HEAD(&queue->tq_queue, ta_link);
 		pending = task->ta_pending;
 		task->ta_pending = 0;
+		task->ta_flags |= TAF_PENDING;
 		mtx_unlock(&queue->tq_mutex);
 
 		task->ta_func(task->ta_context, pending);
 
 		mtx_lock(&queue->tq_mutex);
+		task->ta_flags &= ~TAF_PENDING;
+		wakeup(task);
 	}
 
 	/*
@@ -201,6 +204,17 @@ taskqueue_run(struct taskqueue *queue)
 		mtx_unlock(&queue->tq_mutex);
 }
 
+void
+taskqueue_drain(struct taskqueue *queue, struct task *task)
+{
+	WITNESS_WARN(WARN_GIANTOK | WARN_SLEEPOK, NULL, "taskqueue_drain");
+	mtx_lock(&queue->tq_mutex);
+	while (task->ta_pending != 0 || (task->ta_flags & TAF_PENDING)) {
+		msleep(task, &queue->tq_mutex, PWAIT, "-", 0);
+	}
+	mtx_unlock(&queue->tq_mutex);
+}
+
 static void
 taskqueue_swi_enqueue(void *context)
 {
diff --git a/sys/sys/_task.h b/sys/sys/_task.h
index f459c2370ac5..4eb6d644d450 100644
--- a/sys/sys/_task.h
+++ b/sys/sys/_task.h
@@ -45,6 +45,9 @@ struct task {
 	int	ta_priority;		/* priority of task in queue */
 	task_fn_t *ta_func;		/* task handler */
 	void	*ta_context;		/* argument for handler */
+	int	ta_flags;		/* Flags */
 };
 
+#define TAF_PENDING	0x1		/* Task is being run now */
+
 #endif /* !_SYS__TASK_H_ */
diff --git a/sys/sys/taskqueue.h b/sys/sys/taskqueue.h
index 355f1ce562ae..31bd9f27b242 100644
--- a/sys/sys/taskqueue.h
+++ b/sys/sys/taskqueue.h
@@ -51,6 +51,7 @@ struct taskqueue *taskqueue_create(const char *name, int mflags,
 				    taskqueue_enqueue_fn enqueue,
 				    void *context);
 int	taskqueue_enqueue(struct taskqueue *queue, struct task *task);
+void	taskqueue_drain(struct taskqueue *queue, struct task *task);
 struct taskqueue *taskqueue_find(const char *name);
 void	taskqueue_free(struct taskqueue *queue);
 void	taskqueue_run(struct taskqueue *queue);
@@ -69,6 +70,7 @@ void	taskqueue_thread_enqueue(void *context);
 	(task)->ta_priority = (priority);		\
 	(task)->ta_func = (func);			\
 	(task)->ta_context = (context);			\
+	(task)->ta_flags = 0;				\
 } while (0)
 
 /*