[PATCH kvm-unit-tests v2 3/8] lib/x86/smp: introduce on_cpus and cpus_active

[Andrew Jones <drjones@xxxxxxxxxx>]

Signed-off-by: Andrew Jones <drjones@xxxxxxxxxx>
---
 lib/x86/smp.c | 21 +++++++++++++++++++++
 lib/x86/smp.h |  2 ++
 2 files changed, 23 insertions(+)

diff --git a/lib/x86/smp.c b/lib/x86/smp.c
index 4bdbeaeb8b68..bffb6dec0c65 100644
--- a/lib/x86/smp.c
+++ b/lib/x86/smp.c
@@ -1,5 +1,7 @@
 
 #include <libcflat.h>
+#include "processor.h"
+#include "atomic.h"
 #include "smp.h"
 #include "apic.h"
 #include "fwcfg.h"
@@ -15,6 +17,7 @@ static void *volatile ipi_data;
 static volatile int ipi_done;
 static volatile bool ipi_wait;
 static int _cpu_count;
+static atomic_t active_cpus;
 
 static __attribute__((used)) void ipi()
 {
@@ -27,6 +30,7 @@ static __attribute__((used)) void ipi()
 	apic_write(APIC_EOI, 0);
     }
     function(data);
+    atomic_dec(&active_cpus);
     if (wait) {
 	ipi_done = 1;
 	apic_write(APIC_EOI, 0);
@@ -68,6 +72,7 @@ static void __on_cpu(int cpu, void (*function)(void *data), void *data,
     if (cpu == smp_id())
 	function(data);
     else {
+	atomic_inc(&active_cpus);
 	ipi_done = 0;
 	ipi_function = function;
 	ipi_data = data;
@@ -91,6 +96,21 @@ void on_cpu_async(int cpu, void (*function)(void *data), void *data)
     __on_cpu(cpu, function, data, 0);
 }
 
+void on_cpus(void (*function)(void *data), void *data)
+{
+    int cpu;
+
+    for (cpu = cpu_count() - 1; cpu >= 0; --cpu)
+        on_cpu_async(cpu, function, data);
+
+    while (cpus_active() > 1)
+        pause();
+}
+
+int cpus_active(void)
+{
+    return atomic_read(&active_cpus);
+}
 
 void smp_init(void)
 {
@@ -106,4 +126,5 @@ void smp_init(void)
     for (i = 1; i < cpu_count(); ++i)
         on_cpu(i, setup_smp_id, 0);
 
+    atomic_inc(&active_cpus);
 }
diff --git a/lib/x86/smp.h b/lib/x86/smp.h
index afabac8495f1..1453bb5e6805 100644
--- a/lib/x86/smp.h
+++ b/lib/x86/smp.h
@@ -6,7 +6,9 @@ void smp_init(void);
 
 int cpu_count(void);
 int smp_id(void);
+int cpus_active(void);
 void on_cpu(int cpu, void (*function)(void *data), void *data);
 void on_cpu_async(int cpu, void (*function)(void *data), void *data);
+void on_cpus(void (*function)(void *data), void *data);
 
 #endif
-- 
2.9.4