SERVER-49106 Add poll function to the ServiceExecutor API

1 files changed, 22 insertions, 0 deletions

diff --git a/src/mongo/transport/mock_session.h b/src/mongo/transport/mock_session.h
index d88a0af263e..72d05aa3e66 100644
--- a/src/mongo/transport/mock_session.h
+++ b/src/mongo/transport/mock_session.h
@@ -29,8 +29,11 @@
 
 #pragma once
 
+#include <list>
+
 #include "mongo/base/checked_cast.h"
 #include "mongo/config.h"
+#include "mongo/stdx/mutex.h"
 #include "mongo/transport/session.h"
 #include "mongo/transport/transport_layer_mock.h"
 #include "mongo/util/net/hostandport.h"
@@ -100,6 +103,22 @@ public:
         return Future<Message>::makeReady(sourceMessage());
     }
 
+    Future<void> waitForData() override {
+        auto fp = makePromiseFuture<void>();
+        stdx::lock_guard<Latch> lk(_waitForDataMutex);
+        _waitForDataQueue.emplace_back(std::move(fp.promise));
+        return std::move(fp.future);
+    }
+
+    void signalAvailableData() {
+        stdx::lock_guard<Latch> lk(_waitForDataMutex);
+        if (_waitForDataQueue.size() == 0)
+            return;
+        Promise<void> promise = std::move(_waitForDataQueue.front());
+        _waitForDataQueue.pop_front();
+        promise.emplaceValue();
+    }
+
     Status sinkMessage(Message message) override {
         if (!_tl || _tl->inShutdown()) {
             return TransportLayer::ShutdownStatus;
@@ -154,6 +173,9 @@ protected:
     HostAndPort _local;
     SockAddr _remoteAddr;
     SockAddr _localAddr;
+
+    mutable Mutex _waitForDataMutex = MONGO_MAKE_LATCH("MockSession::_waitForDataMutex");
+    std::list<Promise<void>> _waitForDataQueue;
 };
 
 }  // namespace transport