SERVER-62188 fix memory corruption in the DeadlineFuture

(cherry picked from commit 58d452b07f5f09781a5e3d8ed4e1644c510a3815)

3 files changed, 40 insertions, 25 deletions

diff --git a/src/mongo/db/process_health/deadline_future.h b/src/mongo/db/process_health/deadline_future.h
index ee5768b03c0..5b9e549de0c 100644
--- a/src/mongo/db/process_health/deadline_future.h
+++ b/src/mongo/db/process_health/deadline_future.h
@@ -42,19 +42,45 @@ namespace process_health {
  * - otherwise the outputFuture returns an error.
  */
 template <typename ResultStatus>
-class DeadlineFuture {
+class DeadlineFuture : public std::enable_shared_from_this<DeadlineFuture<ResultStatus>> {
 public:
-    DeadlineFuture(std::shared_ptr<executor::TaskExecutor> executor,
-                   Future<ResultStatus> inputFuture,
-                   Milliseconds timeout)
-        : _executor(executor) {
+    static std::shared_ptr<DeadlineFuture<ResultStatus>> create(
+        std::shared_ptr<executor::TaskExecutor> executor,
+        Future<ResultStatus> inputFuture,
+        Milliseconds timeout) {
+        auto instance = std::shared_ptr<DeadlineFuture<ResultStatus>>(
+            new DeadlineFuture<ResultStatus>(executor));
+        instance->init(instance, std::move(inputFuture), timeout);
+        return instance;
+    }
+
+    ~DeadlineFuture() {
+        auto lk = stdx::lock_guard(_mutex);
+        _executor->cancel(_timeoutCbHandle.get());
+        // The _executor holds the shared ptr on this, the callback will set the promise.
+        invariant(get().isReady());
+    }
+
+    SharedSemiFuture<ResultStatus> get() const {
+        return _outputFuturePromise->getFuture();
+    }
+
+private:
+    DeadlineFuture(std::shared_ptr<executor::TaskExecutor> executor) : _executor(executor) {}
+
+    void init(std::shared_ptr<DeadlineFuture<ResultStatus>> self,
+              Future<ResultStatus> inputFuture,
+              Milliseconds timeout) {
         _outputFuturePromise = std::make_unique<SharedPromise<ResultStatus>>();
 
         auto swCbHandle = _executor->scheduleWorkAt(
-            _executor->now() + timeout, [this](const executor::TaskExecutor::CallbackArgs& cbData) {
+            _executor->now() + timeout,
+            [this, self](const executor::TaskExecutor::CallbackArgs& cbData) {
                 auto lk = stdx::lock_guard(_mutex);
-
                 if (!cbData.status.isOK()) {
+                    if (!get().isReady()) {
+                        _outputFuturePromise->setError(cbData.status);
+                    }
                     return;
                 }
 
@@ -74,33 +100,22 @@ public:
 
         _timeoutCbHandle = swCbHandle.getValue();
 
-        _inputFuture = std::move(inputFuture).onCompletion([this](StatusWith<ResultStatus> status) {
-            {
+        _inputFuture =
+            std::move(inputFuture).onCompletion([this, self](StatusWith<ResultStatus> status) {
                 auto lk = stdx::lock_guard(_mutex);
                 _executor->cancel(_timeoutCbHandle.get());
                 if (!get().isReady()) {
                     _outputFuturePromise->setFrom(status);
                 }
                 return status;
-            }
-        });
-    }
-
-    ~DeadlineFuture() {
-        {
-            auto lk = stdx::lock_guard(_mutex);
-            _executor->cancel(_timeoutCbHandle.get());
-        }
-    }
-
-    SharedSemiFuture<ResultStatus> get() const {
-        return _outputFuturePromise->getFuture();
+            });
     }
 
 private:
+    const std::shared_ptr<executor::TaskExecutor> _executor;
+
     mutable Mutex _mutex =
         MONGO_MAKE_LATCH(HierarchicalAcquisitionLevel(4), "DeadlineFuture::_mutex");
-    const std::shared_ptr<executor::TaskExecutor> _executor;
     Future<ResultStatus> _inputFuture;
     boost::optional<executor::TaskExecutor::CallbackHandle> _timeoutCbHandle;
     std::unique_ptr<SharedPromise<ResultStatus>> _outputFuturePromise;
diff --git a/src/mongo/db/process_health/health_observer_base.cpp b/src/mongo/db/process_health/health_observer_base.cpp
index 3e40bcf7484..8328dcd2311 100644
--- a/src/mongo/db/process_health/health_observer_base.cpp
+++ b/src/mongo/db/process_health/health_observer_base.cpp
@@ -57,7 +57,7 @@ SharedSemiFuture<HealthCheckStatus> HealthObserverBase::periodicCheck(
         _currentlyRunningHealthCheck = true;
     }
 
-    _deadlineFuture = std::make_unique<DeadlineFuture<HealthCheckStatus>>(
+    _deadlineFuture = DeadlineFuture<HealthCheckStatus>::create(
         taskExecutor,
         periodicCheckImpl({cancellationToken, taskExecutor})
             .onCompletion([this](StatusWith<HealthCheckStatus> status) {
diff --git a/src/mongo/db/process_health/health_observer_base.h b/src/mongo/db/process_health/health_observer_base.h
index fc47d36e806..3ed9d29c876 100644
--- a/src/mongo/db/process_health/health_observer_base.h
+++ b/src/mongo/db/process_health/health_observer_base.h
@@ -104,7 +104,7 @@ protected:
 
     // Indicates if there any check running to prevent running checks concurrently.
     bool _currentlyRunningHealthCheck = false;
-    std::unique_ptr<DeadlineFuture<HealthCheckStatus>> _deadlineFuture;
+    std::shared_ptr<const DeadlineFuture<HealthCheckStatus>> _deadlineFuture;
     // Enforces the safety interval.
     Date_t _lastTimeTheCheckWasRun;
     Date_t _lastTimeCheckCompleted;