1 #pragma once 2 3 #include <c10/macros/Export.h> 4 #include <c10/util/UniqueVoidPtr.h> 5 6 #include <atomic> 7 #include <cstdint> 8 #include <memory> 9 #include <shared_mutex> 10 #include <variant> 11 12 namespace c10::impl::cow { 13 14 // A COWDeleterContext object is used as the `ctx` argument for DataPtr 15 // to implement a Copy-on-write (COW) DataPtr. 16 class C10_API COWDeleterContext { 17 public: 18 // Creates an instance, holding the pair of data and original 19 // deleter. 20 // 21 // Note that the deleter will only be called in our destructor if 22 // the last reference to this goes away without getting 23 // materialized. 24 explicit COWDeleterContext(std::unique_ptr<void, DeleterFnPtr> data); 25 26 // Increments the current refcount. 27 void increment_refcount(); 28 29 // See README.md in this directory to understand the locking 30 // strategy. 31 32 // Represents a reference to the context. 33 // 34 // This is returned by decrement_refcount to allow the caller to 35 // copy the data under the shared lock. 36 using NotLastReference = std::shared_lock<std::shared_mutex>; 37 38 // Represents the last reference to the context. 39 // 40 // This will be returned by decrement_refcount when it is the last 41 // reference remaining and after any pending copies have completed. 42 using LastReference = std::unique_ptr<void, DeleterFnPtr>; 43 44 // Decrements the refcount, returning a handle indicating what to 45 // do with it. 46 std::variant<NotLastReference, LastReference> decrement_refcount(); 47 48 private: 49 // The destructor is hidden, this should only ever be used within 50 // UniqueVoidPtr using cow::delete_context as the deleter. 51 ~COWDeleterContext(); 52 53 std::shared_mutex mutex_; 54 std::unique_ptr<void, DeleterFnPtr> data_; 55 std::atomic<std::int64_t> refcount_ = 1; 56 }; 57 58 // `cow_deleter` is used as the `ctx_deleter` for DataPtr to implement a COW 59 // DataPtr. 60 // 61 // Warning: This should only be called on a pointer to a COWDeleterContext that 62 // was allocated on the heap with `new`, because when the refcount reaches 0, 63 // the context is deleted with `delete`. 64 C10_API void cow_deleter(void* ctx); 65 66 } // namespace c10::impl::cow 67