SERVER-32454 Add support for decorations referring to their owners

6 files changed, 53 insertions, 16 deletions

diff --git a/src/mongo/util/decorable.h b/src/mongo/util/decorable.h
index 34fb4a78d87..0545ec2f6ae 100644
--- a/src/mongo/util/decorable.h
+++ b/src/mongo/util/decorable.h
@@ -105,8 +105,13 @@ public:
         return Decoration<T>(getRegistry()->declareDecoration<T>());
     }
 
+    template <typename T>
+    static Decoration<T> declareDecorationWithOwner() {
+        return Decoration<T>(getRegistry()->declareDecorationWithOwner<T, D>());
+    }
+
 protected:
-    Decorable() : _decorations(getRegistry()) {}
+    Decorable() : _decorations(getRegistry(), this) {}
     ~Decorable() = default;
 
 private:
diff --git a/src/mongo/util/decorable_test.cpp b/src/mongo/util/decorable_test.cpp
index c7a234c3c77..356cc9cfb64 100644
--- a/src/mongo/util/decorable_test.cpp
+++ b/src/mongo/util/decorable_test.cpp
@@ -108,10 +108,10 @@ TEST(DecorableTest, SimpleDecoration) {
     const auto dd3 = registry.declareDecoration<int>();
 
     {
-        DecorationContainer decorable1(&registry);
+        DecorationContainer decorable1(&registry, nullptr);
         ASSERT_EQ(2, numConstructedAs);
         ASSERT_EQ(0, numDestructedAs);
-        DecorationContainer decorable2(&registry);
+        DecorationContainer decorable2(&registry, nullptr);
         ASSERT_EQ(4, numConstructedAs);
         ASSERT_EQ(0, numDestructedAs);
 
@@ -144,7 +144,7 @@ TEST(DecorableTest, ThrowingConstructor) {
     registry.declareDecoration<A>();
 
     try {
-        DecorationContainer d(&registry);
+        DecorationContainer d(&registry, nullptr);
     } catch (const AssertionException& ex) {
         ASSERT_EQ(ErrorCodes::Unauthorized, ex.code());
     }
@@ -158,7 +158,7 @@ TEST(DecorableTest, Alignment) {
     const auto firstInt = registry.declareDecoration<int>();
     const auto secondChar = registry.declareDecoration<int>();
     const auto secondInt = registry.declareDecoration<int>();
-    DecorationContainer d(&registry);
+    DecorationContainer d(&registry, nullptr);
     ASSERT_EQ(0U,
               reinterpret_cast<uintptr_t>(&d.getDecoration(firstChar)) %
                   std::alignment_of<char>::value);
@@ -173,5 +173,20 @@ TEST(DecorableTest, Alignment) {
                   std::alignment_of<int>::value);
 }
 
+
+class MyDecorable : public Decorable<MyDecorable> {};
+
+class ClassWithOwnerPtr {
+public:
+    ClassWithOwnerPtr(MyDecorable* decorable) : owner(decorable) {}
+    MyDecorable* owner;
+};
+
+TEST(DecorableTest, DecorationWithOwner) {
+    const auto get = MyDecorable::declareDecorationWithOwner<ClassWithOwnerPtr>();
+    MyDecorable decorable;
+    invariant(get(decorable).owner == &decorable);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/util/decoration_container.cpp b/src/mongo/util/decoration_container.cpp
index a7fc47c07ae..b78cdd92bac 100644
--- a/src/mongo/util/decoration_container.cpp
+++ b/src/mongo/util/decoration_container.cpp
@@ -34,10 +34,10 @@
 
 namespace mongo {
 
-DecorationContainer::DecorationContainer(const DecorationRegistry* registry)
+DecorationContainer::DecorationContainer(const DecorationRegistry* registry, void* owner)
     : _registry(registry),
       _decorationData(new unsigned char[registry->getDecorationBufferSizeBytes()]) {
-    _registry->construct(this);
+    _registry->construct(this, owner);
 }
 
 DecorationContainer::~DecorationContainer() {
diff --git a/src/mongo/util/decoration_container.h b/src/mongo/util/decoration_container.h
index 2fc71c73b17..e4fb89a4cff 100644
--- a/src/mongo/util/decoration_container.h
+++ b/src/mongo/util/decoration_container.h
@@ -87,7 +87,7 @@ public:
      * have any declareDecoration() calls made on it while a DecorationContainer dependent on it
      * is in scope.
      */
-    explicit DecorationContainer(const DecorationRegistry* registry);
+    explicit DecorationContainer(const DecorationRegistry* registry, void* owner);
     ~DecorationContainer();
 
     /**
diff --git a/src/mongo/util/decoration_registry.cpp b/src/mongo/util/decoration_registry.cpp
index f901de67c6a..8ccc838bbe4 100644
--- a/src/mongo/util/decoration_registry.cpp
+++ b/src/mongo/util/decoration_registry.cpp
@@ -47,17 +47,17 @@ DecorationContainer::DecorationDescriptor DecorationRegistry::declareDecoration(
     return result;
 }
 
-void DecorationRegistry::construct(DecorationContainer* decorable) const {
+void DecorationRegistry::construct(DecorationContainer* container, void* owner) const {
     auto iter = _decorationInfo.cbegin();
     try {
         for (; iter != _decorationInfo.cend(); ++iter) {
-            iter->constructor(decorable->getDecoration(iter->descriptor));
+            iter->constructor(container->getDecoration(iter->descriptor), owner);
         }
     } catch (...) {
         try {
             while (iter != _decorationInfo.cbegin()) {
                 --iter;
-                iter->destructor(decorable->getDecoration(iter->descriptor));
+                iter->destructor(container->getDecoration(iter->descriptor));
             }
         } catch (...) {
             std::terminate();
@@ -66,13 +66,13 @@ void DecorationRegistry::construct(DecorationContainer* decorable) const {
     }
 }
 
-void DecorationRegistry::destruct(DecorationContainer* decorable) const {
+void DecorationRegistry::destruct(DecorationContainer* container) const {
     try {
         for (DecorationInfoVector::const_reverse_iterator iter = _decorationInfo.rbegin(),
                                                           end = _decorationInfo.rend();
              iter != end;
              ++iter) {
-            iter->destructor(decorable->getDecoration(iter->descriptor));
+            iter->destructor(container->getDecoration(iter->descriptor));
         }
     } catch (...) {
         std::terminate();
diff --git a/src/mongo/util/decoration_registry.h b/src/mongo/util/decoration_registry.h
index 408d24e1a51..b6a3d879609 100644
--- a/src/mongo/util/decoration_registry.h
+++ b/src/mongo/util/decoration_registry.h
@@ -66,6 +66,18 @@ public:
             sizeof(T), std::alignment_of<T>::value, &constructAt<T>, &destructAt<T>)));
     }
 
+    template <typename T, typename Owner>
+    DecorationContainer::DecorationDescriptorWithType<T> declareDecorationWithOwner() {
+        MONGO_STATIC_ASSERT_MSG(std::is_nothrow_destructible<T>::value,
+                                "Decorations must be nothrow destructible");
+        return DecorationContainer::DecorationDescriptorWithType<T>(
+            std::move(declareDecoration(sizeof(T),
+                                        std::alignment_of<T>::value,
+                                        &constructAtWithOwner<T, Owner>,
+                                        &destructAt<T>)));
+    }
+
+
     size_t getDecorationBufferSizeBytes() const {
         return _totalSizeBytes;
     }
@@ -76,7 +88,7 @@ public:
      *
      * Called by the DecorationContainer constructor. Do not call directly.
      */
-    void construct(DecorationContainer* decorable) const;
+    void construct(DecorationContainer* decorable, void* owner) const;
 
     /**
      * Destroys the decorations declared in this registry on the given instance of "decorable".
@@ -89,7 +101,7 @@ private:
     /**
      * Function that constructs (initializes) a single instance of a decoration.
      */
-    using DecorationConstructorFn = stdx::function<void(void*)>;
+    using DecorationConstructorFn = stdx::function<void(void*, void*)>;
 
     /**
      * Function that destructs (deinitializes) a single instance of a decoration.
@@ -110,10 +122,15 @@ private:
     using DecorationInfoVector = std::vector<DecorationInfo>;
 
     template <typename T>
-    static void constructAt(void* location) {
+    static void constructAt(void* location, void* owner) {
         new (location) T();
     }
 
+    template <typename T, typename Owner>
+    static void constructAtWithOwner(void* location, void* owner) {
+        new (location) T(static_cast<Owner*>(owner));
+    }
+
     template <typename T>
     static void destructAt(void* location) {
         static_cast<T*>(location)->~T();