[FIX] stop message handling in us_manager

stop_machine is now used only to acquire references to all the
instrumented tasks. In such way we assure that they are not going
to be dropped during SWAP cleanup code execution.

Also a special callback STOP_CB_TD was added for task_data since it
may need to cleanup in stop_machine.

Change-Id: I4ea67092e1f0a006c908392e41ae511b9fff5b6d
Signed-off-by: Vyacheslav Cherkashin <v.cherkashin@samsung.com>

7 files changed, 91 insertions, 16 deletions

diff --git a/task_data/task_data.c b/task_data/task_data.c
index 085e0163..6f31155c 100644
--- a/task_data/task_data.c
+++ b/task_data/task_data.c
@@ -208,7 +208,7 @@ static int task_data_init(void)
 		goto out;
 	}
 
-	__task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB, task_data_stop);
+	__task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB_TD, task_data_stop);
 
 	if (__task_data_cbs_stop_h < 0) {
 		ret = __task_data_cbs_stop_h;
diff --git a/us_manager/callbacks.h b/us_manager/callbacks.h
index a952af6b..f6d1b273 100644
--- a/us_manager/callbacks.h
+++ b/us_manager/callbacks.h
@@ -3,7 +3,8 @@
 
 enum callback_t {
 	START_CB = 0,
-	STOP_CB
+	STOP_CB,
+	STOP_CB_TD
 };
 
 /* Gets callback type (on start or on stop) and function pointer.
diff --git a/us_manager/pf/pf_group.c b/us_manager/pf/pf_group.c
index 0057636b..fb92fdb5 100644
--- a/us_manager/pf/pf_group.c
+++ b/us_manager/pf/pf_group.c
@@ -572,12 +572,7 @@ void uninstall_proc(struct sspt_proc *proc)
 {
 	struct task_struct *task = proc->task;
 
-
-	task_lock(task);
-	BUG_ON(task->mm == NULL);
 	sspt_proc_uninstall(proc, task, US_UNREGS_PROBE);
-	task_unlock(task);
-
 	sspt_proc_cleanup(proc);
 }
 
@@ -622,11 +617,6 @@ void install_all(void)
 	/* TODO: to be implemented */
 }
 
-static void on_each_uninstall_proc(struct sspt_proc *proc, void *data)
-{
-	uninstall_proc(proc);
-}
-
 /**
  * @brief Uninstall probes from all processes
  *
@@ -634,11 +624,56 @@ static void on_each_uninstall_proc(struct sspt_proc *proc, void *data)
  */
 void uninstall_all(void)
 {
+	struct list_head *proc_list = sspt_proc_list();
+
 	sspt_proc_write_lock();
-	on_each_proc_no_lock(on_each_uninstall_proc, NULL);
+	while (!list_empty(proc_list)) {
+		struct sspt_proc *proc;
+		proc = list_first_entry(proc_list, struct sspt_proc, list);
+
+		list_del(&proc->list);
+
+		sspt_proc_write_unlock();
+		uninstall_proc(proc);
+		sspt_proc_write_lock();
+	}
 	sspt_proc_write_unlock();
 }
 
+static void __do_get_proc(struct sspt_proc *proc, void *data)
+{
+	get_task_struct(proc->task);
+	proc->__task = proc->task;
+	proc->__mm = get_task_mm(proc->task);
+}
+
+static void __do_put_proc(struct sspt_proc *proc, void *data)
+{
+	if (proc->__mm) {
+		mmput(proc->__mm);
+		proc->__mm = NULL;
+	}
+
+	if (proc->__task) {
+		put_task_struct(proc->__task);
+		proc->__task = NULL;
+	}
+}
+
+void get_all_procs(void)
+{
+	sspt_proc_read_lock();
+	on_each_proc_no_lock(__do_get_proc, NULL);
+	sspt_proc_read_unlock();
+}
+
+void put_all_procs(void)
+{
+	sspt_proc_read_lock();
+	on_each_proc_no_lock(__do_put_proc, NULL);
+	sspt_proc_read_unlock();
+}
+
 /**
  * @brief For debug
  *
diff --git a/us_manager/pf/pf_group.h b/us_manager/pf/pf_group.h
index 00f71c1b..5277ed75 100644
--- a/us_manager/pf/pf_group.h
+++ b/us_manager/pf/pf_group.h
@@ -63,6 +63,9 @@ int pf_unregister_probe(struct pf_group *pfg, struct dentry *dentry,
 void install_all(void);
 void uninstall_all(void);
 
+void get_all_procs(void);
+void put_all_procs(void);
+
 int check_task_on_filters(struct task_struct *task);
 void call_page_fault(struct task_struct *task, unsigned long page_addr);
 void call_mm_release(struct task_struct *task);
diff --git a/us_manager/sspt/sspt_proc.c b/us_manager/sspt/sspt_proc.c
index 85d24f4e..d9473a8b 100644
--- a/us_manager/sspt/sspt_proc.c
+++ b/us_manager/sspt/sspt_proc.c
@@ -38,6 +38,11 @@ static LIST_HEAD(proc_probes_list);
 static DEFINE_RWLOCK(sspt_proc_rwlock);
 
 
+struct list_head *sspt_proc_list()
+{
+	return &proc_probes_list;
+}
+
 /**
  * @brief Global read lock for sspt_proc
  *
@@ -146,8 +151,17 @@ struct sspt_proc *sspt_proc_get(struct sspt_proc *proc)
 
 void sspt_proc_put(struct sspt_proc *proc)
 {
-	if (atomic_dec_and_test(&proc->usage))
+	if (atomic_dec_and_test(&proc->usage)) {
+		if (proc->__mm) {
+			mmput(proc->__mm);
+			proc->__mm = NULL;
+		}
+		if (proc->__task) {
+			put_task_struct(proc->__task);
+			proc->__task = NULL;
+		}
 		kfree(proc);
+	}
 }
 
 /**
diff --git a/us_manager/sspt/sspt_proc.h b/us_manager/sspt/sspt_proc.h
index d82a2712..46b9e44e 100644
--- a/us_manager/sspt/sspt_proc.h
+++ b/us_manager/sspt/sspt_proc.h
@@ -48,6 +48,8 @@ struct sspt_proc {
 	struct list_head list;		/**< For global process list */
 	pid_t tgid;			/**< Thread group ID */
 	struct task_struct *task;	/**< Ptr to the task */
+	struct mm_struct *__mm;
+	struct task_struct *__task;
 	unsigned long r_state_addr;	/**< address of r_state */
 	struct slot_manager *sm;	/**< Ptr to the manager slot */
 	struct list_head file_list;	/**< For sspt_file */
@@ -65,6 +67,9 @@ struct sspt_proc_cb {
 	void (*priv_destroy)(struct sspt_proc *, void *);
 };
 
+
+struct list_head *sspt_proc_list(void);
+
 struct sspt_proc *sspt_proc_create(struct task_struct *task);
 void sspt_proc_cleanup(struct sspt_proc *proc);
 struct sspt_proc *sspt_proc_get(struct sspt_proc *proc);
diff --git a/us_manager/us_manager.c b/us_manager/us_manager.c
index 4328e447..c306f9d8 100644
--- a/us_manager/us_manager.c
+++ b/us_manager/us_manager.c
@@ -24,6 +24,7 @@
 
 #include <linux/module.h>
 #include <linux/mutex.h>
+#include <linux/stop_machine.h>
 #include "pf/pf_group.h"
 #include "sspt/sspt_proc.h"
 #include "probes/probe_info_new.h"
@@ -39,14 +40,30 @@ static DEFINE_MUTEX(mutex_inst);
 static enum status_type status = ST_OFF;
 
 
-static void do_usm_stop(void)
+static int __do_usm_stop(void *data)
 {
-	exec_cbs(STOP_CB);
+	get_all_procs();
+
+	return 0;
+}
+
+static int do_usm_stop(void)
+{
+	int ret;
 
+	exec_cbs(STOP_CB);
 	unregister_helper_top();
+
+	ret = stop_machine(__do_usm_stop, NULL, NULL);
+	if (ret)
+		printk("do_usm_stop failed: %d\n", ret);
+
 	uninstall_all();
 	unregister_helper_bottom();
 	sspt_proc_free_all();
+	exec_cbs(STOP_CB_TD);
+
+	return ret;
 }
 
 static int do_usm_start(void)