1 files changed, 92 insertions, 9 deletions

diff --git a/src/mongo/util/interruptible.h b/src/mongo/util/interruptible.h
index 1e9cc070c64..f7a02e971f9 100644
--- a/src/mongo/util/interruptible.h
+++ b/src/mongo/util/interruptible.h
@@ -107,6 +107,28 @@ protected:
      * Returns the equivalent of Date_t::now() + waitFor for the InterruptibleBase's clock
      */
     virtual Date_t getExpirationDateForWaitForValue(Milliseconds waitFor) = 0;
+
+    struct IgnoreInterruptsState {
+        bool ignoreInterrupts;
+        DeadlineState deadline;
+    };
+
+    /**
+     * Pushes an ignore interruption critical section into the InterruptibleBase.
+     * Until an associated popIgnoreInterrupts() is invoked, the InterruptibleBase should ignore
+     * interruptions related to explicit interruption or previously set deadlines.
+     *
+     * Note that new deadlines can be set after this is called, which will again introduce the
+     * possibility of interruption.
+     *
+     * Returns state needed to pop interruption.
+     */
+    virtual IgnoreInterruptsState pushIgnoreInterrupts() = 0;
+
+    /**
+     * Pops the ignored interruption critical section introduced by push.
+     */
+    virtual void popIgnoreInterrupts(IgnoreInterruptsState iis) = 0;
 };
 
 /**
@@ -172,6 +194,44 @@ private:
         return DeadlineGuard(*this, deadline, error);
     }
 
+    /**
+     * An interruption guard provides a region where interruption is ignored.
+     *
+     * Note that this causes the deadline to be reset to Date_t::max(), but that it can also be
+     * subsequently reduced in size after the fact.
+     */
+    class IgnoreInterruptionsGuard {
+    public:
+        IgnoreInterruptionsGuard(const IgnoreInterruptionsGuard&) = delete;
+        IgnoreInterruptionsGuard& operator=(const IgnoreInterruptionsGuard&) = delete;
+
+        IgnoreInterruptionsGuard(IgnoreInterruptionsGuard&& other)
+            : _interruptible(other._interruptible), _oldState(other._oldState) {
+            other._interruptible = nullptr;
+        }
+
+        IgnoreInterruptionsGuard& operator=(IgnoreInterruptionsGuard&&) = delete;
+
+        ~IgnoreInterruptionsGuard() {
+            if (_interruptible) {
+                _interruptible->popIgnoreInterrupts(_oldState);
+            }
+        }
+
+    private:
+        friend Interruptible;
+
+        explicit IgnoreInterruptionsGuard(Interruptible& interruptible)
+            : _interruptible(&interruptible), _oldState(_interruptible->pushIgnoreInterrupts()) {}
+
+        Interruptible* _interruptible;
+        IgnoreInterruptsState _oldState;
+    };
+
+    IgnoreInterruptionsGuard makeIgnoreInterruptionsGuard() {
+        return IgnoreInterruptionsGuard(*this);
+    }
+
 public:
     class WaitListener;
 
@@ -246,6 +306,24 @@ public:
     }
 
     /**
+     * Invokes the passed callback with an interruption guard active.  Additionally handles the
+     * dance of try/catching the invocation and checking checkForInterrupt with the guard inactive
+     * (to allow a higher level timeout to override a lower level one, or for top level interruption
+     * to propagate)
+     */
+    template <typename Callback>
+    decltype(auto) runWithoutInterruptionExceptAtGlobalShutdown(Callback&& cb) {
+        try {
+            const auto guard = makeIgnoreInterruptionsGuard();
+            return std::forward<Callback>(cb)();
+        } catch (const ExceptionForCat<ErrorCategory::ExceededTimeLimitError>&) {
+            // May throw replacement exception
+            checkForInterrupt();
+            throw;
+        }
+    }
+
+    /**
      * Raises a AssertionException if this operation is in a killed state.
      */
     void checkForInterrupt() {
@@ -494,15 +572,7 @@ class Interruptible::NotInterruptible final : public Interruptible {
             return stdx::cv_status::no_timeout;
         }
 
-        try {
-            // If the system clock's time_point's compiler-dependent resolution is higher than
-            // Date_t's milliseconds, it's possible for the conversion from Date_t to time_point
-            // to overflow and trigger an exception. We catch that here to maintain the noexcept
-            // contract.
-            return cv.wait_until(m, deadline.toSystemTimePoint());
-        } catch (const ExceptionFor<ErrorCodes::DurationOverflow>& ex) {
-            return ex.toStatus();
-        }
+        return cv.wait_until(m, deadline.toSystemTimePoint());
     }
 
     Date_t getDeadline() const override {
@@ -513,6 +583,19 @@ class Interruptible::NotInterruptible final : public Interruptible {
         return Status::OK();
     }
 
+    // It's invalid to call the deadline or ignore interruption guards on a possibly noop
+    // Interruptible.
+    //
+    // The noop Interruptible should only be invoked as a default arg at the bottom of the call
+    // stack (with types that won't modify it's invocation)
+    IgnoreInterruptsState pushIgnoreInterrupts() override {
+        MONGO_UNREACHABLE;
+    }
+
+    void popIgnoreInterrupts(IgnoreInterruptsState) override {
+        MONGO_UNREACHABLE;
+    }
+
     DeadlineState pushArtificialDeadline(Date_t deadline, ErrorCodes::Error error) override {
         MONGO_UNREACHABLE;
     }