[GitHub] [arrow] bkietz commented on a change in pull request #9105: ARROW-11009: [C++] Allow changing default memory pool with an environment variable

[GitBox <gi...@apache.org>]

bkietz commented on a change in pull request #9105:
URL: https://github.com/apache/arrow/pull/9105#discussion_r553562039



##########
File path: cpp/src/arrow/memory_pool.cc
##########
@@ -83,10 +87,94 @@ const char* je_arrow_malloc_conf =
 #endif  // ARROW_JEMALLOC
 
 namespace arrow {
+namespace {
 
 constexpr size_t kAlignment = 64;
 
-namespace {
+constexpr char kDefaultBackendEnvVar[] = "ARROW_DEFAULT_MEMORY_POOL";
+
+enum class MemoryPoolBackend : uint8_t { System, Jemalloc, Mimalloc };
+
+std::string BackendName(MemoryPoolBackend backend) {
+  switch (backend) {
+    case MemoryPoolBackend::System:
+      return "system";
+    case MemoryPoolBackend::Jemalloc:
+      return "jemalloc";
+    case MemoryPoolBackend::Mimalloc:
+      return "mimalloc";
+    default:
+      return "";  // XXX unreachable
+  }
+}
+
+std::vector<MemoryPoolBackend> SupportedBackends() {
+  std::vector<MemoryPoolBackend> backends = {
+#ifdef ARROW_JEMALLOC
+      MemoryPoolBackend::Jemalloc,
+#endif
+#ifdef ARROW_MIMALLOC
+      MemoryPoolBackend::Mimalloc,
+#endif
+      MemoryPoolBackend::System};
+  return backends;
+}
+
+const std::vector<MemoryPoolBackend> supported_backends = SupportedBackends();
+
+util::optional<MemoryPoolBackend> UserSelectedBackend() {
+  auto unsupported_backend = [](const std::string& name) {
+    std::vector<std::string> supported;
+    for (const auto backend : supported_backends) {
+      supported.push_back("'" + BackendName(backend) + "'");
+    }
+    ARROW_LOG(WARNING) << "Unsupported backend '" << name << "' specified in "
+                       << kDefaultBackendEnvVar << " (supported backends are "
+                       << internal::JoinStrings(supported, ", ") << ")";
+  };
+
+  auto maybe_name = internal::GetEnvVar(kDefaultBackendEnvVar);
+  if (!maybe_name.ok()) {
+    return {};
+  }
+  const auto name = *std::move(maybe_name);
+  if (name.empty()) {
+    return {};
+  } else if (name == "system") {
+    return MemoryPoolBackend::System;
+  } else if (name == "jemalloc") {
+#ifdef ARROW_JEMALLOC
+    return MemoryPoolBackend::Jemalloc;
+#else
+    unsupported_backend(name);
+#endif
+  } else if (name == "mimalloc") {
+#ifdef ARROW_MIMALLOC
+    return MemoryPoolBackend::Mimalloc;
+#else
+    unsupported_backend(name);
+#endif
+  } else {
+    unsupported_backend(name);
+  }

Review comment:
       Nit: this would be more readable if you iterated through SupportedBackends and created a `map<name, backend>` then did lookup in that

##########
File path: cpp/src/arrow/memory_pool.cc
##########
@@ -404,14 +499,25 @@ Status mimalloc_memory_pool(MemoryPool** out) {
 #endif
 }
 
+static std::unique_ptr<MemoryPool> default_pool = MemoryPool::CreateDefault();

Review comment:
       This seems to be unused. Did you intend to write
   ```c++
   MemoryPool* default_memory_pool() { return default_pool.get(); }
   ```
   ?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org