SERVER-30014 OperationContextGroup tracks, interrupts OperationContext objects

4 files changed, 368 insertions, 0 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index d3728542f69..eba59679fc9 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -445,6 +445,7 @@ env.Library(
         'operation_context.cpp',
         'service_context.cpp',
         'service_context_noop.cpp',
+        'operation_context_group.cpp'
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/util/clock_sources',
diff --git a/src/mongo/db/operation_context_group.cpp b/src/mongo/db/operation_context_group.cpp
new file mode 100644
index 00000000000..30305044b23
--- /dev/null
+++ b/src/mongo/db/operation_context_group.cpp
@@ -0,0 +1,120 @@
+/**
+ *    Copyright (C) 2017 Mongodb Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects
+ *    for all of the code used other than as permitted herein. If you modify
+ *    file(s) with this exception, you may extend this exception to your
+ *    version of the file(s), but you are not obligated to do so. If you do not
+ *    wish to do so, delete this exception statement from your version. If you
+ *    delete this exception statement from all source files in the program,
+ *    then also delete it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/operation_context_group.h"
+
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+
+namespace mongo {
+
+namespace {
+
+using ContextTable = std::vector<ServiceContext::UniqueOperationContext>;
+
+auto find(ContextTable& contexts, OperationContext* cp) {
+    auto it = std::find_if(
+        contexts.begin(), contexts.end(), [cp](auto const& opCtx) { return opCtx.get() == cp; });
+    invariant(it != contexts.end());
+    return it;
+}
+
+void interruptOne(OperationContext* opCtx, ErrorCodes::Error code) {
+    stdx::lock_guard<Client> lk(*opCtx->getClient());
+    opCtx->getServiceContext()->killOperation(opCtx, code);
+}
+
+}  // namespace
+
+// OperationContextGroup::Context
+
+OperationContextGroup::Context::Context(OperationContext& ctx, OperationContextGroup& group)
+    : _opCtx(ctx), _ctxGroup(group) {}
+
+void OperationContextGroup::Context::discard() {
+    if (!_movedFrom) {
+        stdx::lock_guard<stdx::mutex> lk(_ctxGroup._lock);
+        auto it = find(_ctxGroup._contexts, &_opCtx);
+        _ctxGroup._contexts.erase(it);
+        _movedFrom = true;
+    }
+}
+
+// OperationContextGroup
+
+auto OperationContextGroup::makeOperationContext(Client& client) -> Context {
+    return adopt(client.makeOperationContext());
+}
+
+auto OperationContextGroup::adopt(UniqueOperationContext opCtx) -> Context {
+    auto cp = opCtx.get();
+    invariant(cp);
+    stdx::lock_guard<stdx::mutex> lk(_lock);
+    _contexts.emplace_back(std::move(opCtx));
+    if (_interrupted) {
+        interruptOne(cp, _interrupted);
+    }
+    return Context(*cp, *this);
+}
+
+auto OperationContextGroup::take(Context ctx) -> Context {
+    if (ctx._movedFrom || &ctx._ctxGroup == this) {
+        return ctx;
+    }
+    {
+        stdx::lock_guard<stdx::mutex> lk(_lock);
+        auto it = find(ctx._ctxGroup._contexts, &ctx._opCtx);
+        _contexts.emplace_back(std::move(*it));
+        ctx._ctxGroup._contexts.erase(it);
+    }
+    ctx._movedFrom = true;
+    return Context(ctx._opCtx, *this);
+}
+
+void OperationContextGroup::interrupt(ErrorCodes::Error code) {
+    invariant(code);
+    stdx::lock_guard<stdx::mutex> lk(_lock);
+    _interrupted = code;
+    for (auto&& uniqueOperationContext : _contexts) {
+        interruptOne(uniqueOperationContext.get(), code);
+    }
+}
+
+void OperationContextGroup::resetInterrupt() {
+    stdx::lock_guard<stdx::mutex> lk(_lock);
+    _interrupted = ErrorCodes::Error{};
+}
+
+bool OperationContextGroup::isEmpty() {
+    stdx::lock_guard<stdx::mutex> lk(_lock);
+    return _contexts.empty();
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_context_group.h b/src/mongo/db/operation_context_group.h
new file mode 100644
index 00000000000..c2a0eded364
--- /dev/null
+++ b/src/mongo/db/operation_context_group.h
@@ -0,0 +1,172 @@
+/**
+ *    Copyright (C) 2017 Mongodb Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects
+ *    for all of the code used other than as permitted herein. If you modify
+ *    file(s) with this exception, you may extend this exception to your
+ *    version of the file(s), but you are not obligated to do so. If you do not
+ *    wish to do so, delete this exception statement from your version. If you
+ *    delete this exception statement from all source files in the program,
+ *    then also delete it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+#include "mongo/stdx/mutex.h"
+
+namespace mongo {
+
+/**
+ * OperationContextGroup maintains a collection of operation contexts so that they may all be killed
+ * on a common event (typically stepdown).  Its public member functions serialize access to private
+ * data members.
+ */
+class OperationContextGroup {
+    OperationContextGroup(OperationContextGroup const&) = delete;
+    OperationContextGroup(OperationContextGroup&&) = delete;
+    void operator=(OperationContextGroup const&) = delete;
+    void operator=(OperationContextGroup&&) = delete;
+
+public:
+    using UniqueOperationContext = ServiceContext::UniqueOperationContext;
+    class Context;
+
+    OperationContextGroup() = default;
+
+    ~OperationContextGroup() {
+        invariant(isEmpty());
+    }
+
+    /**
+     * Makes an OperationContext on `client` and returns a Context object to track it.  On
+     * destruction of the returned Context, the OperationContext is destroyed and its corresponding
+     * entry in *this is erased.  If *this has been interrupted already, the new context will be
+     * interrupted immediately (taking and releasing the client lock).
+     */
+    Context makeOperationContext(Client& client);
+
+    /**
+     * Takes ownership of the OperationContext from `ctx`, and returns a Context object to track it.
+     * On destruction of the Context, its entry in *this is erased and its corresponding
+     * OperationContext is destroyed. If *this has been interrupted already, `ctx` will be
+     * interrupted immediately (taking and releasing the client lock).
+     */
+    Context adopt(UniqueOperationContext ctx);
+
+    /**
+     * Moves the OperationContext of `ctx` from its current OperationContextGroup into *this.
+     * Do this to protect an OperationContext from being interrupted along with the rest of its
+     * group, or to expose `ctx` to this->interrupt().  Taking from a Context already in *this is
+     * equivalent to moving from `ctx`. Taking a moved-from Context yields another moved-from
+     * Context.  If *this has been interrupted already, `ctx` will be interrupted immediately
+     * (taking and releasing the client lock).
+     */
+    Context take(Context ctx);
+
+    /*
+     * Interrupts all the OperationContexts maintained in *this. Contexts subsequently added to the
+     * group will be interrupted immediately until resetInterrupt() is called.  The lock is taken on
+     * each affected client while interrupting its operation.
+     *
+     * Note: Takes and releases each context's client lock.
+     */
+    void interrupt(ErrorCodes::Error);
+
+    /*
+     * Unsticks the interrupting state of *this.  Subsequently added contexts are not interrupted.
+     */
+    void resetInterrupt();
+
+    /**
+     * Reports whether the group has any OperationContexts.  This must be true before the destructor
+     * is called.  Its usefulness is typically limited to invariants.
+     */
+    bool isEmpty();
+
+private:
+    friend class Context;
+
+    stdx::mutex _lock;
+    std::vector<UniqueOperationContext> _contexts;
+    ErrorCodes::Error _interrupted{};
+
+};  // class OperationContextGroup
+
+/**
+ * Context tracks one OperationContext*, and on destruction unregisters and destroys the associated
+ * OperationContext.  May be used as if it were an OperationContext*.
+ *
+ * The lifetime of an OperationContextGroup::Context object must not exceed that of its
+ * OperationContextGroup, unless it has been moved from, taken from (see
+ * OperationContextGroup::take), or discarded.
+ */
+class OperationContextGroup::Context {
+    Context() = delete;
+    Context(Context const&) = delete;
+    void operator=(Context const&) = delete;
+    void operator=(Context&&) = delete;
+
+public:
+    Context(Context&& ctx) : _opCtx(ctx._opCtx), _ctxGroup(ctx._ctxGroup) {
+        ctx._movedFrom = true;
+    }
+    ~Context() {
+        discard();
+    }
+
+    /**
+     * Returns a pointer to the tracked OperationContext, or nullptr if *this has been moved from.
+     */
+    OperationContext* opCtx() const {
+        return _movedFrom ? nullptr : &_opCtx;
+    }
+
+    /**
+     * These enable treating a Context as if it were an OperationContext*.
+     */
+    operator OperationContext*() const {
+        dassert(!_movedFrom);
+        return &_opCtx;
+    }
+    OperationContext* operator->() const {  // because op-> will not use the conversion
+        dassert(!_movedFrom);
+        return &_opCtx;
+    }
+
+    /**
+     * Destroys and unregisters the corresponding OperationContext.  After this operation, *this is
+     * an xvalue, and can only be destroyed.
+     */
+    void discard();
+
+private:
+    friend class OperationContextGroup;
+
+    Context(OperationContext& ctx, OperationContextGroup& group);
+
+    bool _movedFrom = false;
+    OperationContext& _opCtx;
+    OperationContextGroup& _ctxGroup;
+
+};  // class OperationContextGroup::Context
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_context_test.cpp b/src/mongo/db/operation_context_test.cpp
index 4a6117dbeea..5174e599748 100644
--- a/src/mongo/db/operation_context_test.cpp
+++ b/src/mongo/db/operation_context_test.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/json.h"
 #include "mongo/db/logical_session_id.h"
 #include "mongo/db/operation_context.h"
+#include "mongo/db/operation_context_group.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/service_context_noop.h"
 #include "mongo/stdx/future.h"
@@ -170,6 +171,80 @@ TEST(OperationContextTest, InitializeOperationSessionInfo_SessionIdAndTransactio
     ASSERT_EQ(100, *opCtx->getTxnNumber());
 }
 
+TEST(OperationContextTest, OpCtxGroup) {
+    OperationContextGroup group1;
+    ASSERT_TRUE(group1.isEmpty());
+    {
+        auto serviceCtx1 = stdx::make_unique<ServiceContextNoop>();
+        auto client1 = serviceCtx1->makeClient("OperationContextTest1");
+        auto opCtx1 = group1.makeOperationContext(*client1);
+        ASSERT_FALSE(group1.isEmpty());
+
+        auto serviceCtx2 = stdx::make_unique<ServiceContextNoop>();
+        auto client2 = serviceCtx2->makeClient("OperationContextTest2");
+        {
+            auto opCtx2 = group1.makeOperationContext(*client2);
+            opCtx1.discard();
+            ASSERT_FALSE(group1.isEmpty());
+        }
+        ASSERT_TRUE(group1.isEmpty());
+
+        auto opCtx3 = group1.makeOperationContext(*client1);
+        auto opCtx4 = group1.makeOperationContext(*client2);
+        ASSERT_TRUE(opCtx3->checkForInterruptNoAssert().isOK());    // use member op->
+        ASSERT_TRUE((*opCtx4).checkForInterruptNoAssert().isOK());  // use conversion to OC*
+        group1.interrupt(ErrorCodes::InternalError);
+        ASSERT_FALSE(opCtx3->checkForInterruptNoAssert().isOK());
+        ASSERT_FALSE((*opCtx4).checkForInterruptNoAssert().isOK());
+
+        auto serviceCtx3 = stdx::make_unique<ServiceContextNoop>();
+        auto client3 = serviceCtx3->makeClient("OperationContextTest3");
+        auto opCtx5 = group1.makeOperationContext(*client3);
+        ASSERT_FALSE(opCtx5->checkForInterruptNoAssert().isOK());  // interrupt is sticky
+    }
+    ASSERT_TRUE(group1.isEmpty());
+
+    {
+        group1.resetInterrupt();
+        auto serviceCtx1 = stdx::make_unique<ServiceContextNoop>();
+        auto client1 = serviceCtx1->makeClient("OperationContextTest3");
+        auto opCtx1 = group1.makeOperationContext(*client1);
+        ASSERT_TRUE(opCtx1->checkForInterruptNoAssert().isOK());  // interrupt unstuck
+    }
+
+    OperationContextGroup group2;
+    {
+        auto serviceCtx = stdx::make_unique<ServiceContextNoop>();
+        auto client = serviceCtx->makeClient("OperationContextTest1");
+        auto opCtx2 = group2.adopt(client->makeOperationContext());
+        ASSERT_FALSE(group2.isEmpty());
+        ASSERT_TRUE(opCtx2->checkForInterruptNoAssert().isOK());
+        group2.interrupt(ErrorCodes::InternalError);
+        ASSERT_FALSE(opCtx2->checkForInterruptNoAssert().isOK());
+        opCtx2.discard();
+        ASSERT(opCtx2.opCtx() == nullptr);
+        ASSERT_TRUE(group2.isEmpty());
+    }
+
+    OperationContextGroup group3;
+    OperationContextGroup group4;
+    {
+        auto serviceCtx = stdx::make_unique<ServiceContextNoop>();
+        auto client3 = serviceCtx->makeClient("OperationContextTest3");
+        auto opCtx3 = group3.makeOperationContext(*client3);
+        auto p3 = opCtx3.opCtx();
+        auto opCtx4 = group4.take(std::move(opCtx3));
+        ASSERT_EQ(p3, opCtx4.opCtx());
+        ASSERT(opCtx3.opCtx() == nullptr);
+        ASSERT_TRUE(group3.isEmpty());
+        ASSERT_FALSE(group4.isEmpty());
+        group3.interrupt(ErrorCodes::InternalError);
+        ASSERT_TRUE(opCtx4->checkForInterruptNoAssert().isOK());
+        group4.interrupt(ErrorCodes::InternalError);
+        ASSERT_FALSE(opCtx4->checkForInterruptNoAssert().isOK());
+    }
+}
+
 class OperationDeadlineTests : public unittest::Test {
 public:
     void setUp() {