From b339a871dec81827cc81d5e084dd0bbd47decac4 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Wed, 21 Sep 2022 22:07:09 +0200
Subject: [PATCH] ARROW-17800: [C++] Fix failures in jemalloc stats tests
 (#14194)

- Provide compatibility for 32-bit platforms
- Avoid memory leak in tests
- Make checks less strict

Authored-by: Antoine Pitrou <antoine@python.org>
Signed-off-by: Antoine Pitrou <antoine@python.org>
---
 cpp/src/arrow/memory_pool_jemalloc.cc | 28 ++++++++++++++++++++-------
 cpp/src/arrow/memory_pool_test.cc     | 23 +++++++++++++---------
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/cpp/src/arrow/memory_pool_jemalloc.cc b/cpp/src/arrow/memory_pool_jemalloc.cc
index 03d2b28ee3e3e..c7d73c84910f6 100644
--- a/cpp/src/arrow/memory_pool_jemalloc.cc
+++ b/cpp/src/arrow/memory_pool_jemalloc.cc
@@ -155,9 +155,8 @@ Status jemalloc_set_decay_ms(int ms) {
 #undef RETURN_IF_JEMALLOC_ERROR
 
 Result<int64_t> jemalloc_get_stat(const char* name) {
-  size_t sz = sizeof(uint64_t);
+  size_t sz;
   int err;
-  uint64_t value;
 
   // Update the statistics cached by mallctl.
   if (std::strcmp(name, "stats.allocated") == 0 ||
@@ -167,16 +166,31 @@ Result<int64_t> jemalloc_get_stat(const char* name) {
       std::strcmp(name, "stats.mapped") == 0 ||
       std::strcmp(name, "stats.retained") == 0) {
     uint64_t epoch;
+    sz = sizeof(epoch);
     mallctl("epoch", &epoch, &sz, &epoch, sz);
   }
 
-  err = mallctl(name, &value, &sz, nullptr, 0);
-
-  if (err) {
-    return arrow::internal::IOErrorFromErrno(err, "Failed retrieving ", &name);
+  // Depending on the stat being queried and on the platform, we could need
+  // to pass a uint32_t or uint64_t pointer. Try both.
+  {
+    uint64_t value = 0;
+    sz = sizeof(value);
+    err = mallctl(name, &value, &sz, nullptr, 0);
+    if (!err) {
+      return value;
+    }
+  }
+  // EINVAL means the given value length (`sz`) was incorrect.
+  if (err == EINVAL) {
+    uint32_t value = 0;
+    sz = sizeof(value);
+    err = mallctl(name, &value, &sz, nullptr, 0);
+    if (!err) {
+      return value;
+    }
   }
 
-  return value;
+  return arrow::internal::IOErrorFromErrno(err, "Failed retrieving ", &name);
 }
 
 Status jemalloc_peak_reset() {
diff --git a/cpp/src/arrow/memory_pool_test.cc b/cpp/src/arrow/memory_pool_test.cc
index 5ac14a44b9a4a..61e0abf3e666b 100644
--- a/cpp/src/arrow/memory_pool_test.cc
+++ b/cpp/src/arrow/memory_pool_test.cc
@@ -180,6 +180,7 @@ TEST(Jemalloc, GetAllocationStats) {
       metadata0, resident0, mapped0, retained0;
   int64_t thread_allocated, thread_deallocated, thread_peak_read, thread_allocated0,
       thread_deallocated0, thread_peak_read0;
+
   MemoryPool* pool = nullptr;
   ABORT_NOT_OK(jemalloc_memory_pool(&pool));
   ASSERT_EQ("jemalloc", pool->backend_name());
@@ -211,14 +212,15 @@ TEST(Jemalloc, GetAllocationStats) {
   ASSERT_OK_AND_ASSIGN(thread_allocated, jemalloc_get_stat("thread.allocated"));
   ASSERT_OK_AND_ASSIGN(thread_deallocated, jemalloc_get_stat("thread.deallocated"));
   ASSERT_OK_AND_ASSIGN(thread_peak_read, jemalloc_get_stat("thread.peak.read"));
+  pool->Free(data, 1023);
 
   // Check allocated stats pre-allocation
-  ASSERT_NEAR(allocated0, 120000, 100000);
-  ASSERT_NEAR(active0, 75000, 70000);
-  ASSERT_NEAR(metadata0, 3000000, 1000000);
-  ASSERT_NEAR(resident0, 3000000, 1000000);
-  ASSERT_NEAR(mapped0, 6500000, 1000000);
-  ASSERT_NEAR(retained0, 1500000, 2000000);
+  ASSERT_GT(allocated0, 0);
+  ASSERT_GT(active0, 0);
+  ASSERT_GT(metadata0, 0);
+  ASSERT_GT(resident0, 0);
+  ASSERT_GT(mapped0, 0);
+  ASSERT_GE(retained0, 0);
 
   // Check allocated stats change due to allocation
   ASSERT_NEAR(allocated - allocated0, 70000, 50000);
@@ -233,13 +235,16 @@ TEST(Jemalloc, GetAllocationStats) {
   ASSERT_EQ(thread_deallocated - thread_deallocated0, 1280);
 
   // Resetting thread peak read metric
-  ASSERT_OK(pool->Allocate(12560, &data));
+  ASSERT_OK(pool->Allocate(100000, &data));
   ASSERT_OK_AND_ASSIGN(thread_peak_read, jemalloc_get_stat("thread.peak.read"));
-  ASSERT_NEAR(thread_peak_read, 15616, 1000);
+  ASSERT_NEAR(thread_peak_read, 100000, 50000);
+  pool->Free(data, 100000);
   ASSERT_OK(jemalloc_peak_reset());
+
   ASSERT_OK(pool->Allocate(1256, &data));
   ASSERT_OK_AND_ASSIGN(thread_peak_read, jemalloc_get_stat("thread.peak.read"));
-  ASSERT_NEAR(thread_peak_read, 1280, 100);
+  ASSERT_NEAR(thread_peak_read, 1256, 100);
+  pool->Free(data, 1256);
 
   // Print statistics to stderr
   ASSERT_OK(jemalloc_stats_print("J"));