From 961af3a858cfbf3e71258c0c185770bb130a8b20 Mon Sep 17 00:00:00 2001
From: Laurens Kuiper <laurens.kuiper@cwi.nl>
Date: Mon, 4 Nov 2024 08:54:36 +0100
Subject: [PATCH] add fallback for thread count if jemalloc cannot identify

---
 extension/jemalloc/jemalloc/README.md      | 5 +++++
 extension/jemalloc/jemalloc/src/jemalloc.c | 5 +++++
 extension/jemalloc/jemalloc_extension.cpp  | 7 +++++++
 3 files changed, 17 insertions(+)

diff --git a/extension/jemalloc/jemalloc/README.md b/extension/jemalloc/jemalloc/README.md
index 9bf08821e76..81b6669327e 100644
--- a/extension/jemalloc/jemalloc/README.md
+++ b/extension/jemalloc/jemalloc/README.md
@@ -65,6 +65,8 @@ Add this to `jemalloc.h`:
 We also supply our own config string in `jemalloc.c`.
 Define this just after the `#include`s.
 ```c++
+#include "malloc_ncpus.h"
+
 #define JE_MALLOC_CONF_BUFFER_SIZE 200
 char JE_MALLOC_CONF_BUFFER[JE_MALLOC_CONF_BUFFER_SIZE];
 ```
@@ -74,6 +76,9 @@ JEMALLOC_ATTR(constructor)
 static void
 jemalloc_constructor(void) {
 	unsigned long long cpu_count = malloc_ncpus();
+	if (cpu_count == 0) {
+		cpu_count = duckdb_malloc_ncpus();
+	}
 	unsigned long long bgt_count = cpu_count / 16;
 	if (bgt_count == 0) {
 		bgt_count = 1;
diff --git a/extension/jemalloc/jemalloc/src/jemalloc.c b/extension/jemalloc/jemalloc/src/jemalloc.c
index 9d57c2e42e7..292d553aa36 100644
--- a/extension/jemalloc/jemalloc/src/jemalloc.c
+++ b/extension/jemalloc/jemalloc/src/jemalloc.c
@@ -25,6 +25,8 @@
 #include "jemalloc/internal/thread_event.h"
 #include "jemalloc/internal/util.h"
 
+#include "malloc_ncpus.h"
+
 /******************************************************************************/
 /* Data. */
 
@@ -4270,6 +4272,9 @@ JEMALLOC_ATTR(constructor)
 static void
 jemalloc_constructor(void) {
 	unsigned long long cpu_count = malloc_ncpus();
+	if (cpu_count == 0) {
+		cpu_count = duckdb_malloc_ncpus();
+	}
 	unsigned long long bgt_count = cpu_count / 16;
 	if (bgt_count == 0) {
 		bgt_count = 1;
diff --git a/extension/jemalloc/jemalloc_extension.cpp b/extension/jemalloc/jemalloc_extension.cpp
index 10a49cf99a1..f6d4db65e57 100644
--- a/extension/jemalloc/jemalloc_extension.cpp
+++ b/extension/jemalloc/jemalloc_extension.cpp
@@ -3,6 +3,9 @@
 
 #include "duckdb/common/allocator.hpp"
 #include "jemalloc/jemalloc.h"
+#include "malloc_ncpus.h"
+
+#include <thread>
 
 namespace duckdb {
 
@@ -114,6 +117,10 @@ std::string JemallocExtension::Version() const {
 
 extern "C" {
 
+unsigned duckdb_malloc_ncpus() {
+	return duckdb::NumericCast<unsigned>(std::thread::hardware_concurrency());
+}
+
 DUCKDB_EXTENSION_API void jemalloc_init(duckdb::DatabaseInstance &db) {
 	duckdb::DuckDB db_wrapper(db);
 	db_wrapper.LoadExtension<duckdb::JemallocExtension>();