SERVER-38521 Use `advanceTime` in network operation cancellation tests to ensure the cancel callbacks run

4 files changed, 16 insertions, 12 deletions

diff --git a/src/mongo/db/transaction_coordinator_futures_util_test.cpp b/src/mongo/db/transaction_coordinator_futures_util_test.cpp
index 7bb67bc89b7..f5ebb46cbf6 100644
--- a/src/mongo/db/transaction_coordinator_futures_util_test.cpp
+++ b/src/mongo/db/transaction_coordinator_futures_util_test.cpp
@@ -374,13 +374,13 @@ TEST_F(AsyncWorkSchedulerTest, ScheduledBlockingWorkInSucceeds) {
 
     {
         executor::NetworkInterfaceMock::InNetworkGuard guard(network());
-        network()->runUntil(network()->now() + Milliseconds{5});
+        network()->advanceTime(network()->now() + Milliseconds{5});
         ASSERT(!future.isReady());
     }
 
     {
         executor::NetworkInterfaceMock::InNetworkGuard guard(network());
-        network()->runUntil(network()->now() + Milliseconds{5});
+        network()->advanceTime(network()->now() + Milliseconds{5});
         ASSERT(future.isReady());
     }
 
@@ -501,13 +501,19 @@ TEST_F(AsyncWorkSchedulerTest, ShutdownInterruptsRemoteCommandsWhichAreBlockedWa
         kShardIds[1], ReadPreferenceSetting{ReadPreference::PrimaryOnly}, BSON("TestCommand" << 1));
 
     auto future2 = async.scheduleRemoteCommand(
-        kShardIds[2], ReadPreferenceSetting{ReadPreference::PrimaryOnly}, BSON("TestCommand" << 1));
+        kShardIds[2], ReadPreferenceSetting{ReadPreference::PrimaryOnly}, BSON("TestCommand" << 2));
 
     ASSERT(!future1.isReady());
     ASSERT(!future2.isReady());
 
     async.shutdown({ErrorCodes::InternalError, "Test internal error"});
 
+    // Ensure that any scheduled cancellations as a result of the shutdown call above run
+    {
+        executor::NetworkInterfaceMock::InNetworkGuard guard(network());
+        network()->advanceTime(network()->now() + Milliseconds(1));
+    }
+
     ASSERT_THROWS_CODE(future1.get(), AssertionException, ErrorCodes::InternalError);
     ASSERT_THROWS_CODE(future2.get(), AssertionException, ErrorCodes::InternalError);
 }
@@ -643,7 +649,7 @@ TEST_F(DoWhileTest, LoopObeysBackoff) {
     // Back-off is 1 millisecond now
     {
         executor::NetworkInterfaceMock::InNetworkGuard guard(network());
-        network()->runUntil(network()->now() + Milliseconds{1});
+        network()->advanceTime(network()->now() + Milliseconds{1});
         ASSERT(!future.isReady());
         ASSERT_EQ(2, numLoops);
     }
@@ -652,14 +658,14 @@ TEST_F(DoWhileTest, LoopObeysBackoff) {
     // loop body to run
     {
         executor::NetworkInterfaceMock::InNetworkGuard guard(network());
-        network()->runUntil(network()->now() + Milliseconds{1});
+        network()->advanceTime(network()->now() + Milliseconds{1});
         ASSERT(!future.isReady());
         ASSERT_EQ(2, numLoops);
     }
 
     {
         executor::NetworkInterfaceMock::InNetworkGuard guard(network());
-        network()->runUntil(network()->now() + Seconds{1});
+        network()->advanceTime(network()->now() + Seconds{1});
         ASSERT(future.isReady());
         ASSERT_EQ(3, numLoops);
     }
diff --git a/src/mongo/executor/network_interface_mock.cpp b/src/mongo/executor/network_interface_mock.cpp
index 641f1095530..9924104ae94 100644
--- a/src/mongo/executor/network_interface_mock.cpp
+++ b/src/mongo/executor/network_interface_mock.cpp
@@ -154,7 +154,7 @@ void NetworkInterfaceMock::cancelCommand(const CallbackHandle& cbHandle, const B
 
 void NetworkInterfaceMock::_interruptWithResponse_inlock(
     const CallbackHandle& cbHandle,
-    const std::vector<NetworkOperationList*> queuesToCheck,
+    const std::vector<NetworkOperationList*>& queuesToCheck,
     const ResponseStatus& response) {
 
     auto matchFn = [&cbHandle](const auto& ops) { return ops.isForCallback(cbHandle); };
diff --git a/src/mongo/executor/network_interface_mock.h b/src/mongo/executor/network_interface_mock.h
index c0f1693a660..0cf12cd900a 100644
--- a/src/mongo/executor/network_interface_mock.h
+++ b/src/mongo/executor/network_interface_mock.h
@@ -281,7 +281,7 @@ public:
      * CallbackCanceled error.
      */
     void _interruptWithResponse_inlock(const TaskExecutor::CallbackHandle& cbHandle,
-                                       const std::vector<NetworkOperationList*> queuesToCheck,
+                                       const std::vector<NetworkOperationList*>& queuesToCheck,
                                        const TaskExecutor::ResponseStatus& response);
 
 private:
diff --git a/src/mongo/executor/network_test_env.cpp b/src/mongo/executor/network_test_env.cpp
index 0a6aa800fed..22c450fe462 100644
--- a/src/mongo/executor/network_test_env.cpp
+++ b/src/mongo/executor/network_test_env.cpp
@@ -44,7 +44,7 @@ NetworkTestEnv::NetworkTestEnv(TaskExecutor* executor, NetworkInterfaceMock* net
     : _executor(executor), _mockNetwork(network) {}
 
 void NetworkTestEnv::onCommand(OnCommandFunction func) {
-    _mockNetwork->enterNetwork();
+    executor::NetworkInterfaceMock::InNetworkGuard guard(_mockNetwork);
 
     const NetworkInterfaceMock::NetworkOperationIterator noi = _mockNetwork->getNextReadyRequest();
     const RemoteCommandRequest& request = noi->getRequest();
@@ -63,11 +63,10 @@ void NetworkTestEnv::onCommand(OnCommandFunction func) {
     }
 
     _mockNetwork->runReadyNetworkOperations();
-    _mockNetwork->exitNetwork();
 }
 
 void NetworkTestEnv::onCommandWithMetadata(OnCommandWithMetadataFunction func) {
-    _mockNetwork->enterNetwork();
+    executor::NetworkInterfaceMock::InNetworkGuard guard(_mockNetwork);
 
     const NetworkInterfaceMock::NetworkOperationIterator noi = _mockNetwork->getNextReadyRequest();
     const RemoteCommandRequest& request = noi->getRequest();
@@ -85,7 +84,6 @@ void NetworkTestEnv::onCommandWithMetadata(OnCommandWithMetadataFunction func) {
     }
 
     _mockNetwork->runReadyNetworkOperations();
-    _mockNetwork->exitNetwork();
 }
 
 void NetworkTestEnv::onFindCommand(OnFindCommandFunction func) {