RUST-2412 Send afterClusterTime on writes in causally-consistent sessions

driver/src/client/executor.rs

@@ -839,7 +839,7 @@ impl Client {
                         session.transaction.state,
                         TransactionState::None | TransactionState::Starting
                     )
-                    && op.read_concern().supported()
+                    && (op.read_concern().supported() || op.is_after_cluster_time_write())
                 {
                     cmd.set_after_cluster_time(session);
                 }

driver/src/operation.rs

@@ -193,6 +193,9 @@ pub(crate) trait Operation {
     /// The noun to this operation's verb.
     fn target(&self) -> OperationTarget;
 
+    /// Whether this is a write operation that needs `afterClusterTime` set.
+    fn is_after_cluster_time_write(&self) -> bool;
+
     #[cfg(feature = "opentelemetry")]
     type Otel: crate::otel::OtelWitness<Op = Self>;
 
@@ -407,70 +410,197 @@ pub(crate) trait OperationWithDefaults: Send + Sync {
 
     fn target(&self) -> OperationTarget;
 
+    fn is_after_cluster_time_write(&self) -> bool {
+        self.write_concern().supported()
+    }
+
+    #[cfg(feature = "opentelemetry")]
+    type Otel: crate::otel::OtelWitness<Op = Self>;
+}
+
+pub(crate) trait OperationDispatch<Kind> {
+    type O;
+    const NAME: &'static CStr;
+    const ZERO_COPY: bool;
+    fn build(&mut self, description: &StreamDescription) -> Result<Command>;
+    fn extract_at_cluster_time(&self, response: &RawDocument) -> Result<Option<Timestamp>>;
+    fn handle_response<'a>(
+        &'a self,
+        response: Cow<'a, RawCommandResponse>,
+        context: ExecutionContext<'a>,
+    ) -> BoxFuture<'a, Result<Self::O>>;
+    fn handle_error(&self, error: Error) -> Result<Self::O>;
+    fn selection_criteria(&self) -> Feature<&SelectionCriteria>;
+    fn read_concern(&self) -> Feature<&ReadConcern>;
+    fn write_concern(&self) -> Feature<&WriteConcern>;
+    fn supports_sessions(&self) -> bool;
+    fn retryability(&self, options: &ClientOptions) -> Retryability;
+    fn is_backpressure_retryable(&self, options: &ClientOptions) -> bool;
+    fn update_for_retry(&mut self, retry: Option<&Retry>);
+    fn override_criteria(&self) -> OverrideCriteriaFn;
+    fn pinned_connection(&self) -> Option<&PinnedConnectionHandle>;
+    fn name(&self) -> &CStr;
+    fn target(&self) -> OperationTarget;
+    fn is_after_cluster_time_write(&self) -> bool;
     #[cfg(feature = "opentelemetry")]
     type Otel: crate::otel::OtelWitness<Op = Self>;
 }
 
-impl<T: OperationWithDefaults> Operation for T
+macro_rules! operation_dispatch {
+    ($trait:path, $tag:ty, $handle_response:ident) => {
+        impl<T: $trait> OperationDispatch<$tag> for T {
+            operation_dispatch_body! { $trait, $handle_response }
+        }
+    };
+}
+
+macro_rules! operation_dispatch_body {
+    ($trait:path, $handle_response:ident) => {
+        type O = <T as $trait>::O;
+        const NAME: &'static CStr = <T as $trait>::NAME;
+        const ZERO_COPY: bool = <T as $trait>::ZERO_COPY;
+        fn build(&mut self, description: &StreamDescription) -> Result<Command> {
+            <T as $trait>::build(self, description)
+        }
+        fn extract_at_cluster_time(&self, response: &RawDocument) -> Result<Option<Timestamp>> {
+            <T as $trait>::extract_at_cluster_time(self, response)
+        }
+        fn handle_response<'a>(
+            &'a self,
+            response: Cow<'a, RawCommandResponse>,
+            context: ExecutionContext<'a>,
+        ) -> BoxFuture<'a, Result<Self::O>> {
+            <T as $trait>::$handle_response(self, response, context)
+        }
+        fn handle_error(&self, error: Error) -> Result<Self::O> {
+            <T as $trait>::handle_error(self, error)
+        }
+        fn selection_criteria(&self) -> Feature<&SelectionCriteria> {
+            <T as $trait>::selection_criteria(self)
+        }
+        fn read_concern(&self) -> Feature<&ReadConcern> {
+            <T as $trait>::read_concern(self)
+        }
+        fn write_concern(&self) -> Feature<&WriteConcern> {
+            <T as $trait>::write_concern(self)
+        }
+        fn supports_sessions(&self) -> bool {
+            <T as $trait>::supports_sessions(self)
+        }
+        fn retryability(&self, options: &ClientOptions) -> Retryability {
+            <T as $trait>::retryability(self, options)
+        }
+        fn is_backpressure_retryable(&self, options: &ClientOptions) -> bool {
+            <T as $trait>::is_backpressure_retryable(self, options)
+        }
+        fn update_for_retry(&mut self, retry: Option<&Retry>) {
+            <T as $trait>::update_for_retry(self, retry)
+        }
+        fn override_criteria(&self) -> OverrideCriteriaFn {
+            <T as $trait>::override_criteria(self)
+        }
+        fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
+            <T as $trait>::pinned_connection(self)
+        }
+        fn name(&self) -> &CStr {
+            <T as $trait>::name(self)
+        }
+        fn target(&self) -> OperationTarget {
+            <T as $trait>::target(self)
+        }
+        fn is_after_cluster_time_write(&self) -> bool {
+            <T as $trait>::is_after_cluster_time_write(self)
+        }
+        #[cfg(feature = "opentelemetry")]
+        type Otel = <T as $trait>::Otel;
+    };
+}
+
+pub(crate) trait OperationImpl {
+    type Kind;
+}
+
+impl<K, T> Operation for T
 where
-    T: Send + Sync,
+    T: OperationImpl<Kind = K> + OperationDispatch<K> + Send + Sync,
 {
-    type O = T::O;
-    const NAME: &'static CStr = T::NAME;
-    const ZERO_COPY: bool = T::ZERO_COPY;
-    fn build(&mut self, description: &StreamDescription) -> Result<Command> {
-        self.build(description)
-    }
-    fn extract_at_cluster_time(&self, response: &RawDocument) -> Result<Option<Timestamp>> {
-        self.extract_at_cluster_time(response)
-    }
+    operation_dispatch_body! { OperationDispatch<K>, handle_response }
+}
+
+pub(crate) struct WithDefaults;
+
+operation_dispatch! { OperationWithDefaults, WithDefaults, handle_response_async }
+
+pub(crate) trait OperationWrapper {
+    // Required
+    type Wrapped: Operation;
+    type O;
+    #[cfg(feature = "opentelemetry")]
+    type Otel: crate::otel::OtelWitness<Op = Self>;
+
+    fn wrapped(&self) -> &Self::Wrapped;
+    fn wrapped_mut(&mut self) -> &mut Self::Wrapped;
+
     fn handle_response<'a>(
         &'a self,
         response: Cow<'a, RawCommandResponse>,
         context: ExecutionContext<'a>,
-    ) -> BoxFuture<'a, Result<Self::O>> {
-        self.handle_response_async(response, context)
+    ) -> BoxFuture<'a, Result<Self::O>>;
+
+    // Delegated to wrapped
+    const NAME: &'static CStr = Self::Wrapped::NAME;
+    const ZERO_COPY: bool = Self::Wrapped::ZERO_COPY;
+    fn build(&mut self, description: &StreamDescription) -> Result<Command> {
+        self.wrapped_mut().build(description)
     }
     fn handle_error(&self, error: Error) -> Result<Self::O> {
-        self.handle_error(error)
+        Err(error)
+    }
+    fn extract_at_cluster_time(&self, response: &RawDocument) -> Result<Option<Timestamp>> {
+        self.wrapped().extract_at_cluster_time(response)
     }
     fn selection_criteria(&self) -> Feature<&SelectionCriteria> {
-        self.selection_criteria()
+        self.wrapped().selection_criteria()
     }
     fn read_concern(&self) -> Feature<&ReadConcern> {
-        self.read_concern()
+        self.wrapped().read_concern()
     }
     fn write_concern(&self) -> Feature<&WriteConcern> {
-        self.write_concern()
+        self.wrapped().write_concern()
     }
     fn supports_sessions(&self) -> bool {
-        self.supports_sessions()
+        self.wrapped().supports_sessions()
     }
     fn retryability(&self, options: &ClientOptions) -> Retryability {
-        self.retryability(options)
+        self.wrapped().retryability(options)
     }
     fn is_backpressure_retryable(&self, options: &ClientOptions) -> bool {
-        self.is_backpressure_retryable(options)
+        self.wrapped().is_backpressure_retryable(options)
     }
     fn update_for_retry(&mut self, retry: Option<&Retry>) {
-        self.update_for_retry(retry)
+        self.wrapped_mut().update_for_retry(retry);
     }
     fn override_criteria(&self) -> OverrideCriteriaFn {
-        self.override_criteria()
+        self.wrapped().override_criteria()
     }
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
-        self.pinned_connection()
+        self.wrapped().pinned_connection()
     }
     fn name(&self) -> &CStr {
-        self.name()
+        self.wrapped().name()
     }
     fn target(&self) -> OperationTarget {
-        self.target()
+        self.wrapped().target()
+    }
+    fn is_after_cluster_time_write(&self) -> bool {
+        self.wrapped().is_after_cluster_time_write()
     }
-    #[cfg(feature = "opentelemetry")]
-    type Otel = <Self as OperationWithDefaults>::Otel;
 }
 
+pub(crate) struct Wrapper;
+
+operation_dispatch! { OperationWrapper, Wrapper, handle_response }
+
 fn should_redact_body(body: &RawDocumentBuf) -> bool {
     if let Some(Ok((command_name, _))) = body.into_iter().next() {
         HELLO_COMMAND_NAMES.contains(command_name.to_lowercase().as_str())

driver/src/operation/abort_transaction.rs

@@ -4,7 +4,7 @@ use crate::{
     client::{session::TransactionPin, Retry},
     cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription},
     error::Result,
-    operation::Retryability,
+    operation::{OperationImpl, Retryability, WithDefaults},
     options::{ClientOptions, SelectionCriteria, WriteConcern},
     Client,
 };
@@ -85,9 +85,17 @@ impl OperationWithDefaults for AbortTransaction {
         crate::operation::OperationTarget::admin(&self.target)
     }
 
+    fn is_after_cluster_time_write(&self) -> bool {
+        false
+    }
+
     #[cfg(feature = "opentelemetry")]
     type Otel = crate::otel::Witness<Self>;
 }
 
+impl OperationImpl for AbortTransaction {
+    type Kind = WithDefaults;
+}
+
 #[cfg(feature = "opentelemetry")]
 impl crate::otel::OtelInfoDefaults for AbortTransaction {}

driver/src/operation/aggregate.rs

@@ -7,7 +7,7 @@ use crate::{
     cmap::{Command, RawCommandResponse, StreamDescription},
     cursor::common::CursorSpecification,
     error::Result,
-    operation::{append_options, OperationTarget, Retryability},
+    operation::{append_options, OperationImpl, OperationTarget, Retryability, WithDefaults},
     options::{AggregateOptions, ClientOptions, ReadPreference, SelectionCriteria, WriteConcern},
 };
 
@@ -173,10 +173,18 @@ impl OperationWithDefaults for Aggregate {
         self.target.clone()
     }
 
+    fn is_after_cluster_time_write(&self) -> bool {
+        false
+    }
+
     #[cfg(feature = "opentelemetry")]
     type Otel = crate::otel::Witness<Self>;
 }
 
+impl OperationImpl for Aggregate {
+    type Kind = WithDefaults;
+}
+
 #[cfg(feature = "opentelemetry")]
 impl crate::otel::OtelInfoDefaults for Aggregate {
     fn output_cursor_id(output: &Self::O) -> Option<i64> {