1 files changed, 49 insertions, 0 deletions
diff --git a/src/mongo/transport/grpc/mock_util.h b/src/mongo/transport/grpc/mock_util.h
index 456437cfe1f..2e0fc9c02ed 100644
--- a/src/mongo/transport/grpc/mock_util.h
+++ b/src/mongo/transport/grpc/mock_util.h
@@ -32,14 +32,63 @@
 #include <map>
 #include <string>
 
+#include <boost/optional.hpp>
+#include <grpcpp/grpcpp.h>
+#include <grpcpp/support/status.h>
+
 #include "mongo/db/operation_context.h"
 #include "mongo/db/service_context.h"
 #include "mongo/util/interruptible.h"
+#include "mongo/util/synchronized_value.h"
 #include "mongo/util/time_support.h"
 
 namespace mongo::transport::grpc {
 
 /**
+ * Class containing the shared cancellation state between a MockServerStream and its corresponding
+ * MockClientStream. This mocks cases in which an RPC is terminated before the server's RPC handler
+ * is able to return a status (e.g. explicit client/server cancellation or a network error).
+ */
+class MockCancellationState {
+public:
+    explicit MockCancellationState(Date_t deadline) : _deadline(deadline) {}
+
+    Date_t getDeadline() const {
+        return _deadline;
+    }
+
+    bool isCancelled() const {
+        return _cancellationStatus->has_value() || isDeadlineExceeded();
+    }
+
+    bool isDeadlineExceeded() const {
+        return getGlobalServiceContext()->getFastClockSource()->now() > _deadline;
+    }
+
+    boost::optional<::grpc::Status> getCancellationStatus() {
+        if (auto status = _cancellationStatus.synchronize(); status->has_value()) {
+            return *status;
+        } else if (isDeadlineExceeded()) {
+            return ::grpc::Status(::grpc::StatusCode::DEADLINE_EXCEEDED, "Deadline exceeded");
+        } else {
+            return boost::none;
+        }
+    }
+
+    void cancel(::grpc::Status status) {
+        auto statusGuard = _cancellationStatus.synchronize();
+        if (statusGuard->has_value()) {
+            return;
+        }
+        *statusGuard = std::move(status);
+    }
+
+private:
+    Date_t _deadline;
+    synchronized_value<boost::optional<::grpc::Status>> _cancellationStatus;
+};
+
+/**
  * Performs the provided lambda and returns its return value. If the lambda's execution is
  * interrupted due to the deadline being exceeded, this returns a default-constructed T instead.
  */