From 66283521fac33ac5716b71a6d2662093aec409c1 Mon Sep 17 00:00:00 2001
From: David Smith <dsmith@redhat.com>
Date: Wed, 12 Mar 2014 09:07:58 -0500
Subject: [PATCH] Improve task_finder/utrace shutdown.

* runtime/stp_utrace.c (utrace_exit): Move the call to
  stp_task_work_exit() up above the calls to free the kmem caches. This
  make sure any running task work items don't have the memory freed out
  from under them.
* runtime/linux/task_finder2.c (__stp_task_finder_cleanup): Move
  utrace_shutdown() call into stap_stop_task_finder().
  (stap_stop_task_finder): Call utrace_shutdown() directly. Wait to make
  sure all tracepoint probes are finished.
---
 runtime/linux/task_finder2.c | 24 +++++++-----------------
 runtime/stp_utrace.c         | 27 +++++++++++++++++++--------
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/runtime/linux/task_finder2.c b/runtime/linux/task_finder2.c
index a3abc16..ef074c5 100644
--- a/runtime/linux/task_finder2.c
+++ b/runtime/linux/task_finder2.c
@@ -418,15 +418,6 @@ stap_utrace_detach_ops(struct utrace_engine_ops *ops)
 	debug_task_finder_report();
 }
 
-static void
-__stp_task_finder_cleanup(void)
-{
-	// The utrace_shutdown() function detaches and deletes
-	// everything for us - we don't have to go through each
-	// engine.
-	utrace_shutdown();
-}
-
 static char *
 __stp_get_mm_path(struct mm_struct *mm, char *buf, int buflen)
 {
@@ -1849,16 +1840,16 @@ stap_stop_task_finder(void)
 	atomic_set(&__stp_task_finder_state, __STP_TF_STOPPING);
 
 	debug_task_finder_report();
-#if 0
-	/* We don't need this since __stp_task_finder_cleanup()
-	 * removes everything by calling utrace_shutdown(). */
-	stap_utrace_detach_ops(&__stp_utrace_task_finder_ops);
-#endif
-	__stp_task_finder_cleanup();
+
+	// The utrace_shutdown() function detaches and cleans up
+	// everything for us - we don't have to go through each
+	// engine. This also means that the attach_count could end up
+	// > 0 (since we don't got through each engine individually).
+	utrace_shutdown();
+
 	debug_task_finder_report();
 	atomic_set(&__stp_task_finder_state, __STP_TF_STOPPED);
 
-#if 0
 	/* Now that all the engines are detached, make sure
 	 * all the callbacks are finished.  If they aren't, we'll
 	 * crash the kernel when the module is removed. */
@@ -1873,7 +1864,6 @@ stap_stop_task_finder(void)
 		printk(KERN_ERR "it took %d polling loops to quit.\n", i);
 #endif
 	debug_task_finder_report();
-#endif
 
 	/* Make sure all outstanding task work requests are canceled. */
 	__stp_tf_cancel_task_work();
diff --git a/runtime/stp_utrace.c b/runtime/stp_utrace.c
index 89ea0e4..a6f363d 100644
--- a/runtime/stp_utrace.c
+++ b/runtime/stp_utrace.c
@@ -286,12 +286,18 @@ static int utrace_exit(void)
 {
 	utrace_shutdown();
 
-	if (utrace_cachep)
-		kmem_cache_destroy(utrace_cachep);
-	if (utrace_engine_cachep)
-		kmem_cache_destroy(utrace_engine_cachep);
-
 	stp_task_work_exit();
+
+	/* After utrace_shutdown() and stp_task_work_exit() (and the
+	 * code in stap_stop_task_finder()), we're *sure* there are no
+	 * tracepoint probes or task work items running or scheduled
+	 * to be run. So, now would be a great time to actually free
+	 * everything. */
+
+	if (utrace_cachep)
+		kmem_cache_destroy(utrace_cachep);
+	if (utrace_engine_cachep)
+		kmem_cache_destroy(utrace_engine_cachep);
 	return 0;
 }
 
@@ -373,9 +379,14 @@ static void utrace_shutdown(void)
 	unregister_trace_sys_exit(utrace_report_syscall_exit, NULL);
 	tracepoint_synchronize_unregister();
 
-	/* After calling tracepoint_synchronize_unregister(), we're
-	 * sure there are no outstanding tracepoint probes being
-	 * called.  So, now would be a great time to free everything. */
+	/* After tracepoint_synchronize_unregister(), we're *sure*
+	 * there will be no new tracepoint probes running. There could
+	 * be currently running tracepoint probes or task work
+	 * items. So, now would be a great time to cleanup.
+	 *
+	 * Currently running items should be OK, since
+	 * utrace_cleanup() just puts the memory back into the utrace
+	 * kmem caches. */
 	
 #ifdef STP_TF_DEBUG
 	printk(KERN_ERR "%s:%d - freeing task-specific\n", __FUNCTION__, __LINE__);
-- 
1.8.3.1