SERVER-39237 Ensure with replaceCollection can clean up temp collection after interrupt

12 files changed, 112 insertions, 12 deletions

diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication.yml
index 8653ff344ea..bb0ad70e109 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication.yml
@@ -102,6 +102,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn.yml
index 124bedf1480..c5429ac8c97 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn.yml
@@ -105,6 +105,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn_with_balancer.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn_with_balancer.yml
index 34d36b6a456..096f6854118 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn_with_balancer.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_local_read_write_multi_stmt_txn_with_balancer.yml
@@ -105,6 +105,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn.yml
index 9532f5931cd..09c4fea9119 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn.yml
@@ -105,6 +105,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn_with_balancer.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn_with_balancer.yml
index 91dc46a83eb..dedebb659a7 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn_with_balancer.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_multi_stmt_txn_with_balancer.yml
@@ -105,6 +105,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_with_balancer.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_with_balancer.yml
index 039bd6cbba8..c34499279f1 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_with_balancer.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_replication_with_balancer.yml
@@ -105,6 +105,9 @@ selector:
   # TODO Unblacklist (SERVER-35538).
   - jstests/concurrency/fsm_workloads/database_versioning.js
 
+  # TODO Unblacklist (SERVER-38852).
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
+
   # serverStatus does not include transaction metrics on mongos.
   - jstests/concurrency/fsm_workloads/multi_statement_transaction_atomicity_isolation_metrics_test.js
 
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns.yml
index 0e4acd4efa7..9516c1b51e1 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns.yml
@@ -149,6 +149,7 @@ selector:
   - jstests/concurrency/fsm_workloads/agg_sort.js
   - jstests/concurrency/fsm_workloads/collmod.js
   - jstests/concurrency/fsm_workloads/collmod_separate_collections.js
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
   - jstests/concurrency/fsm_workloads/invalidated_cursors.js
   - jstests/concurrency/fsm_workloads/kill_multicollection_aggregation.js
   - jstests/concurrency/fsm_workloads/view_catalog.js
diff --git a/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns_and_balancer.yml b/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns_and_balancer.yml
index 8560af6de3c..b86c6e59dda 100644
--- a/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns_and_balancer.yml
+++ b/buildscripts/resmokeconfig/suites/concurrency_sharded_with_stepdowns_and_balancer.yml
@@ -151,6 +151,7 @@ selector:
   - jstests/concurrency/fsm_workloads/agg_sort.js
   - jstests/concurrency/fsm_workloads/collmod.js
   - jstests/concurrency/fsm_workloads/collmod_separate_collections.js
+  - jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
   - jstests/concurrency/fsm_workloads/invalidated_cursors.js
   - jstests/concurrency/fsm_workloads/kill_multicollection_aggregation.js
   - jstests/concurrency/fsm_workloads/view_catalog.js
diff --git a/jstests/concurrency/fsm_workloads/agg_base.js b/jstests/concurrency/fsm_workloads/agg_base.js
index 410fa2295ef..ec598ef00c8 100644
--- a/jstests/concurrency/fsm_workloads/agg_base.js
+++ b/jstests/concurrency/fsm_workloads/agg_base.js
@@ -68,6 +68,10 @@ var $config = (function() {
         assertWhenOwnColl.eq(this.numDocs / 2, db[collName].find({flag: true}).itcount());
     }
 
+    function teardown(db, collName, cluster) {
+        // By default, do nothing on teardown. Workload tests may implement this function.
+    }
+
     return {
         // Using few threads and iterations because each iteration is fairly expensive compared to
         // other workloads' iterations. (Each does a collection scan over a few thousand documents
@@ -79,5 +83,6 @@ var $config = (function() {
         transitions: transitions,
         data: data,
         setup: setup,
+        teardown: teardown,
     };
 })();
diff --git a/jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js b/jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
new file mode 100644
index 00000000000..d36769776fd
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/agg_out_interrupt_cleanup.js
@@ -0,0 +1,59 @@
+/**
+ * Tests $out stage of aggregate command concurrently with killOp. Ensures that all the temporary
+ * collections created during aggreate command are deleted.
+ *
+ * @tags: [uses_curop_agg_stage]
+ */
+'use strict';
+load('jstests/concurrency/fsm_libs/extend_workload.js');  // for extendWorkload
+load('jstests/concurrency/fsm_workloads/agg_base.js');    // for $config
+
+var $config = extendWorkload($config, function($config, $super) {
+
+    $config.states.aggregate = function aggregate(db, collName) {
+        // $out to the same collection so that concurrent aggregate commands would cause congestion.
+        db[collName].runCommand(
+            {aggregate: collName, pipeline: [{$out: "interrupt_temp_out"}], cursor: {}});
+    };
+
+    $config.states.killOp = function killOp(db, collName) {
+        // The aggregate command could be running different commands internally (renameCollection,
+        // insertDocument, etc.) depending on which stage of execution it is in. So, get all the
+        // operations that are running against the input, output or temp collections.
+        const activeCurOpsFilter = {
+            op: "command",
+            active: true,
+            $or: [
+                {"ns": db.getName() + ".interrupt_temp_out"},              // For output collection.
+                {"ns": db.getName() + "." + collName},                     // For input collection.
+                {"ns": {$regex: "^" + db.getName() + "\.tmp\.agg_out.*"}}  // For temp during $out.
+            ],
+            "command.drop": {
+                $exists: false
+            }  // Exclude 'drop' command from the filter to make sure that we don't kill the the
+               // drop command which is responsible for dropping the temporary collection.
+        };
+
+        const currentOpOutput = db.getSiblingDB('admin')
+                                    .aggregate([{$currentOp: {}}, {$match: activeCurOpsFilter}])
+                                    .toArray();
+        for (let op of currentOpOutput) {
+            assert(op.hasOwnProperty('opid'));
+            assertAlways.commandWorked(db.getSiblingDB('admin').killOp(op.opid));
+        }
+    };
+
+    $config.teardown = function teardown(db, collName, cluster) {
+        // Ensure that no temporary collection is left behind.
+        assertAlways.eq(db.getCollectionNames().filter(col => col.includes('tmp.agg_out')).length,
+                        0);
+    };
+
+    $config.transitions = {
+        query: {aggregate: 1.0},
+        aggregate: {aggregate: 0.8, killOp: 0.2},
+        killOp: {aggregate: 0.8, killOp: 0.2}
+    };
+
+    return $config;
+});
diff --git a/src/mongo/db/pipeline/document_source_out_replace_coll.cpp b/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
index b53970f5005..f98275fb03b 100644
--- a/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
+++ b/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
@@ -28,14 +28,40 @@
  */
 
 #include "mongo/platform/basic.h"
-#include "mongo/rpc/get_status_from_command_result.h"
 
 #include "mongo/db/pipeline/document_source_out_replace_coll.h"
 
+#include "mongo/rpc/get_status_from_command_result.h"
+
 namespace mongo {
 
 static AtomicWord<unsigned> aggOutCounter;
 
+DocumentSourceOutReplaceColl::~DocumentSourceOutReplaceColl() {
+    DESTRUCTOR_GUARD(
+        // Make sure we drop the temp collection if anything goes wrong. Errors are ignored
+        // here because nothing can be done about them. Additionally, if this fails and the
+        // collection is left behind, it will be cleaned up next time the server is started.
+        if (_tempNs.size()) {
+            auto cleanupClient =
+                pExpCtx->opCtx->getServiceContext()->makeClient("$out_replace_coll_cleanup");
+            AlternativeClientRegion acr(cleanupClient);
+            // Create a new operation context so that any interrputs on the current operation will
+            // not affect the dropCollection operation below.
+            auto cleanupOpCtx = cc().makeOperationContext();
+
+            LocalReadConcernBlock readLocal(cleanupOpCtx.get());
+
+            pExpCtx->mongoProcessInterface->setOperationContext(cleanupOpCtx.get());
+
+            // Reset the operation context back to original once dropCollection is done.
+            ON_BLOCK_EXIT(
+                [this] { pExpCtx->mongoProcessInterface->setOperationContext(pExpCtx->opCtx); });
+
+            pExpCtx->mongoProcessInterface->directClient()->dropCollection(_tempNs.ns());
+        });
+}
+
 void DocumentSourceOutReplaceColl::initializeWriteNs() {
     LocalReadConcernBlock readLocal(pExpCtx->opCtx);
 
diff --git a/src/mongo/db/pipeline/document_source_out_replace_coll.h b/src/mongo/db/pipeline/document_source_out_replace_coll.h
index ec5c14f2a47..db0dde23ec0 100644
--- a/src/mongo/db/pipeline/document_source_out_replace_coll.h
+++ b/src/mongo/db/pipeline/document_source_out_replace_coll.h
@@ -42,17 +42,7 @@ class DocumentSourceOutReplaceColl final : public DocumentSourceOut {
 public:
     using DocumentSourceOut::DocumentSourceOut;
 
-    ~DocumentSourceOutReplaceColl() {
-        DESTRUCTOR_GUARD(
-            // Make sure we drop the temp collection if anything goes wrong. Errors are ignored
-            // here because nothing can be done about them. Additionally, if this fails and the
-            // collection is left behind, it will be cleaned up next time the server is started.
-            if (_tempNs.size()) {
-                LocalReadConcernBlock readLocal(pExpCtx->opCtx);
-
-                pExpCtx->mongoProcessInterface->directClient()->dropCollection(_tempNs.ns());
-            });
-    }
+    ~DocumentSourceOutReplaceColl();
 
     /**
      * Sets up a temp collection which contains the same indexes and options as the output