[RFC][PATCH v2 07/11] context_tracking: Add an atomic sequence/state count

@@ -9,19 +9,47 @@
 
 #include <asm/ptrace.h>
 
+enum ct_work {
+	CT_WORK_n = 0,
+};
+
+/*
+ * context_tracking::seq
+ *
+ * bit0 - work
+ * bit1 - nmi
+ * bit2 - user
+ *
+ */
+enum ct_seq_state {
+	CT_SEQ_WORK = 0x01,
+	CT_SEQ_NMI  = 0x02,
+	CT_SEQ_USER = 0x04,
+	CT_SEQ      = 0x08,
+};
+
+static __always_inline bool __context_tracking_seq_in_user(unsigned int seq)
+{
+	return (seq & (CT_SEQ_USER | CT_SEQ_NMI)) == CT_SEQ_USER;
+}
 
 #ifdef CONFIG_CONTEXT_TRACKING
+
 extern void context_tracking_cpu_set(int cpu);
 
 /* Called with interrupts disabled.  */
 extern void __context_tracking_enter(enum ctx_state state);
 extern void __context_tracking_exit(enum ctx_state state);
+extern void __context_tracking_nmi_enter(void);
+extern void __context_tracking_nmi_exit(void);
 
 extern void context_tracking_enter(enum ctx_state state);
 extern void context_tracking_exit(enum ctx_state state);
 extern void context_tracking_user_enter(void);
 extern void context_tracking_user_exit(void);
 
+extern bool context_tracking_set_cpu_work(unsigned int cpu, unsigned int work);
+
 static inline void ct_user_enter(void)
 {
 	if (context_tracking_enabled())

@@ -47,6 +75,17 @@ static __always_inline void ct_user_exit
 		__context_tracking_exit(CONTEXT_USER);
 }
 
+static __always_inline void ct_nmi_enter_irqoff(void)
+{
+	if (context_tracking_enabled())
+		__context_tracking_nmi_enter();
+}
+static __always_inline void ct_nmi_exit_irqoff(void)
+{
+	if (context_tracking_enabled())
+		__context_tracking_nmi_exit();
+}
+
 static inline enum ctx_state exception_enter(void)
 {
 	enum ctx_state prev_ctx;

@@ -97,19 +136,40 @@ static __always_inline enum ctx_state ct
 	return context_tracking_enabled() ?
 		this_cpu_read(context_tracking.state) : CONTEXT_DISABLED;
 }
+
+static __always_inline unsigned int __context_tracking_cpu_seq(unsigned int cpu)
+{
+	return arch_atomic_read(per_cpu_ptr(&context_tracking.seq, cpu));
+}
+
 #else
 static inline void ct_user_enter(void) { }
 static inline void ct_user_exit(void) { }
 static inline void ct_user_enter_irqoff(void) { }
 static inline void ct_user_exit_irqoff(void) { }
+static inline void ct_nmi_enter_irqoff(void) { }
+static inline void ct_nmi_exit_irqoff(void) { }
 static inline enum ctx_state exception_enter(void) { return 0; }
 static inline void exception_exit(enum ctx_state prev_ctx) { }
 static inline enum ctx_state ct_state(void) { return CONTEXT_DISABLED; }
 static inline bool context_tracking_guest_enter(void) { return false; }
 static inline void context_tracking_guest_exit(void) { }
 
+static inline bool context_tracking_set_cpu_work(unsigned int cpu, unsigned int work) { return false; }
+
+static __always_inline unsigned int __context_tracking_cpu_seq(unsigned int cpu)
+{
+	return 0;
+}
+
 #endif /* !CONFIG_CONTEXT_TRACKING */
 
+static __always_inline bool context_tracking_cpu_in_user(unsigned int cpu)
+{
+	unsigned int seq = __context_tracking_cpu_seq(cpu);
+	return __context_tracking_seq_in_user(seq);
+}
+
 #define CT_WARN_ON(cond) WARN_ON(context_tracking_enabled() && (cond))
 
 #ifdef CONFIG_CONTEXT_TRACKING_FORCE
--- a/include/linux/context_tracking_state.h
+++ b/include/linux/context_tracking_state.h

@@ -4,6 +4,7 @@
 
 #include <linux/percpu.h>
 #include <linux/static_key.h>
+#include <linux/types.h>
 
 struct context_tracking {
 	/*

@@ -13,6 +14,8 @@ struct context_tracking {
 	 * may be further optimized using static keys.
 	 */
 	bool active;
+	atomic_t seq;
+	atomic_t work;
 	int recursion;
 	enum ctx_state {
 		CONTEXT_DISABLED = -1,	/* returned by ct_state() if unknown */
--- a/kernel/context_tracking.c
+++ b/kernel/context_tracking.c

@@ -50,6 +50,85 @@ static __always_inline void context_trac
 	__this_cpu_dec(context_tracking.recursion);
 }
 
+/* CT_WORK_n, must be noinstr, non-blocking, NMI safe and deal with spurious calls */
+static noinstr void ct_exit_user_work(struct context_tracking *ct)
+{
+	unsigned int work = arch_atomic_read(&ct->work);
+
+#if 0
+	if (work & CT_WORK_n) {
+		/* NMI happens here and must still do/finish CT_WORK_n */
+		do_work_n();
+
+		smp_mb__before_atomic();
+		arch_atomic_andnot(CT_WORK_n, &ct->work);
+	}
+#endif
+
+	smp_mb__before_atomic();
+	arch_atomic_andnot(CT_SEQ_WORK, &ct->seq);
+}
+
+/* all CPU local */
+
+static __always_inline unsigned int ct_seq_nmi_enter(struct context_tracking *ct)
+{
+	unsigned int seq = arch_atomic_add_return(CT_SEQ_NMI, &ct->seq);
+	if (seq & CT_SEQ_WORK) /* NMI-enter is USER-exit */
+		ct_exit_user_work(ct);
+	return seq;
+}
+
+static __always_inline unsigned int ct_seq_nmi_exit(struct context_tracking *ct)
+{
+	arch_atomic_set(&ct->work, 0);
+	return arch_atomic_add_return(CT_SEQ - CT_SEQ_NMI, &ct->seq);
+}
+
+static __always_inline unsigned int ct_seq_user_enter(struct context_tracking *ct)
+{
+	arch_atomic_set(&ct->work, 0);
+	return arch_atomic_add_return(CT_SEQ_USER, &ct->seq);
+}
+
+static __always_inline unsigned int ct_seq_user_exit(struct context_tracking *ct)
+{
+	unsigned int seq = arch_atomic_add_return(CT_SEQ - CT_SEQ_USER, &ct->seq);
+	if (seq & CT_SEQ_WORK)
+		ct_exit_user_work(ct);
+	return seq;
+}
+
+/* remote */
+
+/*
+ * When returns true: guaratees that CPu will call @work
+ */
+static bool ct_seq_set_user_work(struct context_tracking *ct, unsigned int work)
+{
+	unsigned int seq;
+	bool ret = false;
+
+	if (!context_tracking_enabled() || !ct->active)
+		return false;
+
+	preempt_disable();
+	seq = atomic_read(&ct->seq);
+	if (__context_tracking_seq_in_user(seq)) {
+		/* ctrl-dep */
+		atomic_or(work, &ct->work);
+		ret = atomic_try_cmpxchg(&ct->seq, &seq, seq|CT_SEQ_WORK);
+	}
+	preempt_enable();
+
+	return ret;
+}
+
+bool context_tracking_set_cpu_work(unsigned int cpu, unsigned int work)
+{
+	return ct_seq_set_user_work(per_cpu_ptr(&context_tracking, cpu), work);
+}
+
 /**
  * context_tracking_enter - Inform the context tracking that the CPU is going
  *                          enter user or guest space mode.

@@ -83,6 +162,7 @@ void noinstr __context_tracking_enter(en
 				instrumentation_end();
 			}
 			rcu_user_enter();
+			ct_seq_user_enter(raw_cpu_ptr(&context_tracking));
 		}
 		/*
 		 * Even if context tracking is disabled on this CPU, because it's outside

@@ -154,6 +234,7 @@ void noinstr __context_tracking_exit(enu
 			 * We are going to run code that may use RCU. Inform
 			 * RCU core about that (ie: we may need the tick again).
 			 */
+			ct_seq_user_exit(raw_cpu_ptr(&context_tracking));
 			rcu_user_exit();
 			if (state == CONTEXT_USER) {
 				instrumentation_begin();

@@ -188,6 +269,18 @@ void context_tracking_user_exit(void)
 }
 NOKPROBE_SYMBOL(context_tracking_user_exit);
 
+void noinstr __context_tracking_nmi_enter(void)
+{
+	if (__this_cpu_read(context_tracking.active))
+		ct_seq_nmi_enter(raw_cpu_ptr(&context_tracking));
+}
+
+void noinstr __context_tracking_nmi_exit(void)
+{
+	if (__this_cpu_read(context_tracking.active))
+		ct_seq_nmi_exit(raw_cpu_ptr(&context_tracking));
+}
+
 void __init context_tracking_cpu_set(int cpu)
 {
 	static __initdata bool initialized = false;
--- a/kernel/entry/common.c
+++ b/kernel/entry/common.c

@@ -451,6 +451,7 @@ irqentry_state_t noinstr irqentry_nmi_en
 	__nmi_enter();
 	lockdep_hardirqs_off(CALLER_ADDR0);
 	lockdep_hardirq_enter();
+	ct_nmi_enter_irqoff();
 	rcu_nmi_enter();
 
 	instrumentation_begin();

@@ -472,6 +473,7 @@ void noinstr irqentry_nmi_exit(struct pt
 	instrumentation_end();
 
 	rcu_nmi_exit();
+	ct_nmi_exit_irqoff();
 	lockdep_hardirq_exit();
 	if (irq_state.lockdep)
 		lockdep_hardirqs_on(CALLER_ADDR0);