Implement an object lifetime predictor in TCMalloc.

This library stores a database that maps stack traces to lifetimes (short/long) of objects allocated at this particular stack trace, and allows efficient lookup of these lifetimes.

PiperOrigin-RevId: 391660025
Change-Id: Ia6308f58099c4d7cd3d52a36fe8d573194dd3433

Author: martinmaas

tcmalloc/internal/BUILD

@@ -82,6 +82,39 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "lifetime_predictions",
+    hdrs = ["lifetime_predictions.h"],
+    copts = TCMALLOC_DEFAULT_COPTS,
+    visibility = [
+        "//tcmalloc:__subpackages__",
+    ],
+    deps = [
+        ":linked_list",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:malloc_internal",
+        "@com_google_absl//absl/debugging:stacktrace",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/time",
+    ],
+)
+
+cc_test(
+    name = "lifetime_predictions_test",
+    srcs = ["lifetime_predictions_test.cc"],
+    copts = TCMALLOC_DEFAULT_COPTS,
+    deps = [
+        ":lifetime_predictions",
+        ":logging",
+        "@com_github_google_benchmark//:benchmark",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/flags:flag",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "linked_list",
     hdrs = ["linked_list.h"],

tcmalloc/internal/lifetime_predictions.h

@@ -0,0 +1,252 @@
+// Copyright 2020 The TCMalloc Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef TCMALLOC_INTERNAL_LIFETIME_PREDICTIONS_H_
+#define TCMALLOC_INTERNAL_LIFETIME_PREDICTIONS_H_
+
+#include <algorithm>
+#include <cstdlib>
+#include <functional>
+
+#include "absl/algorithm/container.h"
+#include "absl/base/const_init.h"
+#include "absl/base/internal/low_level_alloc.h"
+#include "absl/base/internal/spinlock.h"
+#include "absl/debugging/stacktrace.h"
+#include "absl/hash/hash.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "tcmalloc/internal/linked_list.h"
+
+GOOGLE_MALLOC_SECTION_BEGIN
+namespace tcmalloc {
+namespace tcmalloc_internal {
+
+// Counts how many times we observed objects with a particular stack trace
+// that were short lived/long lived. Each LifetimeStats object is associated
+// with a particular allocation site (i.e., allocation stack trace) and each
+// allocation site has at most one LifetimeStats object. All accesses to
+// LifetimeStats objects need to be synchronized via the page heap lock.
+class LifetimeStats : public TList<LifetimeStats>::Elem {
+ public:
+  enum class Certainty { kLowCertainty, kHighCertainty };
+  enum class Prediction { kShortLived, kLongLived };
+
+  void Update(Prediction prediction) {
+    if (prediction == Prediction::kShortLived) {
+      short_lived_++;
+    } else {
+      long_lived_++;
+    }
+  }
+
+  Prediction Predict(Certainty certainty) {
+    if (certainty == Certainty::kLowCertainty) {
+      return (short_lived_ > long_lived_) ? Prediction::kShortLived
+                                          : Prediction::kLongLived;
+    } else {
+      // If little data was collected, predict as long-lived (current behavior).
+      return (short_lived_ > (long_lived_ + 10)) ? Prediction::kShortLived
+                                                 : Prediction::kLongLived;
+    }
+  }
+
+  // Reference counts are protected by LifetimeDatabase::table_lock_.
+
+  // Increments the reference count of this entry.
+  void IncRef() { ++refcount_; }
+
+  // Returns true if and only if the reference count reaches 0.
+  bool DecRef() { return --refcount_ == 0; }
+
+ private:
+  uint64_t refcount_ = 1;
+  uint64_t short_lived_ = 0;
+  uint64_t long_lived_ = 0;
+};
+
+// Manages stack traces and statistics about their associated lifetimes. Since
+// the database can fill up, old entries are evicted. Evicted entries need to
+// survive as long as the last lifetime tracker referencing them and are thus
+// reference-counted.
+class LifetimeDatabase {
+ public:
+  struct Key {
+    int depth;  // Number of PC values stored in array below
+    void* stack[kMaxStackDepth];
+
+    // Statically instantiate at the start of the allocation to acquire
+    // the allocation stack trace.
+    Key() { depth = absl::GetStackTrace(stack, kMaxStackDepth, 1); }
+
+    template <typename H>
+    friend H AbslHashValue(H h, const Key& c) {
+      return H::combine(H::combine_contiguous(std::move(h), c.stack, c.depth),
+                        c.depth);
+    }
+
+    bool operator==(const Key& other) const {
+      if (depth != other.depth) {
+        return false;
+      }
+      return std::equal(stack, stack + depth, other.stack);
+    }
+  };
+
+  // Captures statistics associated with the low-level allocator backing the
+  // memory used by the database.
+  struct ArenaStats {
+    uint64_t bytes_allocated;
+  };
+
+  static constexpr int kMaxDatabaseSize = 1024;
+
+  LifetimeDatabase() {}
+  ~LifetimeDatabase() {}
+
+  // Not copyable or movable
+  LifetimeDatabase(const LifetimeDatabase&) = delete;
+  LifetimeDatabase& operator=(const LifetimeDatabase&) = delete;
+
+  // Identifies the current stack trace and returns a handle to the lifetime
+  // statistics associated with this stack trace. May run outside the page heap
+  // lock -- we therefore need to do our own locking. This increments the
+  // reference count of the lifetime stats object and the caller is responsible
+  // for calling RemoveLifetimeStatsReference when finished with the object.
+  LifetimeStats* LookupOrAddLifetimeStats(Key* k) {
+    absl::base_internal::SpinLockHolder h(&table_lock_);
+    auto it = table_.find(*k);
+    LifetimeStats* s;
+    if (it == table_.end()) {
+      MaybeEvictLRU();
+      // Allocate a new entry using the low-level allocator, which is safe
+      // to call from within TCMalloc.
+      s = stats_allocator_.allocate(1);
+      new (s) LifetimeStats();
+      table_.insert(std::make_pair(*k, s));
+      stats_fifo_.append(s);
+    } else {
+      s = it->second;
+      UpdateLRU(s);
+    }
+    s->IncRef();
+    return s;
+  }
+
+  void RemoveLifetimeStatsReference(LifetimeStats* s) {
+    absl::base_internal::SpinLockHolder h(&table_lock_);
+    if (s->DecRef()) {
+      stats_allocator_.deallocate(s, 1);
+    }
+  }
+
+  size_t size() const {
+    absl::base_internal::SpinLockHolder h(&table_lock_);
+    return table_.size();
+  }
+
+  size_t evictions() const {
+    absl::base_internal::SpinLockHolder h(&table_lock_);
+    return n_evictions_;
+  }
+
+  static ArenaStats* arena_stats() {
+    static ArenaStats stats = {0};
+    return &stats;
+  }
+
+ protected:
+  static const int kMaxStackDepth = 64;
+
+  static absl::base_internal::LowLevelAlloc::Arena* GetArena() {
+    static absl::base_internal::LowLevelAlloc::Arena* arena =
+        absl::base_internal::LowLevelAlloc::NewArena(0);
+    return arena;
+  }
+
+  static uint64_t bytes_allocated_ ABSL_GUARDED_BY(table_lock_);
+
+  void UpdateLRU(LifetimeStats* stats)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(table_lock_) {
+    stats_fifo_.remove(stats);
+    stats_fifo_.append(stats);
+  }
+
+  // If an entry is evicted, it is returned (nullptr otherwise).
+  void MaybeEvictLRU() ABSL_EXCLUSIVE_LOCKS_REQUIRED(table_lock_) {
+    if (table_.size() < kMaxDatabaseSize) {
+      return;
+    }
+    n_evictions_++;
+    LifetimeStats* evict = stats_fifo_.first();
+    stats_fifo_.remove(evict);
+    for (auto it = table_.begin(); it != table_.end(); ++it) {
+      if (it->second == evict) {
+        table_.erase(it);
+        if (evict->DecRef()) {
+          stats_allocator_.deallocate(evict, 1);
+        }
+        return;
+      }
+    }
+    CHECK_CONDITION(false);  // Should not happen
+  }
+
+ private:
+  template <typename T>
+  class MyAllocator : public std::allocator<T> {
+   public:
+    template <typename U>
+    struct rebind {
+      using other = MyAllocator<U>;
+    };
+
+    MyAllocator() noexcept {}
+
+    template <typename U>
+    explicit MyAllocator(const MyAllocator<U>&) noexcept {}
+
+    T* allocate(size_t num_objects, const void* = nullptr) {
+      size_t bytes = num_objects * sizeof(T);
+      arena_stats()->bytes_allocated += bytes;
+      return static_cast<T*>(absl::base_internal::LowLevelAlloc::AllocWithArena(
+          bytes, GetArena()));
+    }
+
+    void deallocate(T* p, size_t num_objects) {
+      size_t bytes = num_objects * sizeof(T);
+      arena_stats()->bytes_allocated -= bytes;
+      absl::base_internal::LowLevelAlloc::Free(p);
+    }
+  };
+
+  MyAllocator<LifetimeStats> stats_allocator_ ABSL_GUARDED_BY(table_lock_);
+  mutable absl::base_internal::SpinLock table_lock_{
+      absl::kConstInit, absl::base_internal::SCHEDULE_KERNEL_ONLY};
+
+  // Stores the current mapping from allocation site to LifetimeStats.
+  std::unordered_map<Key, LifetimeStats*, absl::Hash<Key>, std::equal_to<Key>,
+                     MyAllocator<std::pair<const Key, LifetimeStats*>>>
+      table_ ABSL_GUARDED_BY(table_lock_);
+
+  // Stores the entries ordered by how many times they have been accessed.
+  TList<LifetimeStats> stats_fifo_ ABSL_GUARDED_BY(table_lock_);
+  size_t n_evictions_ ABSL_GUARDED_BY(table_lock_) = 0;
+};
+
+}  // namespace tcmalloc_internal
+}  // namespace tcmalloc
+GOOGLE_MALLOC_SECTION_END
+
+#endif  // TCMALLOC_INTERNAL_LIFETIME_PREDICTIONS_H_