From 7b8cbd9cf0c618bf06068b19eaa7c1c9e43a97f6 Mon Sep 17 00:00:00 2001
From: Iago Abal <mail@iagoabal.eu>
Date: Mon, 5 Jul 2021 11:34:38 +0200
Subject: [PATCH] Avoid looping forever if core pinning fails

This is known to happen, under some circumstances, on Apple M1.

Plus code cleanup.

Closes #103
---
 src/setcore_stubs.c | 63 +++++++++++++++++++++------------------------
 1 file changed, 30 insertions(+), 33 deletions(-)

diff --git a/src/setcore_stubs.c b/src/setcore_stubs.c
index 6c7cedd..d8e031a 100644
--- a/src/setcore_stubs.c
+++ b/src/setcore_stubs.c
@@ -18,47 +18,44 @@ CAMLprim value numcores(value unit) {
   return Val_int(numcores);
 }
 
-CAMLprim value setcore(value which) {
-  int numcores = sysconf( _SC_NPROCESSORS_ONLN );
-  int w = Int_val(which) % numcores; // stay in the space of existing cores
+int setcorew(int w) {
 #if HAVE_DECL_SCHED_SETAFFINITY
-  cpu_set_t cpus;   
-#endif
-#if HAVE_MACH_THREAD_POLICY_H
+  cpu_set_t cpus;
+  CPU_ZERO(&cpus); 
+  CPU_SET (w,&cpus);
+  //fprintf(stderr,"Trying to pin to cpu %d out of %d reported by the system\n",w,numcores);
+  return sched_setaffinity(getpid(), sizeof(cpu_set_t), &cpus);
+#elif HAVE_MACH_THREAD_POLICY_H
   thread_affinity_policy_data_t affinityData;
+  affinityData.affinity_tag = w;
+  return thread_policy_set(mach_thread_self(),
+                    THREAD_AFFINITY_POLICY,
+                    &affinityData,
+                    THREAD_AFFINITY_POLICY_COUNT);
+#else
+  return 0; // fake success
 #endif
+}
+
+CAMLprim value setcore(value which) {
+  int numcores = sysconf( _SC_NPROCESSORS_ONLN );
+  int w = Int_val(which) % numcores; // stay in the space of existing cores
   int retcode;
   int finished=0;
   if (numcores <= 1) // only one core in the system, no need to attempt pinning
     return Val_unit;
-  while (finished==0)
+  do
+  {
+    retcode = setcorew(w);
+    if(retcode != 0)
     {
-#if HAVE_DECL_SCHED_SETAFFINITY
-      CPU_ZERO(&cpus); 
-      CPU_SET (w,&cpus);
-      //fprintf(stderr,"Trying to pin to cpu %d out of %d reported by the system\n",w,numcores);
-      retcode = sched_setaffinity(getpid(), sizeof(cpu_set_t), &cpus);
-      if(retcode != 0) {
-	fprintf(stderr,"Failed pinning to cpu %d, trying %d/2\n",w, w); 
-	w=w/2;
-      }
-      else
-#endif
-#if HAVE_MACH_THREAD_POLICY_H
-      affinityData.affinity_tag = w;
-      retcode = thread_policy_set(mach_thread_self(),
-                        THREAD_AFFINITY_POLICY,
-                        &affinityData,
-                        THREAD_AFFINITY_POLICY_COUNT);
-      if(retcode) {
-        fprintf(stderr,"MAC OS X: Failed pinning to cpu %d, trying %d/2\n",w, w);
-        w=w/2;
-      }
-      else 
-#endif
-	{ //fprintf(stderr,"Succeeded pinning to cpu %d\n",w); 
-	  finished=1;
-	}
+      fprintf(stderr,"Failed pinning to cpu %d, trying %d/2\n",w, w);
+      w=w/2;
+    }
+    else
+    { //fprintf(stderr,"Succeeded pinning to cpu %d\n",w);
+      finished=1;
     }
+  } while (finished==0 && w>0);
   return Val_unit;
 }