@@ -53,6 +53,8 @@
 #include <linux/io.h> /* vmap_page_range() */
 #include <linux/fs_context.h>
 
+#include <trace/events/sched.h>
+
 #include <asm/setup.h> /* COMMAND_LINE_SIZE */
 
 #include "trace.h"

@@ -5984,6 +5986,32 @@ struct trace_user_buf {
 static DEFINE_MUTEX(trace_user_buffer_mutex);
 static struct trace_user_buf_info *trace_user_buffer;
 
+static DEFINE_PER_CPU(unsigned long, sched_switch_cnt);
+
+/*
+ * The per CPU buffer trace_user_buffer is written to optimstically.
+ * The counter sched_switch_cnt is taken, preemption is enabled,
+ * the copying of the user space memory is placed into the trace_user_buffer,
+ * Preeption is re-enabled and the count is read again. If the count is greater
+ * than one from its previous reading, it means that another user space
+ * task scheduled in and the buffer is unreliable for use.
+ */
+static void
+probe_sched_switch(void *ignore, bool preempt,
+		   struct task_struct *prev, struct task_struct *next,
+		   unsigned int prev_state)
+{
+	/*
+	 * The buffer can only be corrupted by another user space task.
+	 * Ignore kernel tasks that may be scheduled in order to process
+	 * the faulting memory.
+	 */
+	if (!is_user_task(next))
+		return;
+
+	this_cpu_inc(sched_switch_cnt);
+}
+
 /**
  * trace_user_fault_destroy - free up allocated memory of a trace user buffer
  * @tinfo: The descriptor to free up

@@ -6003,6 +6031,8 @@ void trace_user_fault_destroy(struct trace_user_buf_info *tinfo)
 		kfree(buf);
 	}
 	free_percpu(tinfo->tbuf);
+
+	unregister_trace_sched_switch(probe_sched_switch, NULL);
 }
 
 static int user_fault_buffer_enable(struct trace_user_buf_info *tinfo, size_t size)

@@ -6053,11 +6083,17 @@ static int user_buffer_init(struct trace_user_buf_info **tinfo, size_t size)
 
 	lockdep_assert_held(&trace_user_buffer_mutex);
 
+	ret = register_trace_sched_switch(probe_sched_switch, NULL);
+	if (ret < 0)
+		return ret;
+
 	if (!*tinfo) {
 		alloc = true;
 		*tinfo = kzalloc_obj(**tinfo);
-		if (!*tinfo)
+		if (!*tinfo) {
+			unregister_trace_sched_switch(probe_sched_switch, NULL);
 			return -ENOMEM;
+		}
 	}
 
 	ret = user_fault_buffer_enable(*tinfo, size);

@@ -6241,7 +6277,7 @@ char *trace_user_fault_read(struct trace_user_buf_info *tinfo,
 			return NULL;
 
 		/* Read the current CPU context switch counter */
-		cnt = nr_context_switches_cpu(cpu);
+		cnt = this_cpu_read(sched_switch_cnt);
 
 		/*
 		 * Preemption is going to be enabled, but this task must

@@ -6272,12 +6308,19 @@ char *trace_user_fault_read(struct trace_user_buf_info *tinfo,
 			return NULL;
 
 		/*
-		 * Preemption is disabled again, now check the per CPU context
-		 * switch counter. If it doesn't match, then another user space
-		 * process may have schedule in and corrupted our buffer. In that
-		 * case the copying must be retried.
+		 * Preemption is disabled again, now check the sched_switch_cnt.
+		 * If it increased by two or more, then another user space process
+		 * may have schedule in and corrupted our buffer. In that case
+		 * the copying must be retried.
+		 *
+		 * Note, if this task was scheduled out and only kernel threads
+		 * were scheduled in (maybe to process the fault), then the
+		 * counter would increment again when this task scheduled in.
+		 * If this task scheduled out and another user task scheduled
+		 * in, this task would still need to be scheduled back in and
+		 * the counter would increment by at least two.
 		 */
-	} while (nr_context_switches_cpu(cpu) != cnt);
+	} while (this_cpu_read(sched_switch_cnt) > cnt + 1);
 
 	return buffer;
 }