From 3b8e8078408feb49cd41272935f3df2e35faa092 Mon Sep 17 00:00:00 2001
From: Joseph Hunkeler <jhunkeler@gmail.com>
Date: Mon, 24 Jun 2024 10:51:26 -0400
Subject: Add get_cpu_count()

* Exposes STASIS_CPU_COUNT and CPU_COUNT to the runtime environment
---
 include/utils.h |  6 ++++++
 src/delivery.c  | 10 ++++++++++
 src/utils.c     |  7 +++++++
 3 files changed, 23 insertions(+)

diff --git a/include/utils.h b/include/utils.h
index 2c80e77..eee2e30 100644
--- a/include/utils.h
+++ b/include/utils.h
@@ -351,4 +351,10 @@ int redact_sensitive(const char **to_redact, size_t to_redact_size, char *src, c
  */
 struct StrList *listdir(const char *path);
 
+/**
+ * Get CPU count
+ * @return CPU count on success, zero on error
+ */
+long get_cpu_count();
+
 #endif //STASIS_UTILS_H
diff --git a/src/delivery.c b/src/delivery.c
index fdc2516..5c507bf 100644
--- a/src/delivery.c
+++ b/src/delivery.c
@@ -353,7 +353,17 @@ int delivery_init_platform(struct Delivery *ctx) {
         tolower_s(ctx->system.platform[DELIVERY_PLATFORM_RELEASE]);
     }
 
+    long cpu_count = get_cpu_count();
+    if (!cpu_count) {
+        fprintf(stderr, "Unable to determine CPU count. Falling back to 1.\n");
+        cpu_count = 1;
+    }
+    char ncpus[100] = {0};
+    sprintf(ncpus, "%ld", cpu_count);
+
     // Declare some important bits as environment variables
+    setenv("CPU_COUNT", ncpus, 1);
+    setenv("STASIS_CPU_COUNT", ncpus, 1);
     setenv("STASIS_ARCH", ctx->system.arch, 1);
     setenv("STASIS_PLATFORM", ctx->system.platform[DELIVERY_PLATFORM], 1);
     setenv("STASIS_CONDA_ARCH", ctx->system.arch, 1);
diff --git a/src/utils.c b/src/utils.c
index fb22f59..2143c52 100644
--- a/src/utils.c
+++ b/src/utils.c
@@ -758,3 +758,10 @@ struct StrList *listdir(const char *path) {
     return node;
 }
 
+long get_cpu_count() {
+#if defined(STASIS_OS_LINUX) || defined(STASIS_OS_DARWIN)
+        return sysconf(_SC_NPROCESSORS_ONLN);
+#else
+        return 0;
+#endif
+}
\ No newline at end of file
-- 
cgit