[object store refactor 1/n] Introduce IAllocator and PlasmaAllocator

BUILD.bazel

@@ -339,6 +339,7 @@ cc_library(
     ],
     hdrs = [
         "src/ray/object_manager/common.h",
+        "src/ray/object_manager/plasma/allocator.h",
         "src/ray/object_manager/plasma/create_request_queue.h",
         "src/ray/object_manager/plasma/eviction_policy.h",
         "src/ray/object_manager/plasma/plasma_allocator.h",

src/ray/object_manager/object_manager.cc

@@ -930,15 +930,19 @@ std::string ObjectManager::DebugString() const {
 void ObjectManager::RecordMetrics() const {
   stats::ObjectStoreAvailableMemory().Record(config_.object_store_memory - used_memory_);
   stats::ObjectStoreUsedMemory().Record(used_memory_);
-  stats::ObjectStoreFallbackMemory().Record(plasma::PlasmaAllocator::FallbackAllocated());
+  stats::ObjectStoreFallbackMemory().Record(
+      // TODO(scv119): this is not thread safe.
+      plasma::PlasmaAllocator::GetInstance().FallbackAllocated());
   stats::ObjectStoreLocalObjects().Record(local_objects_.size());
   stats::ObjectManagerPullRequests().Record(pull_manager_->NumActiveRequests());
 }
 
 void ObjectManager::FillObjectStoreStats(rpc::GetNodeStatsReply *reply) const {
   auto stats = reply->mutable_store_stats();
   stats->set_object_store_bytes_used(used_memory_);
-  stats->set_object_store_bytes_fallback(plasma::PlasmaAllocator::FallbackAllocated());
+  stats->set_object_store_bytes_fallback(
+      // TODO(scv119): this is not thread safe.
+      plasma::PlasmaAllocator::GetInstance().FallbackAllocated());
   stats->set_object_store_bytes_avail(config_.object_store_memory);
   stats->set_num_local_objects(local_objects_.size());
   stats->set_consumed_bytes(plasma::plasma_store_runner->GetConsumedBytes());

src/ray/object_manager/plasma/allocator.h

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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
+//
+//   http://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.
+#pragma once
+
+#include "absl/types/optional.h"
+#include "ray/object_manager/plasma/common.h"
+#include "ray/object_manager/plasma/compat.h"
+
+namespace plasma {
+
+// IAllocator is responsible for allocating/deallocating memories.
+// This class is not thread safe.
+class IAllocator {
+ public:
+  virtual ~IAllocator() = default;
+
+  /// Allocates size bytes and returns allocated memory.
+  ///
+  /// \param bytes Number of bytes.
+  /// \return allocated memory. returns empty if not enough space.
+  virtual absl::optional<Allocation> Allocate(size_t bytes) = 0;
+
+  // Same as Allocate, but allocates pages from the filesystem. The footprint limit
+  // is not enforced for these allocations, but allocations here are still tracked
+  // and count towards the limit.
+  //
+  // TODO(scv119) ideally we should have mem/disk allocator implementations
+  // so we don't need FallbackAllocate. However the dlmalloc has some limitation
+  // prevents us from doing so.
+  virtual absl::optional<Allocation> FallbackAllocate(size_t bytes) = 0;
+
+  /// Frees the memory space pointed to by mem, which must have been returned by
+  /// a previous call to Allocate/FallbackAllocate or it yield undefined behavior.
+  ///
+  /// \param allocation allocation to free.
+  virtual void Free(const Allocation &allocation) = 0;
+
+  /// Sets the memory footprint limit for this allocator.
+  ///
+  /// \param bytes memory footprint limit in bytes.
+  virtual void SetFootprintLimit(size_t bytes) = 0;
+
+  /// Get the memory footprint limit for this allocator.
+  ///
+  /// \return memory footprint limit in bytes.
+  virtual int64_t GetFootprintLimit() const = 0;
+
+  /// Get the number of bytes allocated so far.
+  /// \return Number of bytes allocated so far.
+  virtual int64_t Allocated() const = 0;
+
+  /// Get the number of bytes fallback allocated so far.
+  /// \return Number of bytes fallback allocated so far.
+  virtual int64_t FallbackAllocated() const = 0;
+};
+
+}  // namespace plasma

src/ray/object_manager/plasma/common.h

@@ -45,23 +45,38 @@ enum class ObjectState : int {
   PLASMA_SEALED = 2,
 };
 
+// Represents a chunk of allocated memory.
+struct Allocation {
+  /// Pointer to the allocated memory.
+  void *address;
+  /// Num bytes of the allocated memory.
+  int64_t size;
+  /// The file descriptor of the memory mapped file where the memory allocated.
+  MEMFD_TYPE fd;
+  /// The offset in bytes in the memory mapped file of the allocated memory.
+  ptrdiff_t offset;
+  /// Device number of the allocated memory.
+  int device_num;
+  /// the total size of this mapped memory.
+  int64_t mmap_size;
+
+  Allocation(void *address, int64_t size, MEMFD_TYPE fd, ptrdiff_t offset, int device_num,
+             int64_t mmap_size)
+      : address(address),
+        size(size),
+        fd(std::move(fd)),
+        offset(offset),
+        device_num(device_num),
+        mmap_size(mmap_size) {}
+};
+
 /// This type is used by the Plasma store. It is here because it is exposed to
 /// the eviction policy.
 struct ObjectTableEntry {
-  ObjectTableEntry();
-
-  ~ObjectTableEntry();
+  ObjectTableEntry(Allocation allocation);
 
-  /// Memory mapped file containing the object.
-  MEMFD_TYPE fd;
-  /// Device number.
-  int device_num;
-  /// Size of the underlying map.
-  int64_t map_size;
-  /// Offset from the base of the mmap.
-  ptrdiff_t offset;
-  /// Pointer to the object data. Needed to free the object.
-  uint8_t *pointer;
+  /// Allocation Info;
+  Allocation allocation;
   /// Size of the object in bytes.
   int64_t data_size;
   /// Size of the object metadata in bytes.

src/ray/object_manager/plasma/eviction_policy.cc

@@ -90,8 +90,11 @@ int64_t LRUCache::ChooseObjectsToEvict(int64_t num_bytes_required,
   return bytes_evicted;
 }
 
-EvictionPolicy::EvictionPolicy(PlasmaStoreInfo *store_info, int64_t max_size)
-    : pinned_memory_bytes_(0), store_info_(store_info), cache_("global lru", max_size) {}
+EvictionPolicy::EvictionPolicy(PlasmaStoreInfo *store_info, const IAllocator &allocator)
+    : pinned_memory_bytes_(0),
+      store_info_(store_info),
+      allocator_(allocator),
+      cache_("global lru", allocator_.GetFootprintLimit()) {}
 
 int64_t EvictionPolicy::ChooseObjectsToEvict(int64_t num_bytes_required,
                                              std::vector<ObjectID> *objects_to_evict) {
@@ -111,18 +114,16 @@ void EvictionPolicy::ObjectCreated(const ObjectID &object_id, bool is_create) {
 int64_t EvictionPolicy::RequireSpace(int64_t size,
                                      std::vector<ObjectID> *objects_to_evict) {
   // Check if there is enough space to create the object.
-  int64_t required_space =
-      PlasmaAllocator::Allocated() + size - PlasmaAllocator::GetFootprintLimit();
+  int64_t required_space = allocator_.Allocated() + size - allocator_.GetFootprintLimit();
   // Try to free up at least as much space as we need right now but ideally
   // up to 20% of the total capacity.
-  int64_t space_to_free =
-      std::max(required_space, PlasmaAllocator::GetFootprintLimit() / 5);
+  int64_t space_to_free = std::max(required_space, allocator_.GetFootprintLimit() / 5);
   // Choose some objects to evict, and update the return pointers.
   int64_t num_bytes_evicted = ChooseObjectsToEvict(space_to_free, objects_to_evict);
   RAY_LOG(DEBUG) << "There is not enough space to create this object, so evicting "
                  << objects_to_evict->size() << " objects to free up "
                  << num_bytes_evicted << " bytes. The number of bytes in use (before "
-                 << "this eviction) is " << PlasmaAllocator::Allocated() << ".";
+                 << "this eviction) is " << allocator_.Allocated() << ".";
   return required_space - num_bytes_evicted;
 }
 

src/ray/object_manager/plasma/eviction_policy.h

@@ -24,6 +24,7 @@
 #include <utility>
 #include <vector>
 
+#include "ray/object_manager/plasma/allocator.h"
 #include "ray/object_manager/plasma/common.h"
 #include "ray/object_manager/plasma/plasma.h"
 
@@ -96,8 +97,8 @@ class EvictionPolicy {
   ///
   /// \param store_info Information about the Plasma store that is exposed
   ///        to the eviction policy.
-  /// \param max_size Max size in bytes total of objects to store.
-  explicit EvictionPolicy(PlasmaStoreInfo *store_info, int64_t max_size);
+  /// \param allocator Memory allocator.
+  explicit EvictionPolicy(PlasmaStoreInfo *store_info, const IAllocator &allocator);
 
   /// Destroy an eviction policy.
   virtual ~EvictionPolicy() {}
@@ -173,6 +174,9 @@ class EvictionPolicy {
 
   /// Pointer to the plasma store info.
   PlasmaStoreInfo *store_info_;
+
+  const IAllocator &allocator_;
+
   /// Datastructure for the LRU cache.
   LRUCache cache_;
 };

src/ray/object_manager/plasma/malloc.cc

@@ -32,32 +32,25 @@ static ptrdiff_t pointer_distance(void const *pfrom, void const *pto) {
   return (unsigned char const *)pto - (unsigned char const *)pfrom;
 }
 
-void GetMallocMapinfo(void *addr, MEMFD_TYPE *fd, int64_t *map_size, ptrdiff_t *offset) {
+bool GetMallocMapinfo(const void *const addr, MEMFD_TYPE *fd, int64_t *map_size,
+                      ptrdiff_t *offset) {
   // TODO(rshin): Implement a more efficient search through mmap_records.
   for (const auto &entry : mmap_records) {
     if (addr >= entry.first && addr < pointer_advance(entry.first, entry.second.size)) {
       fd->first = entry.second.fd.first;
       fd->second = entry.second.fd.second;
       *map_size = entry.second.size;
       *offset = pointer_distance(entry.first, addr);
-      return;
+      return true;
     }
   }
+
   fd->first = INVALID_FD;
   fd->second = INVALID_UNIQUE_FD_ID;
   *map_size = 0;
   *offset = 0;
-}
 
-int64_t GetMmapSize(MEMFD_TYPE fd) {
-  for (const auto &entry : mmap_records) {
-    if (entry.second.fd == fd) {
-      return entry.second.size;
-    }
-  }
-  RAY_LOG(FATAL) << "failed to find entry in mmap_records for fd " << fd.first << " "
-                 << fd.second;
-  return -1;  // This code is never reached.
+  return false;
 }
 
 }  // namespace plasma

src/ray/object_manager/plasma/malloc.h

@@ -31,14 +31,6 @@ namespace plasma {
 /// (in the client we cannot guarantee that these mmaps are contiguous).
 constexpr int64_t kMmapRegionsGap = sizeof(size_t);
 
-void GetMallocMapinfo(void *addr, MEMFD_TYPE *fd, int64_t *map_length, ptrdiff_t *offset);
-
-/// Get the mmap size corresponding to a specific file descriptor.
-///
-/// \param fd The file descriptor to look up.
-/// \return The size of the corresponding memory-mapped file.
-int64_t GetMmapSize(MEMFD_TYPE fd);
-
 struct MmapRecord {
   MEMFD_TYPE fd;
   int64_t size;
@@ -49,4 +41,11 @@ struct MmapRecord {
 /// and size.
 extern std::unordered_map<void *, MmapRecord> mmap_records;
 
+/// private function, only used by PlasmaAllocator to look up Mmap information
+/// given an address allocated by dlmalloc.
+///
+/// \return true if look up succeed. false means the address is not allocated
+/// by dlmalloc.
+bool GetMallocMapinfo(const void *const addr, MEMFD_TYPE *fd, int64_t *map_length,
+                      ptrdiff_t *offset);
 }  // namespace plasma

src/ray/object_manager/plasma/plasma.cc

@@ -21,15 +21,14 @@
 
 namespace plasma {
 
-ObjectTableEntry::ObjectTableEntry() : pointer(nullptr), ref_count(0) {}
-
-ObjectTableEntry::~ObjectTableEntry() { pointer = nullptr; }
+ObjectTableEntry::ObjectTableEntry(Allocation allocation)
+    : allocation(std::move(allocation)), ref_count(0) {}
 
 ObjectTableEntry *GetObjectTableEntry(PlasmaStoreInfo *store_info,
                                       const ObjectID &object_id) {
   auto it = store_info->objects.find(object_id);
   if (it == store_info->objects.end()) {
-    return NULL;
+    return nullptr;
   }
   return it->second.get();
 }