SERVER-37717 Fix race between killOp and Baton::detach()

4 files changed, 31 insertions, 1 deletions

diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index d5decb651d3..c09b342d1d1 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -366,7 +366,7 @@ void OperationContext::markKilled(ErrorCodes::Error killCode) {
 
     // If we have a baton, we need to wake it up.  The baton itself will check for interruption
     if (_baton) {
-        _baton->schedule([] {});
+        _baton->notify();
     }
 }
 
diff --git a/src/mongo/transport/baton.h b/src/mongo/transport/baton.h
index d26d936c99c..549e84b32a6 100644
--- a/src/mongo/transport/baton.h
+++ b/src/mongo/transport/baton.h
@@ -87,6 +87,14 @@ public:
     virtual void schedule(stdx::function<void()> func) = 0;
 
     /**
+     * Wakes the Baton up if it is currently blocked, or ensures that the next time it tries to
+     * block it is woken immediately.
+     *
+     * Does not actually schedule any work.
+     */
+    virtual void notify() noexcept = 0;
+
+    /**
      * Adds a session, returning a future which activates on read/write-ability of the session.
      */
     enum class Type {
diff --git a/src/mongo/transport/baton_asio_linux.h b/src/mongo/transport/baton_asio_linux.h
index 29edbcb9502..0b41bf5a87b 100644
--- a/src/mongo/transport/baton_asio_linux.h
+++ b/src/mongo/transport/baton_asio_linux.h
@@ -201,6 +201,10 @@ public:
         }
     }
 
+    void notify() noexcept override {
+        _efd.notify();
+    }
+
     bool run(OperationContext* opCtx, boost::optional<Date_t> deadline) override {
         std::vector<SharedPromise<void>> toFulfill;
 
diff --git a/src/mongo/transport/transport_layer_asio_integration_test.cpp b/src/mongo/transport/transport_layer_asio_integration_test.cpp
index 49eac963f76..2eb25acd696 100644
--- a/src/mongo/transport/transport_layer_asio_integration_test.cpp
+++ b/src/mongo/transport/transport_layer_asio_integration_test.cpp
@@ -35,6 +35,7 @@
 #include "mongo/client/async_client.h"
 #include "mongo/client/connection_string.h"
 #include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
 #include "mongo/db/service_context.h"
 #include "mongo/stdx/thread.h"
 #include "mongo/transport/session.h"
@@ -138,6 +139,23 @@ TEST(TransportLayerASIO, ShortReadsAndWritesWork) {
     }
 }
 
+TEST(TransportLayerASIO, KillOpWithBatonDoesntCrash) {
+    auto sc = getGlobalServiceContext();
+    auto client = sc->makeClient(__FILE__);
+    auto opCtx = client->makeOperationContext();
+
+    auto baton = sc->getTransportLayer()->makeBaton(opCtx.get());
+    if (!baton)
+        return;  // This is a test of baton functionality.
+
+    {
+        stdx::lock_guard<Client> lk(*client);
+        opCtx->markKilled();
+    }
+
+    baton->detach();  // Used to go boom. No longer does.
+}
+
 TEST(TransportLayerASIO, asyncConnectTimeoutCleansUpSocket) {
     auto connectionString = unittest::getFixtureConnectionString();
     auto server = connectionString.getServers().front();