folly/concurrency/test/AtomicSharedPtrCounted.h

   1 /*
   2  * Copyright 2017-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #pragma once
  17
  18 struct counted_shared_tag {};
  19 template <template <typename> class Atom = std::atomic>
  20 struct intrusive_shared_count {
  21   intrusive_shared_count() {
  22     counts.store(0);
  23   }
  24   void add_ref(uint64_t count = 1) {
  25     counts.fetch_add(count);
  26   }
  27
  28   uint64_t release_ref(uint64_t count = 1) {
  29     return counts.fetch_sub(count);
  30   }
  31   Atom<uint64_t> counts;
  32 };
  33
  34 template <template <typename> class Atom = std::atomic>
  35 struct counted_ptr_base {
  36  protected:
  37   static intrusive_shared_count<Atom>* getRef(void* pt) {
  38     char* p = (char*)pt;
  39     p -= sizeof(intrusive_shared_count<Atom>);
  40     return (intrusive_shared_count<Atom>*)p;
  41   }
  42 };
  43
  44 // basically shared_ptr, but only supports make_counted, and provides
  45 // access to add_ref / release_ref with a count.  Alias not supported.
  46 template <typename T, template <typename> class Atom = std::atomic>
  47 class counted_ptr : public counted_ptr_base<Atom> {
  48  public:
  49   T* p_;
  50   counted_ptr() : p_(nullptr) {}
  51   counted_ptr(counted_shared_tag, T* p) : p_(p) {
  52     if (p_) {
  53       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  54     }
  55   }
  56
  57   counted_ptr(const counted_ptr& o) : p_(o.p_) {
  58     if (p_) {
  59       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  60     }
  61   }
  62   counted_ptr& operator=(const counted_ptr& o) {
  63     if (p_ && counted_ptr_base<Atom>::getRef(p_)->release_ref() == 1) {
  64       p_->~T();
  65       free(counted_ptr_base<Atom>::getRef(p_));
  66     }
  67     p_ = o.p_;
  68     if (p_) {
  69       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  70     }
  71     return *this;
  72   }
  73   explicit counted_ptr(T* p) : p_(p) {
  74     CHECK(!p);
  75   }
  76   ~counted_ptr() {
  77     if (p_ && counted_ptr_base<Atom>::getRef(p_)->release_ref() == 1) {
  78       p_->~T();
  79       free(counted_ptr_base<Atom>::getRef(p_));
  80     }
  81   }
  82   typename std::add_lvalue_reference<T>::type operator*() const {
  83     return *p_;
  84   }
  85
  86   T* get() const {
  87     return p_;
  88   }
  89   T* operator->() const {
  90     return p_;
  91   }
  92   explicit operator bool() const {
  93     return p_ == nullptr ? false : true;
  94   }
  95   bool operator==(const counted_ptr<T, Atom>& p) const {
  96     return get() == p.get();
  97   }
  98 };
  99
 100 template <
 101     template <typename> class Atom = std::atomic,
 102     typename T,
 103     typename... Args>
 104 counted_ptr<T, Atom> make_counted(Args&&... args) {
 105   char* mem = (char*)malloc(sizeof(T) + sizeof(intrusive_shared_count<Atom>));
 106   if (!mem) {
 107     throw std::bad_alloc();
 108   }
 109   new (mem) intrusive_shared_count<Atom>();
 110   T* ptr = (T*)(mem + sizeof(intrusive_shared_count<Atom>));
 111   new (ptr) T(std::forward<Args>(args)...);
 112   return counted_ptr<T, Atom>(counted_shared_tag(), ptr);
 113 }
 114
 115 template <template <typename> class Atom = std::atomic>
 116 class counted_ptr_internals : public counted_ptr_base<Atom> {
 117  public:
 118   template <typename T, typename... Args>
 119   static counted_ptr<T, Atom> make_ptr(Args&&... args) {
 120     return make_counted<Atom, T>(std::forward<Args...>(args...));
 121   }
 122   template <typename T>
 123   using CountedPtr = counted_ptr<T, Atom>;
 124   typedef void counted_base;
 125
 126   template <typename T>
 127   static counted_base* get_counted_base(const counted_ptr<T, Atom>& bar) {
 128     return bar.p_;
 129   }
 130
 131   template <typename T>
 132   static T* get_shared_ptr(counted_base* base) {
 133     return (T*)base;
 134   }
 135
 136   template <typename T>
 137   static T* release_ptr(counted_ptr<T, Atom>& p) {
 138     auto res = p.p_;
 139     p.p_ = nullptr;
 140     return res;
 141   }
 142
 143   template <typename T>
 144   static counted_ptr<T, Atom> get_shared_ptr_from_counted_base(
 145       counted_base* base,
 146       bool inc = true) {
 147     auto res = counted_ptr<T, Atom>(counted_shared_tag(), (T*)(base));
 148     if (!inc) {
 149       release_shared<T>(base, 1);
 150     }
 151     return res;
 152   }
 153
 154   static void inc_shared_count(counted_base* base, int64_t count) {
 155     counted_ptr_base<Atom>::getRef(base)->add_ref(count);
 156   }
 157
 158   template <typename T>
 159   static void release_shared(counted_base* base, uint64_t count) {
 160     if (count == counted_ptr_base<Atom>::getRef(base)->release_ref(count)) {
 161       ((T*)base)->~T();
 162       free(counted_ptr_base<Atom>::getRef(base));
 163     }
 164   }
 165 };