1 files changed, 87 insertions, 7 deletions

diff --git a/module/zfs/zthr.c b/module/zfs/zthr.c
index 532e8ce0f..53c0a0b3d 100644
--- a/module/zfs/zthr.c
+++ b/module/zfs/zthr.c
@@ -207,12 +207,15 @@ struct zthr {
 	/* flag set to true if we are canceling the zthr */
 	boolean_t	zthr_cancel;
 
+	/* flag set to true if we are waiting for the zthr to finish */
+	boolean_t	zthr_haswaiters;
+	kcondvar_t	zthr_wait_cv;
 	/*
 	 * maximum amount of time that the zthr is spent sleeping;
 	 * if this is 0, the thread doesn't wake up until it gets
 	 * signaled.
 	 */
-	hrtime_t	zthr_wait_time;
+	hrtime_t	zthr_sleep_timeout;
 
 	/* consumer-provided callbacks & data */
 	zthr_checkfunc_t	*zthr_checkfunc;
@@ -239,14 +242,18 @@ zthr_procedure(void *arg)
 			 * order to prevent this process from incorrectly
 			 * contributing to the system load average when idle.
 			 */
-			if (t->zthr_wait_time == 0) {
+			if (t->zthr_sleep_timeout == 0) {
 				cv_wait_sig(&t->zthr_cv, &t->zthr_state_lock);
 			} else {
 				(void) cv_timedwait_sig_hires(&t->zthr_cv,
-				    &t->zthr_state_lock, t->zthr_wait_time,
+				    &t->zthr_state_lock, t->zthr_sleep_timeout,
 				    MSEC2NSEC(1), 0);
 			}
 		}
+		if (t->zthr_haswaiters) {
+			t->zthr_haswaiters = B_FALSE;
+			cv_broadcast(&t->zthr_wait_cv);
+		}
 	}
 
 	/*
@@ -280,12 +287,13 @@ zthr_create_timer(zthr_checkfunc_t *checkfunc, zthr_func_t *func,
 	mutex_init(&t->zthr_state_lock, NULL, MUTEX_DEFAULT, NULL);
 	mutex_init(&t->zthr_request_lock, NULL, MUTEX_DEFAULT, NULL);
 	cv_init(&t->zthr_cv, NULL, CV_DEFAULT, NULL);
+	cv_init(&t->zthr_wait_cv, NULL, CV_DEFAULT, NULL);
 
 	mutex_enter(&t->zthr_state_lock);
 	t->zthr_checkfunc = checkfunc;
 	t->zthr_func = func;
 	t->zthr_arg = arg;
-	t->zthr_wait_time = max_sleep;
+	t->zthr_sleep_timeout = max_sleep;
 
 	t->zthr_thread = thread_create(NULL, 0, zthr_procedure, t,
 	    0, &p0, TS_RUN, minclsyspri);
@@ -303,6 +311,7 @@ zthr_destroy(zthr_t *t)
 	mutex_destroy(&t->zthr_request_lock);
 	mutex_destroy(&t->zthr_state_lock);
 	cv_destroy(&t->zthr_cv);
+	cv_destroy(&t->zthr_wait_cv);
 	kmem_free(t, sizeof (*t));
 }
 
@@ -355,9 +364,8 @@ zthr_cancel(zthr_t *t)
 	 *
 	 * [1] The thread has already been cancelled, therefore
 	 *     there is nothing for us to do.
-	 * [2] The thread is sleeping, so we broadcast the CV first
-	 *     to wake it up and then we set the flag and we are
-	 *     waiting for it to exit.
+	 * [2] The thread is sleeping so we set the flag, broadcast
+	 *     the CV and wait for it to exit.
 	 * [3] The thread is doing work, in which case we just set
 	 *     the flag and wait for it to finish.
 	 * [4] The thread was just created/resumed, in which case
@@ -397,6 +405,7 @@ zthr_resume(zthr_t *t)
 	ASSERT3P(&t->zthr_checkfunc, !=, NULL);
 	ASSERT3P(&t->zthr_func, !=, NULL);
 	ASSERT(!t->zthr_cancel);
+	ASSERT(!t->zthr_haswaiters);
 
 	/*
 	 * There are 4 states that we find the zthr in at this point
@@ -451,3 +460,74 @@ zthr_iscancelled(zthr_t *t)
 	mutex_exit(&t->zthr_state_lock);
 	return (cancelled);
 }
+
+/*
+ * Wait for the zthr to finish its current function. Similar to
+ * zthr_iscancelled, you can use zthr_has_waiters to have the zthr_func end
+ * early. Unlike zthr_cancel, the thread is not destroyed. If the zthr was
+ * sleeping or cancelled, return immediately.
+ */
+void
+zthr_wait_cycle_done(zthr_t *t)
+{
+	mutex_enter(&t->zthr_state_lock);
+
+	/*
+	 * Since we are holding the zthr_state_lock at this point
+	 * we can find the state in one of the following 5 states:
+	 *
+	 * [1] The thread has already cancelled, therefore
+	 *     there is nothing for us to do.
+	 * [2] The thread is sleeping so we set the flag, broadcast
+	 *     the CV and wait for it to exit.
+	 * [3] The thread is doing work, in which case we just set
+	 *     the flag and wait for it to finish.
+	 * [4] The thread was just created/resumed, in which case
+	 *     the behavior is similar to [3].
+	 * [5] The thread is the middle of being cancelled, which is
+	 *     similar to [3]. We'll wait for the cancel, which is
+	 *     waiting for the zthr func.
+	 *
+	 * Since requests are serialized, by the time that we get
+	 * control back we expect that the zthr has completed it's
+	 * zthr_func.
+	 */
+	if (t->zthr_thread != NULL) {
+		t->zthr_haswaiters = B_TRUE;
+
+		/* broadcast in case the zthr is sleeping */
+		cv_broadcast(&t->zthr_cv);
+
+		while ((t->zthr_haswaiters) && (t->zthr_thread != NULL))
+			cv_wait(&t->zthr_wait_cv, &t->zthr_state_lock);
+
+		ASSERT(!t->zthr_haswaiters);
+	}
+
+	mutex_exit(&t->zthr_state_lock);
+}
+
+/*
+ * This function is intended to be used by the zthr itself
+ * to check if another thread is waiting on it to finish
+ *
+ * returns TRUE if we have been asked to finish.
+ *
+ * returns FALSE otherwise.
+ */
+boolean_t
+zthr_has_waiters(zthr_t *t)
+{
+	ASSERT3P(t->zthr_thread, ==, curthread);
+
+	mutex_enter(&t->zthr_state_lock);
+
+	/*
+	 * Similarly to zthr_iscancelled(), we only grab the
+	 * zthr_state_lock so that the zthr itself can use this
+	 * to check for the request.
+	 */
+	boolean_t has_waiters = t->zthr_haswaiters;
+	mutex_exit(&t->zthr_state_lock);
+	return (has_waiters);
+}