tracing: Have the snapshot trigger use the mapping helper functions

As the data pointer for individual ips will soon be removed and no longer
passed to the callback function probe handlers, convert the snapshot
trigger counter over to the new ftrace_func_mapper helper functions.
Signed-off-by: Steven Rostedt (VMware) <rostedt@goodmis.org>

kernel/trace/trace.c

@@ -6745,13 +6745,19 @@ static void
 ftrace_count_snapshot(unsigned long ip, unsigned long parent_ip,
 		      struct ftrace_probe_ops *ops, void **data)
 {
-	unsigned long *count = (long *)data;
+	struct ftrace_func_mapper *mapper = ops->private_data;
+	long *count = NULL;
 
-	if (!*count)
-		return;
+	if (mapper)
+		count = (long *)ftrace_func_mapper_find_ip(mapper, ip);
+
+	if (count) {
+
+		if (*count <= 0)
+			return;
 
-	if (*count != -1)
 		(*count)--;
+	}
 
 	tracing_snapshot();
 }
@@ -6760,20 +6766,42 @@ static int
 ftrace_snapshot_print(struct seq_file *m, unsigned long ip,
 		      struct ftrace_probe_ops *ops, void *data)
 {
-	long count = (long)data;
+	struct ftrace_func_mapper *mapper = ops->private_data;
+	long *count = NULL;
 
 	seq_printf(m, "%ps:", (void *)ip);
 
 	seq_puts(m, "snapshot");
 
-	if (count == -1)
-		seq_puts(m, ":unlimited\n");
+	if (mapper)
+		count = (long *)ftrace_func_mapper_find_ip(mapper, ip);
+
+	if (count)
+		seq_printf(m, ":count=%ld\n", *count);
 	else
-		seq_printf(m, ":count=%ld\n", count);
+		seq_puts(m, ":unlimited\n");
 
 	return 0;
 }
 
+static int
+ftrace_snapshot_init(struct ftrace_probe_ops *ops, unsigned long ip,
+		     void **data)
+{
+	struct ftrace_func_mapper *mapper = ops->private_data;
+
+	return ftrace_func_mapper_add_ip(mapper, ip, *data);
+}
+
+static void
+ftrace_snapshot_free(struct ftrace_probe_ops *ops, unsigned long ip,
+		     void **_data)
+{
+	struct ftrace_func_mapper *mapper = ops->private_data;
+
+	ftrace_func_mapper_remove_ip(mapper, ip);
+}
+
 static struct ftrace_probe_ops snapshot_probe_ops = {
 	.func			= ftrace_snapshot,
 	.print			= ftrace_snapshot_print,
@@ -6782,6 +6810,8 @@ static struct ftrace_probe_ops snapshot_probe_ops = {
 static struct ftrace_probe_ops snapshot_count_probe_ops = {
 	.func			= ftrace_count_snapshot,
 	.print			= ftrace_snapshot_print,
+	.init			= ftrace_snapshot_init,
+	.free			= ftrace_snapshot_free,
 };
 
 static int
@@ -6812,6 +6842,12 @@ ftrace_trace_snapshot_callback(struct ftrace_hash *hash,
 	if (!strlen(number))
 		goto out_reg;
 
+	if (!ops->private_data) {
+		ops->private_data = allocate_ftrace_func_mapper();
+		if (!ops->private_data)
+			return -ENOMEM;
+	}
+
 	/*
 	 * We use the callback data field (which is a pointer)
 	 * as our counter.