SERVER-67251 strictly bind SSM::Impl lifetime to its SSM

3 files changed, 36 insertions, 18 deletions

diff --git a/src/mongo/transport/service_entry_point_impl.cpp b/src/mongo/transport/service_entry_point_impl.cpp
index 1ab09091481..524dc94a964 100644
--- a/src/mongo/transport/service_entry_point_impl.cpp
+++ b/src/mongo/transport/service_entry_point_impl.cpp
@@ -169,11 +169,10 @@ size_t getSupportedMax() {
 class ServiceEntryPointImpl::Sessions {
 public:
     struct Entry {
-        explicit Entry(std::unique_ptr<transport::ServiceStateMachine> ssm) : ssm{std::move(ssm)} {}
-        std::unique_ptr<transport::ServiceStateMachine> ssm;
+        explicit Entry(std::shared_ptr<transport::ServiceStateMachine> ssm) : ssm{std::move(ssm)} {}
+        std::shared_ptr<transport::ServiceStateMachine> ssm;
         ClientSummary summary{ssm->client()};
     };
-
     using ByClientMap = stdx::unordered_map<Client*, Entry>;
     using iterator = ByClientMap::iterator;
 
@@ -199,7 +198,7 @@ public:
                 _lk, deadline.toSystemTimePoint(), [&] { return _src->_byClient.empty(); });
         }
 
-        iterator insert(std::unique_ptr<transport::ServiceStateMachine> ssm) {
+        iterator insert(std::shared_ptr<transport::ServiceStateMachine> ssm) {
             Client* cli = ssm->client();
             auto [it, ok] = _src->_byClient.insert({cli, Entry(std::move(ssm))});
             invariant(ok);
@@ -313,7 +312,7 @@ void ServiceEntryPointImpl::startSession(transport::SessionHandle session) {
             transport::ServiceExecutorContext::set(&*client, std::move(seCtx));
         }
 
-        auto ssm = std::make_unique<transport::ServiceStateMachine>(std::move(client));
+        auto ssm = transport::ServiceStateMachine::make(std::move(client));
         iter = sync.insert(std::move(ssm));
         if (!quiet()) {
             LOGV2(22943,
diff --git a/src/mongo/transport/service_state_machine.cpp b/src/mongo/transport/service_state_machine.cpp
index 15e71fe89ac..6a8050bfbf2 100644
--- a/src/mongo/transport/service_state_machine.cpp
+++ b/src/mongo/transport/service_state_machine.cpp
@@ -121,8 +121,7 @@ Message makeExhaustMessage(Message requestMsg, DbResponse* dbresponse) {
 }
 }  // namespace
 
-class ServiceStateMachine::Impl final
-    : public std::enable_shared_from_this<ServiceStateMachine::Impl> {
+class ServiceStateMachine::Impl {
 public:
     /*
      * Any state may transition to EndSession in case of an error, otherwise the valid state
@@ -142,8 +141,9 @@ public:
                      // state() if this is the current state.
     };
 
-    explicit Impl(ServiceContext::UniqueClient client)
-        : _state{State::Created},
+    Impl(ServiceStateMachine* ssm, ServiceContext::UniqueClient client)
+        : _ssm{ssm},
+          _state{State::Created},
           _serviceContext{client->getServiceContext()},
           _sep{_serviceContext->getServiceEntryPoint()},
           _clientStrand{ClientStrand::make(std::move(client))} {}
@@ -229,6 +229,13 @@ public:
     }
 
 private:
+    /** Alias: refers to this Impl, but holds a ref to the enclosing SSM. */
+    std::shared_ptr<Impl> shared_from_this() {
+        return {_ssm->shared_from_this(), this};
+    }
+
+    ServiceStateMachine* const _ssm;
+
     AtomicWord<State> _state{State::Created};
 
     ServiceContext* const _serviceContext;
@@ -544,8 +551,10 @@ void ServiceStateMachine::Impl::cleanupSession(const Status& status) {
     invariant(_state.swap(State::Ended) != State::Ended);
 }
 
-ServiceStateMachine::ServiceStateMachine(ServiceContext::UniqueClient client)
-    : _impl{std::make_shared<Impl>(std::move(client))} {}
+ServiceStateMachine::ServiceStateMachine(PassKeyTag, ServiceContext::UniqueClient client)
+    : _impl{std::make_unique<Impl>(this, std::move(client))} {}
+
+ServiceStateMachine::~ServiceStateMachine() = default;
 
 Client* ServiceStateMachine::client() const {
     return _impl->client();
diff --git a/src/mongo/transport/service_state_machine.h b/src/mongo/transport/service_state_machine.h
index c830abba1b3..71ed2be3546 100644
--- a/src/mongo/transport/service_state_machine.h
+++ b/src/mongo/transport/service_state_machine.h
@@ -44,8 +44,15 @@ namespace transport {
  * lifecycle of each user request as a state machine. It is the glue between the stateless
  * ServiceEntryPoint and TransportLayer that ties network and database logic together for a
  * user.
+ *
+ * A `ServiceStateMachine` must be managed by a `shared_ptr`, so we force all instances
+ * to be created by the static `make` function.
  */
-class ServiceStateMachine {
+class ServiceStateMachine : public std::enable_shared_from_this<ServiceStateMachine> {
+    struct PassKeyTag {
+        explicit PassKeyTag() = default;
+    };
+    class Impl;
     ServiceStateMachine(ServiceStateMachine&) = delete;
     ServiceStateMachine& operator=(ServiceStateMachine&) = delete;
 
@@ -53,12 +60,15 @@ class ServiceStateMachine {
     ServiceStateMachine& operator=(ServiceStateMachine&&) = delete;
 
 public:
-    class Impl;
+    /** Factory function: The only public way to create instances. */
+    static std::shared_ptr<ServiceStateMachine> make(ServiceContext::UniqueClient client) {
+        return std::make_shared<ServiceStateMachine>(PassKeyTag{}, std::move(client));
+    }
 
-    /*
-     * Construct a ServiceStateMachine for a given Client.
-     */
-    explicit ServiceStateMachine(ServiceContext::UniqueClient client);
+    /** Public must use `make` to create instances. */
+    ServiceStateMachine(PassKeyTag, ServiceContext::UniqueClient client);
+
+    ~ServiceStateMachine();
 
     /** Returns the Client given in the constructor. */
     Client* client() const;
@@ -85,7 +95,7 @@ public:
     void terminateIfTagsDontMatch(transport::Session::TagMask tags);
 
 private:
-    std::shared_ptr<Impl> _impl;
+    std::unique_ptr<Impl> _impl;
 };
 
 }  // namespace transport