add fallback for thread count if jemalloc cannot identify

lnkuiper · Nov 4, 2024 · 961af3a · 961af3a
1 parent 1986445
commit 961af3a
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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>();