1 files changed, 50 insertions, 8 deletions

diff --git a/src/mongo/transport/grpc/mock_server_stream.h b/src/mongo/transport/grpc/mock_server_stream.h
index eb7177f46cc..88fbe51d5eb 100644
--- a/src/mongo/transport/grpc/mock_server_stream.h
+++ b/src/mongo/transport/grpc/mock_server_stream.h
@@ -32,13 +32,14 @@
 #include <map>
 #include <string>
 
+#include <grpcpp/grpcpp.h>
+
 #include "mongo/transport/grpc/bidirectional_pipe.h"
 #include "mongo/transport/grpc/metadata.h"
+#include "mongo/transport/grpc/mock_util.h"
 #include "mongo/transport/grpc/server_stream.h"
-#include "mongo/util/cancellation.h"
 #include "mongo/util/future.h"
 #include "mongo/util/net/hostandport.h"
-#include "mongo/util/time_support.h"
 
 namespace mongo::transport::grpc {
 
@@ -47,16 +48,19 @@ public:
     ~MockServerStream() = default;
 
     boost::optional<SharedBuffer> read() override;
+
     bool write(ConstSharedBuffer msg) override;
 
-    explicit MockServerStream(HostAndPort hostAndPort,
-                              Milliseconds timeout,
+    explicit MockServerStream(HostAndPort remote,
                               Promise<MetadataContainer>&& initialMetadataPromise,
+                              Promise<::grpc::Status>&& rpcTerminationStatusPromise,
+                              std::shared_ptr<MockCancellationState> rpcCancellationState,
                               BidirectionalPipe::End&& serverPipeEnd,
                               MetadataView clientMetadata);
 
 private:
     friend class MockServerContext;
+    friend class MockRPC;
 
     class InitialMetadata {
     public:
@@ -86,13 +90,51 @@ private:
     };
 
     bool isCancelled() const;
-    void close();
 
-    CancellationSource _cancellationSource;
-    Date_t _deadline;
+    /**
+     * Cancel the RPC associated with this stream. This is used for mocking situations in
+     * which an RPC handler was never able to return a final status to the client (e.g. manual
+     * cancellation or a network interruption).
+     *
+     * This method has no effect if the stream is already terminated.
+     */
+    void cancel(::grpc::Status status);
+
+    /**
+     * Closes the stream and sends the final return status of the RPC to the client. This is the
+     * mocked equivalent of an RPC handler returning a status.
+     *
+     * This does not mark the stream as cancelled.
+     *
+     * This method must only be called once, and this stream must not be used after this method has
+     * been called.
+     */
+    void sendReturnStatus(::grpc::Status status);
+
+    HostAndPort _remote;
     InitialMetadata _initialMetadata;
+
+    /**
+     * _rpcReturnStatus is set in sendReturnStatus(), and it is used to mock returning a status from
+     * an RPC handler. sendReturnStatus itself is called via MockRPC::sendReturnStatus().
+     */
+    Promise<::grpc::Status> _rpcReturnStatus;
+
+    /**
+     * _finalStatusReturned is also set in sendReturnStatus(), and it is used to denote that a
+     * status has been returned and the stream should no longer be used.
+     */
+    synchronized_value<bool> _finalStatusReturned;
+
+    /**
+     * _rpcCancellationState is set via cancel(), which is called by either
+     * MockServerContext::tryCancel() or MockRPC::cancel(). It is used to mock situations in which a
+     * server RPC handler is unable to return a status to the client (e.g. explicit cancellation or
+     * a network interruption).
+     */
+    std::shared_ptr<MockCancellationState> _rpcCancellationState;
+
     BidirectionalPipe::End _pipe;
     MetadataView _clientMetadata;
-    HostAndPort _hostAndPort;
 };
 }  // namespace mongo::transport::grpc