From fc3a6eb74ae5fa829a7dd89620567f3c1823053e Mon Sep 17 00:00:00 2001
From: Peter Dillinger <peterd@fb.com>
Date: Mon, 29 Nov 2021 10:52:32 -0800
Subject: [PATCH] Fix/improve 'must free heap allocations' code (#9209)

Summary:
Added missing include, and cleaned up to make same mistake less
likely in future (minimize conditional compilation)

Fixes https://github.com/facebook/rocksdb/issues/9183

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9209

Test Plan: added to existing test

Reviewed By: mrambacher

Differential Revision: D32631390

Pulled By: pdillinger

fbshipit-source-id: 63a0501855cf5fac9e22ca1e5c4f53725dbf3f93
---
 cache/cache_test.cc  |  3 +++
 cache/clock_cache.cc | 11 ++++++-----
 cache/lru_cache.cc   | 11 ++++++-----
 port/lang.h          |  2 ++
 4 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/cache/cache_test.cc b/cache/cache_test.cc
index 66881b5c9..8562490ee 100644
--- a/cache/cache_test.cc
+++ b/cache/cache_test.cc
@@ -609,6 +609,9 @@ TEST_P(CacheTest, SetCapacity) {
   for (size_t i = 5; i < 10; i++) {
     cache->Release(handles[i]);
   }
+
+  // Make sure this doesn't crash or upset ASAN/valgrind
+  cache->DisownData();
 }
 
 TEST_P(LRUCacheTest, SetStrictCapacityLimit) {
diff --git a/cache/clock_cache.cc b/cache/clock_cache.cc
index d2bcab0e2..2e501805e 100644
--- a/cache/clock_cache.cc
+++ b/cache/clock_cache.cc
@@ -33,11 +33,11 @@ std::shared_ptr<Cache> NewClockCache(
 #ifndef ROCKSDB_USE_RTTI
 #define TBB_USE_EXCEPTIONS 0
 #endif
-#include "tbb/concurrent_hash_map.h"
-
 #include "cache/sharded_cache.h"
+#include "port/lang.h"
 #include "port/malloc.h"
 #include "port/port.h"
+#include "tbb/concurrent_hash_map.h"
 #include "util/autovector.h"
 #include "util/mutexlock.h"
 
@@ -809,9 +809,10 @@ class ClockCache final : public ShardedCache {
   }
 
   void DisownData() override {
-#ifndef MUST_FREE_HEAP_ALLOCATIONS
-    shards_ = nullptr;
-#endif
+    // Leak data only if that won't generate an ASAN/valgrind warning
+    if (!kMustFreeHeapAllocations) {
+      shards_ = nullptr;
+    }
   }
 
   void WaitAll(std::vector<Handle*>& /*handles*/) override {}
diff --git a/cache/lru_cache.cc b/cache/lru_cache.cc
index bab58a2a2..bb9186f1e 100644
--- a/cache/lru_cache.cc
+++ b/cache/lru_cache.cc
@@ -15,6 +15,7 @@
 
 #include "monitoring/perf_context_imp.h"
 #include "monitoring/statistics.h"
+#include "port/lang.h"
 #include "util/mutexlock.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -704,11 +705,11 @@ uint32_t LRUCache::GetHash(Handle* handle) const {
 }
 
 void LRUCache::DisownData() {
-// Do not drop data if compile with ASAN to suppress leak warning.
-#ifndef MUST_FREE_HEAP_ALLOCATIONS
-  shards_ = nullptr;
-  num_shards_ = 0;
-#endif
+  // Leak data only if that won't generate an ASAN/valgrind warning
+  if (!kMustFreeHeapAllocations) {
+    shards_ = nullptr;
+    num_shards_ = 0;
+  }
 }
 
 size_t LRUCache::TEST_GetLRUSize() {
diff --git a/port/lang.h b/port/lang.h
index 4a601d41c..5d1067e1f 100644
--- a/port/lang.h
+++ b/port/lang.h
@@ -42,8 +42,10 @@
 //   STATIC_AVOID_DESTRUCTION(Foo, foo)(arg1, arg2);
 #ifdef MUST_FREE_HEAP_ALLOCATIONS
 #define STATIC_AVOID_DESTRUCTION(Type, name) static Type name
+constexpr bool kMustFreeHeapAllocations = true;
 #else
 #define STATIC_AVOID_DESTRUCTION(Type, name) static Type& name = *new Type
+constexpr bool kMustFreeHeapAllocations = false;
 #endif
 
 // TSAN (Thread sanitizer)