SERVER-50273 Add DecorableCopyable to support copyable types where decorations are also copied

4 files changed, 370 insertions, 68 deletions

diff --git a/src/mongo/util/decorable.h b/src/mongo/util/decorable.h
index b8191ac6ebf..0b7a082cb17 100644
--- a/src/mongo/util/decorable.h
+++ b/src/mongo/util/decorable.h
@@ -145,4 +145,88 @@ private:
     DecorationContainer<D> _decorations;
 };
 
+template <typename D>
+class DecorableCopyable {
+public:
+    template <typename T>
+    class Decoration {
+    public:
+        Decoration() = delete;
+
+        T& operator()(D& d) const {
+            return static_cast<DecorableCopyable&>(d)._decorations.getDecoration(this->_raw);
+        }
+
+        T& operator()(D* const d) const {
+            return (*this)(*d);
+        }
+
+        const T& operator()(const D& d) const {
+            return static_cast<const DecorableCopyable&>(d)._decorations.getDecoration(this->_raw);
+        }
+
+        const T& operator()(const D* const d) const {
+            return (*this)(*d);
+        }
+
+        const D* owner(const T* const t) const {
+            return static_cast<const D*>(getOwnerImpl(t));
+        }
+
+        D* owner(T* const t) const {
+            return static_cast<D*>(getOwnerImpl(t));
+        }
+
+        const D& owner(const T& t) const {
+            return *owner(&t);
+        }
+
+        D& owner(T& t) const {
+            return *owner(&t);
+        }
+
+    private:
+        const DecorableCopyable* getOwnerImpl(const T* const t) const {
+            return *reinterpret_cast<const DecorableCopyable* const*>(
+                reinterpret_cast<const unsigned char* const>(t) - _raw._raw._index);
+        }
+
+        DecorableCopyable* getOwnerImpl(T* const t) const {
+            return const_cast<DecorableCopyable*>(getOwnerImpl(const_cast<const T*>(t)));
+        }
+
+        friend class DecorableCopyable;
+
+        explicit Decoration(
+            typename DecorationContainer<D>::template DecorationDescriptorWithType<T> raw)
+            : _raw(std::move(raw)) {}
+
+        typename DecorationContainer<D>::template DecorationDescriptorWithType<T> _raw;
+    };
+
+    template <typename T>
+    static Decoration<T> declareDecoration() {
+        return Decoration<T>(getRegistry()->template declareDecorationCopyable<T>());
+    }
+
+protected:
+    DecorableCopyable() : _decorations(this, getRegistry()) {}
+    ~DecorableCopyable() = default;
+
+    DecorableCopyable(const DecorableCopyable& other)
+        : _decorations(this, getRegistry(), other._decorations) {}
+    DecorableCopyable& operator=(const DecorableCopyable& rhs) {
+        getRegistry()->copyAssign(&_decorations, &rhs._decorations);
+        return *this;
+    }
+
+private:
+    static DecorationRegistry<D>* getRegistry() {
+        static DecorationRegistry<D>* theRegistry = new DecorationRegistry<D>();
+        return theRegistry;
+    }
+
+    DecorationContainer<D> _decorations;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/util/decorable_test.cpp b/src/mongo/util/decorable_test.cpp
index 3d388849506..99e5a2af501 100644
--- a/src/mongo/util/decorable_test.cpp
+++ b/src/mongo/util/decorable_test.cpp
@@ -43,6 +43,8 @@ namespace mongo {
 namespace {
 
 static int numConstructedAs;
+static int numCopyConstructedAs;
+static int numCopyAssignedAs;
 static int numDestructedAs;
 
 class A {
@@ -50,6 +52,14 @@ public:
     A() : value(0) {
         ++numConstructedAs;
     }
+    A(const A& other) : value(other.value) {
+        ++numCopyConstructedAs;
+    }
+    A& operator=(const A& rhs) {
+        value = rhs.value;
+        ++numCopyAssignedAs;
+        return *this;
+    }
     ~A() {
         ++numDestructedAs;
     }
@@ -66,6 +76,7 @@ public:
 };
 
 class MyDecorable : public Decorable<MyDecorable> {};
+class MyCopyableDecorable : public DecorableCopyable<MyCopyableDecorable> {};
 
 TEST(DecorableTest, DecorableType) {
     const auto dd1 = MyDecorable::declareDecoration<A>();
@@ -101,77 +112,101 @@ TEST(DecorableTest, DecorableType) {
 }
 
 TEST(DecorableTest, SimpleDecoration) {
-    numConstructedAs = 0;
-    numDestructedAs = 0;
-    DecorationRegistry<MyDecorable> registry;
-    const auto dd1 = registry.declareDecoration<A>();
-    const auto dd2 = registry.declareDecoration<A>();
-    const auto dd3 = registry.declareDecoration<int>();
-
-    {
-        DecorationContainer<MyDecorable> decorable1(nullptr, &registry);
-        ASSERT_EQ(2, numConstructedAs);
-        ASSERT_EQ(0, numDestructedAs);
-        DecorationContainer<MyDecorable> decorable2(nullptr, &registry);
-        ASSERT_EQ(4, numConstructedAs);
-        ASSERT_EQ(0, numDestructedAs);
-
-        ASSERT_EQ(0, decorable1.getDecoration(dd1).value);
-        ASSERT_EQ(0, decorable1.getDecoration(dd2).value);
-        ASSERT_EQ(0, decorable2.getDecoration(dd1).value);
-        ASSERT_EQ(0, decorable2.getDecoration(dd2).value);
-        ASSERT_EQ(0, decorable2.getDecoration(dd3));
-        decorable1.getDecoration(dd1).value = 1;
-        decorable1.getDecoration(dd2).value = 2;
-        decorable2.getDecoration(dd1).value = 3;
-        decorable2.getDecoration(dd2).value = 4;
-        decorable2.getDecoration(dd3) = 5;
-        ASSERT_EQ(1, decorable1.getDecoration(dd1).value);
-        ASSERT_EQ(2, decorable1.getDecoration(dd2).value);
-        ASSERT_EQ(3, decorable2.getDecoration(dd1).value);
-        ASSERT_EQ(4, decorable2.getDecoration(dd2).value);
-        ASSERT_EQ(5, decorable2.getDecoration(dd3));
-    }
-    ASSERT_EQ(4, numDestructedAs);
+    auto test = [](auto decorable) {
+        using decorable_t = decltype(decorable);
+
+        numConstructedAs = 0;
+        numDestructedAs = 0;
+        DecorationRegistry<decorable_t> registry;
+        const auto dd1 = registry.template declareDecoration<A>();
+        const auto dd2 = registry.template declareDecoration<A>();
+        const auto dd3 = registry.template declareDecoration<int>();
+
+        {
+            decorable_t* decorablePtr = nullptr;
+            DecorationContainer<decorable_t> decorable1(decorablePtr, &registry);
+            ASSERT_EQ(2, numConstructedAs);
+            ASSERT_EQ(0, numDestructedAs);
+            DecorationContainer<decorable_t> decorable2(decorablePtr, &registry);
+            ASSERT_EQ(4, numConstructedAs);
+            ASSERT_EQ(0, numDestructedAs);
+
+            ASSERT_EQ(0, decorable1.getDecoration(dd1).value);
+            ASSERT_EQ(0, decorable1.getDecoration(dd2).value);
+            ASSERT_EQ(0, decorable2.getDecoration(dd1).value);
+            ASSERT_EQ(0, decorable2.getDecoration(dd2).value);
+            ASSERT_EQ(0, decorable2.getDecoration(dd3));
+            decorable1.getDecoration(dd1).value = 1;
+            decorable1.getDecoration(dd2).value = 2;
+            decorable2.getDecoration(dd1).value = 3;
+            decorable2.getDecoration(dd2).value = 4;
+            decorable2.getDecoration(dd3) = 5;
+            ASSERT_EQ(1, decorable1.getDecoration(dd1).value);
+            ASSERT_EQ(2, decorable1.getDecoration(dd2).value);
+            ASSERT_EQ(3, decorable2.getDecoration(dd1).value);
+            ASSERT_EQ(4, decorable2.getDecoration(dd2).value);
+            ASSERT_EQ(5, decorable2.getDecoration(dd3));
+        }
+        ASSERT_EQ(4, numDestructedAs);
+    };
+
+    test(MyDecorable());
+    test(MyCopyableDecorable());
 }
 
 TEST(DecorableTest, ThrowingConstructor) {
-    numConstructedAs = 0;
-    numDestructedAs = 0;
-
-    DecorationRegistry<MyDecorable> registry;
-    registry.declareDecoration<A>();
-    registry.declareDecoration<ThrowA>();
-    registry.declareDecoration<A>();
-
-    try {
-        DecorationContainer<MyDecorable> d(nullptr, &registry);
-    } catch (const AssertionException& ex) {
-        ASSERT_EQ(ErrorCodes::Unauthorized, ex.code());
-    }
-    ASSERT_EQ(1, numConstructedAs);
-    ASSERT_EQ(1, numDestructedAs);
+    auto test = [](auto decorable) {
+        using decorable_t = decltype(decorable);
+
+        numConstructedAs = 0;
+        numDestructedAs = 0;
+
+        DecorationRegistry<decorable_t> registry;
+        registry.template declareDecoration<A>();
+        registry.template declareDecoration<ThrowA>();
+        registry.template declareDecoration<A>();
+
+        try {
+            decorable_t* decorablePtr = nullptr;
+            DecorationContainer<decorable_t> d(decorablePtr, &registry);
+        } catch (const AssertionException& ex) {
+            ASSERT_EQ(ErrorCodes::Unauthorized, ex.code());
+        }
+        ASSERT_EQ(1, numConstructedAs);
+        ASSERT_EQ(1, numDestructedAs);
+    };
+
+    test(MyDecorable());
+    test(MyCopyableDecorable());
 }
 
 TEST(DecorableTest, Alignment) {
-    DecorationRegistry<MyDecorable> registry;
-    const auto firstChar = registry.declareDecoration<char>();
-    const auto firstInt = registry.declareDecoration<int>();
-    const auto secondChar = registry.declareDecoration<int>();
-    const auto secondInt = registry.declareDecoration<int>();
-    DecorationContainer<MyDecorable> d(nullptr, &registry);
-    ASSERT_EQ(0U,
-              reinterpret_cast<uintptr_t>(&d.getDecoration(firstChar)) %
-                  std::alignment_of<char>::value);
-    ASSERT_EQ(0U,
-              reinterpret_cast<uintptr_t>(&d.getDecoration(secondChar)) %
-                  std::alignment_of<char>::value);
-    ASSERT_EQ(0U,
-              reinterpret_cast<uintptr_t>(&d.getDecoration(firstInt)) %
-                  std::alignment_of<int>::value);
-    ASSERT_EQ(0U,
-              reinterpret_cast<uintptr_t>(&d.getDecoration(secondInt)) %
-                  std::alignment_of<int>::value);
+    auto test = [](auto decorable) {
+        using decorable_t = decltype(decorable);
+
+        DecorationRegistry<decorable_t> registry;
+        const auto firstChar = registry.template declareDecoration<char>();
+        const auto firstInt = registry.template declareDecoration<int>();
+        const auto secondChar = registry.template declareDecoration<int>();
+        const auto secondInt = registry.template declareDecoration<int>();
+        decorable_t* decorablePtr = nullptr;
+        DecorationContainer<decorable_t> d(decorablePtr, &registry);
+        ASSERT_EQ(0U,
+                  reinterpret_cast<uintptr_t>(&d.getDecoration(firstChar)) %
+                      std::alignment_of<char>::value);
+        ASSERT_EQ(0U,
+                  reinterpret_cast<uintptr_t>(&d.getDecoration(secondChar)) %
+                      std::alignment_of<char>::value);
+        ASSERT_EQ(0U,
+                  reinterpret_cast<uintptr_t>(&d.getDecoration(firstInt)) %
+                      std::alignment_of<int>::value);
+        ASSERT_EQ(0U,
+                  reinterpret_cast<uintptr_t>(&d.getDecoration(secondInt)) %
+                      std::alignment_of<int>::value);
+    };
+
+    test(MyDecorable());
+    test(MyCopyableDecorable());
 }
 
 struct DecoratedOwnerChecker : public Decorable<DecoratedOwnerChecker> {
@@ -239,5 +274,33 @@ TEST(DecorableTest, DecorationWithOwner) {
     ASSERT_EQ(&owner, &DecorationWithOwner::get.owner(decoration));
 }
 
+TEST(DecorableTest, DecorableCopyableType) {
+    numCopyConstructedAs = 0;
+    numCopyAssignedAs = 0;
+    const auto dd1 = MyCopyableDecorable::declareDecoration<A>();
+    const auto dd2 = MyCopyableDecorable::declareDecoration<int>();
+    {
+        MyCopyableDecorable decorable1;
+        dd1(decorable1).value = 1;
+        dd2(decorable1) = 2;
+
+        MyCopyableDecorable decorable2(decorable1);
+        ASSERT_EQ(1, numCopyConstructedAs);
+        ASSERT_EQ(0, numCopyAssignedAs);
+        ASSERT_EQ(dd1(decorable1).value, dd1(decorable2).value);
+        ASSERT_EQ(dd2(decorable1), dd2(decorable2));
+
+        MyCopyableDecorable decorable3;
+        ASSERT_NE(dd1(decorable1).value, dd1(decorable3).value);
+        ASSERT_NE(dd2(decorable1), dd2(decorable3));
+
+        decorable3 = decorable1;
+        ASSERT_EQ(1, numCopyConstructedAs);
+        ASSERT_EQ(1, numCopyAssignedAs);
+        ASSERT_EQ(dd1(decorable1).value, dd1(decorable3).value);
+        ASSERT_EQ(dd2(decorable1), dd2(decorable3));
+    }
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/util/decoration_container.h b/src/mongo/util/decoration_container.h
index e144ebac5ae..373b199e056 100644
--- a/src/mongo/util/decoration_container.h
+++ b/src/mongo/util/decoration_container.h
@@ -40,6 +40,9 @@ class DecorationRegistry;
 template <typename DecoratedType>
 class Decorable;
 
+template <typename DecoratedType>
+class DecorableCopyable;
+
 /**
  * An container for decorations.
  */
@@ -61,6 +64,7 @@ public:
         friend DecorationContainer;
         friend DecorationRegistry<DecoratedType>;
         friend Decorable<DecoratedType>;
+        friend DecorableCopyable<DecoratedType>;
 
         explicit DecorationDescriptor(size_t index) : _index(index) {}
 
@@ -81,6 +85,7 @@ public:
         friend DecorationContainer;
         friend DecorationRegistry<DecoratedType>;
         friend Decorable<DecoratedType>;
+        friend DecorableCopyable<DecoratedType>;
 
         explicit DecorationDescriptorWithType(DecorationDescriptor raw) : _raw(std::move(raw)) {}
 
@@ -109,6 +114,47 @@ public:
         _registry->construct(this);
     }
 
+    /**
+     * Constructs a copyable decorable built based on the given "registry."
+     *
+     * See above for more details
+     */
+    explicit DecorationContainer(DecorableCopyable<DecoratedType>* const decorated,
+                                 const DecorationRegistry<DecoratedType>* const registry)
+        : _registry(registry),
+          _decorationData(new unsigned char[registry->getDecorationBufferSizeBytes()]) {
+        // Because the decorations live in the externally allocated storage buffer at
+        // `_decorationData`, there needs to be a way to get back from a known location within this
+        // buffer to the type which owns those decorations.  We place a pointer to ourselves, a
+        // "back link" in the front of this storage buffer, as this is the easiest "well known
+        // location" to compute.
+        DecorableCopyable<DecoratedType>** const backLink =
+            reinterpret_cast<DecorableCopyable<DecoratedType>**>(_decorationData.get());
+        *backLink = decorated;
+        _registry->construct(this);
+    }
+
+    /**
+     * Constructs a copyable decorable built based on the given "registry."
+     *
+     * All decorations are copy constructed from provided DecorationContainer.
+     */
+    explicit DecorationContainer(DecorableCopyable<DecoratedType>* const decorated,
+                                 const DecorationRegistry<DecoratedType>* const registry,
+                                 const DecorationContainer& other)
+        : _registry(registry),
+          _decorationData(new unsigned char[registry->getDecorationBufferSizeBytes()]) {
+        // Because the decorations live in the externally allocated storage buffer at
+        // `_decorationData`, there needs to be a way to get back from a known location within this
+        // buffer to the type which owns those decorations.  We place a pointer to ourselves, a
+        // "back link" in the front of this storage buffer, as this is the easiest "well known
+        // location" to compute.
+        DecorableCopyable<DecoratedType>** const backLink =
+            reinterpret_cast<DecorableCopyable<DecoratedType>**>(_decorationData.get());
+        *backLink = decorated;
+        _registry->copyConstruct(this, &other);
+    }
+
     ~DecorationContainer() {
         _registry->destroy(this);
     }
diff --git a/src/mongo/util/decoration_registry.h b/src/mongo/util/decoration_registry.h
index 6e4f8ed55db..c434b0c350e 100644
--- a/src/mongo/util/decoration_registry.h
+++ b/src/mongo/util/decoration_registry.h
@@ -69,8 +69,34 @@ public:
                                 "Decorations must be nothrow destructible");
         return
             typename DecorationContainer<DecoratedType>::template DecorationDescriptorWithType<T>(
-                std::move(declareDecoration(
-                    sizeof(T), std::alignment_of<T>::value, &constructAt<T>, &destroyAt<T>)));
+                std::move(declareDecoration(sizeof(T),
+                                            std::alignment_of<T>::value,
+                                            &constructAt<T>,
+                                            nullptr,
+                                            nullptr,
+                                            &destroyAt<T>)));
+    }
+
+    /**
+     * Declares a copyable decoration of type T, constructed with T's default constructor, and
+     * returns a descriptor for accessing that decoration.
+     *
+     * It also binds T's copy constructor and copy assignment operator.
+     *
+     * NOTE: T's destructor must not throw exceptions.
+     */
+    template <typename T>
+    auto declareDecorationCopyable() {
+        MONGO_STATIC_ASSERT_MSG(std::is_nothrow_destructible<T>::value,
+                                "Decorations must be nothrow destructible");
+        return
+            typename DecorationContainer<DecoratedType>::template DecorationDescriptorWithType<T>(
+                std::move(declareDecoration(sizeof(T),
+                                            std::alignment_of<T>::value,
+                                            &constructAt<T>,
+                                            &copyConstructAt<T>,
+                                            &copyAssignAt<T>,
+                                            &destroyAt<T>)));
     }
 
     size_t getDecorationBufferSizeBytes() const {
@@ -110,6 +136,60 @@ public:
     }
 
     /**
+     * Copy constructs the decorations declared in this registry on the given instance of
+     * "decorable" from another DecorationContainer.
+     *
+     * Called by the DecorationContainer constructor. Do not call directly.
+     */
+    void copyConstruct(DecorationContainer<DecoratedType>* const container,
+                       const DecorationContainer<DecoratedType>* const other) const {
+        using std::cbegin;
+
+        auto iter = cbegin(_decorationInfo);
+
+        auto cleanupFunction = [&iter, container, this ]() noexcept->void {
+            using std::crend;
+            std::for_each(std::make_reverse_iterator(iter),
+                          crend(this->_decorationInfo),
+                          [&](auto&& decoration) {
+                              decoration.destructor(
+                                  container->getDecoration(decoration.descriptor));
+                          });
+        };
+
+        auto cleanup = makeGuard(std::move(cleanupFunction));
+
+        using std::cend;
+
+        for (; iter != cend(_decorationInfo); ++iter) {
+            iter->copyConstructor(container->getDecoration(iter->descriptor),
+                                  other->getDecoration(iter->descriptor));
+        }
+
+        cleanup.dismiss();
+    }
+
+    /**
+     * Copy assigns the decorations declared in this registry on the given instance of
+     * "decorable" from another DecorationContainer.
+     *
+     * Called by the DecorableCopyable copy assignment operator. Do not call directly.
+     */
+    void copyAssign(DecorationContainer<DecoratedType>* const container,
+                    const DecorationContainer<DecoratedType>* const rhs) const {
+        using std::cbegin;
+
+        auto iter = cbegin(_decorationInfo);
+
+        using std::cend;
+
+        for (; iter != cend(_decorationInfo); ++iter) {
+            iter->copyAssignment(container->getDecoration(iter->descriptor),
+                                 rhs->getDecoration(iter->descriptor));
+        }
+    }
+
+    /**
      * Destroys the decorations declared in this registry on the given instance of "decorable".
      *
      * Called by the DecorationContainer destructor.  Do not call directly.
@@ -129,6 +209,16 @@ private:
     using DecorationConstructorFn = void (*)(void*);
 
     /**
+     * Function that copy constructs a single instance of a decoration from another instance.
+     */
+    using DecorationCopyConstructorFn = void (*)(void*, const void*);
+
+    /**
+     * Function that copy assigns a single instance of a decoration from another instance.
+     */
+    using DecorationCopyAssignmentFn = void (*)(void*, const void*);
+
+    /**
      * Function that destroys (deinitializes) a single instance of a decoration.
      */
     using DecorationDestructorFn = void (*)(void*);
@@ -138,13 +228,19 @@ private:
         DecorationInfo(
             typename DecorationContainer<DecoratedType>::DecorationDescriptor inDescriptor,
             DecorationConstructorFn inConstructor,
+            DecorationCopyConstructorFn inCopyConstructor,
+            DecorationCopyAssignmentFn inCopyAssignment,
             DecorationDestructorFn inDestructor)
             : descriptor(std::move(inDescriptor)),
               constructor(std::move(inConstructor)),
+              copyConstructor(std::move(inCopyConstructor)),
+              copyAssignment(std::move(inCopyAssignment)),
               destructor(std::move(inDestructor)) {}
 
         typename DecorationContainer<DecoratedType>::DecorationDescriptor descriptor;
         DecorationConstructorFn constructor;
+        DecorationCopyConstructorFn copyConstructor;
+        DecorationCopyAssignmentFn copyAssignment;
         DecorationDestructorFn destructor;
     };
 
@@ -156,6 +252,16 @@ private:
     }
 
     template <typename T>
+    static void copyConstructAt(void* location, const void* other) {
+        new (location) T(*static_cast<const T*>(other));
+    }
+
+    template <typename T>
+    static void copyAssignAt(void* location, const void* other) {
+        *static_cast<T*>(location) = *static_cast<const T*>(other);
+    }
+
+    template <typename T>
     static void destroyAt(void* location) {
         static_cast<T*>(location)->~T();
     }
@@ -170,13 +276,16 @@ private:
         const size_t sizeBytes,
         const size_t alignBytes,
         const DecorationConstructorFn constructor,
+        const DecorationCopyConstructorFn copyConstructor,
+        const DecorationCopyAssignmentFn copyAssignment,
         const DecorationDestructorFn destructor) {
         const size_t misalignment = _totalSizeBytes % alignBytes;
         if (misalignment) {
             _totalSizeBytes += alignBytes - misalignment;
         }
         typename DecorationContainer<DecoratedType>::DecorationDescriptor result(_totalSizeBytes);
-        _decorationInfo.push_back(DecorationInfo(result, constructor, destructor));
+        _decorationInfo.push_back(
+            DecorationInfo(result, constructor, copyConstructor, copyAssignment, destructor));
         _totalSizeBytes += sizeBytes;
         return result;
     }