Extend CompressedPtr to work with multiple tiers

Now it's size is 8 bytes intead of 4.

Original CompressedPtr stored only some offset with a memory Allocator.
For multi-tier implementation, this is not enough. We must also store
tierId and when uncompressing, select a proper allocator.

An alternative could be to just resign from CompressedPtr but they
are leveraged to allow the cache to be mapped to different addresses on shared memory.

Changing CompressedPtr impacted CacheItem size - it increased from 32 to 44 bytes.

Author: igchor

cachelib/allocator/CacheAllocator.h

@@ -1179,7 +1179,8 @@ class CacheAllocator : public CacheBase {
                  sizeof(typename RefcountWithFlags::Value) + sizeof(uint32_t) +
                  sizeof(uint32_t) + sizeof(KAllocation)) == sizeof(Item),
                 "vtable overhead");
-  static_assert(32 == sizeof(Item), "item overhead is 32 bytes");
+  // XXX: this will fail due to CompressedPtr change
+  // static_assert(32 == sizeof(Item), "item overhead is 32 bytes");
 
   // make sure there is no overhead in ChainedItem on top of a regular Item
   static_assert(sizeof(Item) == sizeof(ChainedItem),
@@ -1798,7 +1799,7 @@ class CacheAllocator : public CacheBase {
   }
 
   typename Item::PtrCompressor createPtrCompressor() const {
-    return allocator_[0 /* TODO */]->createPtrCompressor<Item>();
+    return typename Item::PtrCompressor(allocator_);
   }
 
   // helper utility to throttle and optionally log.

cachelib/allocator/CacheItem.h

@@ -141,6 +141,7 @@ class CACHELIB_PACKED_ATTR CacheItem {
    * to be mapped to different addresses on shared memory.
    */
   using CompressedPtr = facebook::cachelib::CompressedPtr;
+  using SingleTierPtrCompressor = MemoryAllocator::SingleTierPtrCompressor<Item>;
   using PtrCompressor = MemoryAllocator::PtrCompressor<Item>;
 
   // Get the required size for a cache item given the size of memory

cachelib/allocator/memory/AllocationClass.cpp

@@ -50,7 +50,7 @@ AllocationClass::AllocationClass(ClassId classId,
       poolId_(poolId),
       allocationSize_(allocSize),
       slabAlloc_(s),
-      freedAllocations_{slabAlloc_.createPtrCompressor<FreeAlloc>()} {
+      freedAllocations_{slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()} {
   checkState();
 }
 
@@ -102,7 +102,7 @@ AllocationClass::AllocationClass(
       currSlab_(s.getSlabForIdx(*object.currSlabIdx_ref())),
       slabAlloc_(s),
       freedAllocations_(*object.freedAllocationsObject_ref(),
-                        slabAlloc_.createPtrCompressor<FreeAlloc>()),
+                        slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()),
       canAllocate_(*object.canAllocate_ref()) {
   if (!slabAlloc_.isRestorable()) {
     throw std::logic_error("The allocation class cannot be restored.");
@@ -356,9 +356,9 @@ std::pair<bool, std::vector<void*>> AllocationClass::pruneFreeAllocs(
   // allocated slab, release any freed allocations belonging to this slab.
   // Set the bit to true if the corresponding allocation is freed, false
   // otherwise.
-  FreeList freeAllocs{slabAlloc_.createPtrCompressor<FreeAlloc>()};
-  FreeList notInSlab{slabAlloc_.createPtrCompressor<FreeAlloc>()};
-  FreeList inSlab{slabAlloc_.createPtrCompressor<FreeAlloc>()};
+  FreeList freeAllocs{slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()};
+  FreeList notInSlab{slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()};
+  FreeList inSlab{slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()};
 
   lock_->lock_combine([&]() {
     // Take the allocation class free list offline

cachelib/allocator/memory/AllocationClass.h

@@ -446,7 +446,7 @@ class AllocationClass {
   struct CACHELIB_PACKED_ATTR FreeAlloc {
     using CompressedPtr = facebook::cachelib::CompressedPtr;
     using PtrCompressor =
-        facebook::cachelib::PtrCompressor<FreeAlloc, SlabAllocator>;
+        facebook::cachelib::SingleTierPtrCompressor<FreeAlloc, SlabAllocator>;
     SListHook<FreeAlloc> hook_{};
   };
 

cachelib/allocator/memory/CompressedPtr.h

@@ -27,6 +27,9 @@ namespace cachelib {
 
 class SlabAllocator;
 
+template <typename PtrType, typename AllocatorContainer>
+class PtrCompressor;
+
 // the following are for pointer compression for the memory allocator.  We
 // compress pointers by storing the slab index and the alloc index of the
 // allocation inside the slab. With slab worth kNumSlabBits of data, if we
@@ -41,7 +44,7 @@ class SlabAllocator;
 // decompress a CompressedPtr than compress a pointer while creating one.
 class CACHELIB_PACKED_ATTR CompressedPtr {
  public:
-  using PtrType = uint32_t;
+  using PtrType = uint64_t;
   // Thrift doesn't support unsigned type
   using SerializedPtrType = int64_t;
 
@@ -83,14 +86,14 @@ class CACHELIB_PACKED_ATTR CompressedPtr {
  private:
   // null pointer representation. This is almost never guaranteed to be a
   // valid pointer that we can compress to.
-  static constexpr PtrType kNull = 0xffffffff;
+  static constexpr PtrType kNull = 0x00000000ffffffff;
 
   // default construct to null.
   PtrType ptr_{kNull};
 
   // create a compressed pointer for a valid memory allocation.
-  CompressedPtr(uint32_t slabIdx, uint32_t allocIdx)
-      : ptr_(compress(slabIdx, allocIdx)) {}
+  CompressedPtr(uint32_t slabIdx, uint32_t allocIdx, TierId tid = 0)
+      : ptr_(compress(slabIdx, allocIdx, tid)) {}
 
   constexpr explicit CompressedPtr(PtrType ptr) noexcept : ptr_{ptr} {}
 
@@ -100,40 +103,60 @@ class CACHELIB_PACKED_ATTR CompressedPtr {
   static constexpr unsigned int kNumAllocIdxBits =
       Slab::kNumSlabBits - Slab::kMinAllocPower;
 
+  // Use topmost 32 bits for TierId
+  // XXX: optimize
+  static constexpr unsigned int kNumTierIdxOffset = 32;
+
   static constexpr PtrType kAllocIdxMask = ((PtrType)1 << kNumAllocIdxBits) - 1;
 
+  // kNumTierIdxBits most significant bits
+  static constexpr PtrType kTierIdxMask = (((PtrType)1 << kNumTierIdxOffset) - 1) << (NumBits<PtrType>::value - kNumTierIdxOffset);
+
   // Number of bits for the slab index. This will be the top 16 bits of the
   // compressed ptr.
   static constexpr unsigned int kNumSlabIdxBits =
-      NumBits<PtrType>::value - kNumAllocIdxBits;
+      NumBits<PtrType>::value - kNumTierIdxOffset - kNumAllocIdxBits; 
 
-  // Compress the given slabIdx and allocIdx into a 32-bit compressed
+  // Compress the given slabIdx and allocIdx into a 64-bit compressed
   // pointer.
-  static PtrType compress(uint32_t slabIdx, uint32_t allocIdx) noexcept {
+  static PtrType compress(uint32_t slabIdx, uint32_t allocIdx, TierId tid) noexcept {
     XDCHECK_LE(allocIdx, kAllocIdxMask);
     XDCHECK_LT(slabIdx, (1u << kNumSlabIdxBits) - 1);
-    return (slabIdx << kNumAllocIdxBits) + allocIdx;
+    return (static_cast<uint64_t>(tid) << kNumTierIdxOffset) + (slabIdx << kNumAllocIdxBits) + allocIdx;
   }
 
   // Get the slab index of the compressed ptr
   uint32_t getSlabIdx() const noexcept {
     XDCHECK(!isNull());
-    return static_cast<uint32_t>(ptr_ >> kNumAllocIdxBits);
+    auto noTierIdPtr = ptr_ & ~kTierIdxMask;
+    return static_cast<uint32_t>(noTierIdPtr >> kNumAllocIdxBits);
   }
 
   // Get the allocation index of the compressed ptr
   uint32_t getAllocIdx() const noexcept {
     XDCHECK(!isNull());
-    return static_cast<uint32_t>(ptr_ & kAllocIdxMask);
+    auto noTierIdPtr = ptr_ & ~kTierIdxMask;
+    return static_cast<uint32_t>(noTierIdPtr & kAllocIdxMask);
+  }
+
+  uint32_t getTierId() const noexcept {
+    XDCHECK(!isNull());
+    return static_cast<uint32_t>(ptr_ >> kNumTierIdxOffset);
+  }
+
+  void setTierId(TierId tid) noexcept {
+    ptr_ += static_cast<uint64_t>(tid) << kNumTierIdxOffset;
   }
 
   friend SlabAllocator;
+  template <typename CPtrType, typename AllocatorContainer>
+  friend class PtrCompressor;
 };
 
 template <typename PtrType, typename AllocatorT>
-class PtrCompressor {
+class SingleTierPtrCompressor {
  public:
-  explicit PtrCompressor(const AllocatorT& allocator) noexcept
+  explicit SingleTierPtrCompressor(const AllocatorT& allocator) noexcept
       : allocator_(allocator) {}
 
   const CompressedPtr compress(const PtrType* uncompressed) const {
@@ -144,17 +167,61 @@ class PtrCompressor {
     return static_cast<PtrType*>(allocator_.unCompress(compressed));
   }
 
-  bool operator==(const PtrCompressor& rhs) const noexcept {
+  bool operator==(const SingleTierPtrCompressor& rhs) const noexcept {
     return &allocator_ == &rhs.allocator_;
   }
 
-  bool operator!=(const PtrCompressor& rhs) const noexcept {
+  bool operator!=(const SingleTierPtrCompressor& rhs) const noexcept {
     return !(*this == rhs);
   }
 
  private:
   // memory allocator that does the pointer compression.
   const AllocatorT& allocator_;
 };
+
+template <typename PtrType, typename AllocatorContainer>
+class PtrCompressor {
+ public:
+  explicit PtrCompressor(const AllocatorContainer& allocators) noexcept
+      : allocators_(allocators) {}
+
+  const CompressedPtr compress(const PtrType* uncompressed) const {
+    if (uncompressed == nullptr)
+      return CompressedPtr{};
+
+    TierId tid;
+    for (tid = 0; tid < allocators_.size(); tid++) {
+      if (allocators_[tid]->isMemoryInAllocator(static_cast<const void*>(uncompressed)))
+        break;
+    }
+
+    auto cptr = allocators_[tid]->compress(uncompressed);
+    cptr.setTierId(tid);
+
+    return cptr;
+  }
+
+  PtrType* unCompress(const CompressedPtr compressed) const {
+    if (compressed.isNull()) {
+      return nullptr;
+    }
+
+    auto &allocator = *allocators_[compressed.getTierId()];
+    return static_cast<PtrType*>(allocator.unCompress(compressed));
+  }
+
+  bool operator==(const PtrCompressor& rhs) const noexcept {
+    return &allocators_ == &rhs.allocators_;
+  }
+
+  bool operator!=(const PtrCompressor& rhs) const noexcept {
+    return !(*this == rhs);
+  }
+
+ private:
+  // memory allocator that does the pointer compression.
+  const AllocatorContainer& allocators_;
+};
 } // namespace cachelib
 } // namespace facebook

cachelib/allocator/memory/MemoryAllocator.h

@@ -516,12 +516,13 @@ class MemoryAllocator {
   using CompressedPtr = facebook::cachelib::CompressedPtr;
   template <typename PtrType>
   using PtrCompressor =
-      facebook::cachelib::PtrCompressor<PtrType, SlabAllocator>;
+      facebook::cachelib::PtrCompressor<PtrType,
+      std::vector<std::unique_ptr<MemoryAllocator>>>;
 
   template <typename PtrType>
-  PtrCompressor<PtrType> createPtrCompressor() {
-    return slabAllocator_.createPtrCompressor<PtrType>();
-  }
+  using SingleTierPtrCompressor =
+      facebook::cachelib::PtrCompressor<PtrType,
+      SlabAllocator>;
 
   // compress a given pointer to a valid allocation made out of this allocator
   // through an allocate() or nullptr. Calling this otherwise with invalid

cachelib/allocator/memory/SlabAllocator.cpp

@@ -527,6 +527,8 @@ serialization::SlabAllocatorObject SlabAllocator::saveState() {
 // for benchmarking purposes.
 const unsigned int kMarkerBits = 6;
 CompressedPtr SlabAllocator::compressAlt(const void* ptr) const {
+  // XXX: do we need to set tierId here?
+
   if (ptr == nullptr) {
     return CompressedPtr{};
   }
@@ -538,6 +540,8 @@ CompressedPtr SlabAllocator::compressAlt(const void* ptr) const {
 }
 
 void* SlabAllocator::unCompressAlt(const CompressedPtr cPtr) const {
+  // XXX: do we need to set tierId here?
+
   if (cPtr.isNull()) {
     return nullptr;
   }

cachelib/allocator/memory/SlabAllocator.h

@@ -308,8 +308,8 @@ class SlabAllocator {
   }
 
   template <typename PtrType>
-  PtrCompressor<PtrType, SlabAllocator> createPtrCompressor() const {
-    return PtrCompressor<PtrType, SlabAllocator>(*this);
+  SingleTierPtrCompressor<PtrType, SlabAllocator> createSingleTierPtrCompressor() const {
+    return SingleTierPtrCompressor<PtrType, SlabAllocator>(*this);
   }
 
   // returns starting address of memory we own.

cachelib/allocator/tests/AllocatorResizeTest.h

@@ -1098,7 +1098,7 @@ class AllocatorResizeTest : public AllocatorTest<AllocatorT> {
         size_t allocBytes = 0;
         for (size_t k = 0; k < expectedIters * Slab::kSize / sz; k++) {
           const auto key = this->getRandomNewKey(alloc, keyLen);
-          auto handle = util::allocateAccessible(alloc, poolId, key, sz - 45);
+          auto handle = util::allocateAccessible(alloc, poolId, key, sz - 45 - 9 /* TODO: compressed ptr size */);
           if (!handle.get()) {
             break;
           }
@@ -1110,7 +1110,7 @@ class AllocatorResizeTest : public AllocatorTest<AllocatorT> {
         for (size_t k = 0; k < expectedIters * Slab::kSize / sz; k++) {
           const auto key = this->getRandomNewKey(alloc, keyLen);
           size_t allocBytes = 0;
-          auto handle = util::allocateAccessible(alloc, poolId, key, sz - 45);
+          auto handle = util::allocateAccessible(alloc, poolId, key, sz - 45 - 9 /* TODO: compressed ptr size */);
           allocBytes += handle->getSize();
         }
       }