RUST-2138 Add optional support for OpenTelemetry (#1495)

Author: abr-egn

.evergreen/run-tests.sh

@@ -6,7 +6,7 @@ set -o pipefail
 source .evergreen/env.sh
 source .evergreen/cargo-test.sh
 
-FEATURE_FLAGS+=("tracing-unstable" "cert-key-password")
+FEATURE_FLAGS+=("tracing-unstable" "cert-key-password" "opentelemetry" "error-backtrace")
 
 if [ "$OPENSSL" = true ]; then
   FEATURE_FLAGS+=("openssl-tls")

Cargo.lock

@@ -78,6 +78,12 @@ tracing-unstable = ["dep:tracing", "dep:log", "bson3?/serde_json-1"]
 # compatible with the preview version.
 text-indexes-unstable = []
 
+# Enable support for opentelemetry instrumentation
+opentelemetry = ["dep:opentelemetry"]
+
+# Capture backtraces in errors.  This can be slow, memory intensive, and very verbose.
+error-backtrace = []
+
 [dependencies]
 base64 = "0.22"
 bitflags = "2"
@@ -101,6 +107,7 @@ mongodb-internal-macros = { path = "macros", version = "3.3.0" }
 num_cpus = { version = "1.13.1", optional = true }
 openssl = { version = "0.10.38", optional = true }
 openssl-probe = { version = "0.1.5", optional = true }
+opentelemetry = { version = "0.31.0", optional = true }
 pem = { version = "3.0.4", optional = true }
 percent-encoding = "2.0.0"
 pkcs8 = { version = "0.10.2", features = ["encryption", "pkcs5"], optional = true }
@@ -118,7 +125,7 @@ take_mut = "0.2.2"
 thiserror = "1.0.24"
 tokio-openssl = { version = "0.6.3", optional = true }
 tracing = { version = "0.1.36", optional = true }
-typed-builder = "0.20.0"
+typed-builder = "0.22.0"
 webpki-roots = "1"
 zstd = { version = "0.11.2", optional = true }
 macro_magic = "0.5.1"
@@ -216,6 +223,7 @@ futures = "0.3"
 hex = "0.4"
 home = "0.5"
 lambda_runtime = "0.6.0"
+opentelemetry_sdk = { version = "0.31.0", features = ["testing"] }
 pkcs8 = { version = "0.10.2", features = ["3des", "des-insecure", "sha1-insecure"] }
 pretty_assertions = "1.3.0"
 serde = { version = ">= 0.0.0", features = ["rc"] }

Cargo.toml

@@ -58,6 +58,7 @@ features = ["sync"]
 | `azure-oidc`         | Enable support for Azure OIDC environment authentication. |
 | `gcp-oidc`           | Enable support for GCP OIDC environment authentication. |
 | `text-indexes-unstable` | Enables support for text indexes in explicit encryption. This feature is in preview and should be used for experimental workloads only. This feature is unstable and its security is not guaranteed until released as Generally Available (GA). The GA version of this feature may not be backwards compatible with the preview version. |
+| `error-backtrace`    | Capture backtraces in `Error` values.  This can be slow, memory intensive, and very verbose. |
 
 ## Web Framework Examples
 

README.md

@@ -341,6 +341,38 @@ pub(crate) mod option_u64_as_i64 {
     }
 }
 
+/// Truncates the given string at the closest UTF-8 character boundary >= the provided length.
+/// If the new length is >= the current length, does nothing.
+#[cfg(any(feature = "tracing-unstable", feature = "opentelemetry"))]
+pub(crate) fn truncate_on_char_boundary(s: &mut String, new_len: usize) {
+    let original_len = s.len();
+    if original_len > new_len {
+        // to avoid generating invalid UTF-8, find the first index >= max_length_bytes that is
+        // the end of a character.
+        // TODO: RUST-1496 we should use ceil_char_boundary here but it's currently nightly-only.
+        // see: https://doc.rust-lang.org/std/string/struct.String.html#method.ceil_char_boundary
+        let mut truncate_index = new_len;
+        // is_char_boundary returns true when the provided value == the length of the string, so
+        // if we reach the end of the string this loop will terminate.
+        while !s.is_char_boundary(truncate_index) {
+            truncate_index += 1;
+        }
+        s.truncate(truncate_index);
+        // due to the "rounding up" behavior we might not actually end up truncating anything.
+        // if we did, spec requires we add a trailing "...".
+        if truncate_index < original_len {
+            s.push_str("...")
+        }
+    }
+}
+
+#[cfg(any(feature = "tracing-unstable", feature = "opentelemetry"))]
+pub(crate) fn doc_to_json_str(doc: crate::bson::Document, max_length_bytes: usize) -> String {
+    let mut ext_json = Bson::Document(doc).into_relaxed_extjson().to_string();
+    truncate_on_char_boundary(&mut ext_json, max_length_bytes);
+    ext_json
+}
+
 #[cfg(test)]
 mod test {
     use crate::bson_util::num_decimal_digits;

src/bson_util.rs

@@ -143,6 +143,8 @@ struct ClientInner {
     end_sessions_token: std::sync::Mutex<AsyncDropToken>,
     #[cfg(feature = "in-use-encryption")]
     csfle: tokio::sync::RwLock<Option<csfle::ClientState>>,
+    #[cfg(feature = "opentelemetry")]
+    tracer: opentelemetry::global::BoxedTracer,
     #[cfg(test)]
     disable_command_events: AtomicBool,
 }
@@ -181,6 +183,9 @@ impl Client {
             tx: Some(cleanup_tx),
         });
 
+        #[cfg(feature = "opentelemetry")]
+        let tracer = options.tracer();
+
         let inner = TrackingArc::new(ClientInner {
             topology: Topology::new(options.clone())?,
             session_pool: ServerSessionPool::new(),
@@ -193,6 +198,8 @@ impl Client {
             end_sessions_token,
             #[cfg(feature = "in-use-encryption")]
             csfle: Default::default(),
+            #[cfg(feature = "opentelemetry")]
+            tracer,
             #[cfg(test)]
             disable_command_events: AtomicBool::new(false),
         });
@@ -668,6 +675,11 @@ impl Client {
                 .await;
         }
     }
+
+    #[cfg(feature = "opentelemetry")]
+    pub(crate) fn tracer(&self) -> &opentelemetry::global::BoxedTracer {
+        &self.inner.tracer
+    }
 }
 
 #[derive(Clone, Debug)]

src/client.rs

@@ -3,6 +3,8 @@ use crate::bson::RawDocumentBuf;
 use crate::bson::{doc, RawBsonRef, RawDocument, Timestamp};
 #[cfg(feature = "in-use-encryption")]
 use futures_core::future::BoxFuture;
+#[cfg(feature = "opentelemetry")]
+use opentelemetry::context::FutureExt;
 use serde::de::DeserializeOwned;
 use std::sync::LazyLock;
 
@@ -14,6 +16,8 @@ use std::{
 };
 
 use super::{options::ServerAddress, session::TransactionState, Client, ClientSession};
+#[cfg(not(feature = "opentelemetry"))]
+use crate::otel::OtelFutureStub as _;
 use crate::{
     bson::Document,
     change_stream::{
@@ -106,7 +110,26 @@ impl Client {
         op: &mut T,
         session: impl Into<Option<&mut ClientSession>>,
     ) -> Result<ExecutionDetails<T>> {
-        // Validate inputs that can be checked before server selection and connection checkout.
+        let session = session.into();
+        #[cfg(feature = "opentelemetry")]
+        let span = self.start_operation_span(op, session.as_deref());
+        let inner = self.execute_operation_with_details_inner(op, session);
+        #[cfg(feature = "opentelemetry")]
+        let inner = inner.with_context(span.context.clone());
+        let result = inner.await;
+        #[cfg(feature = "opentelemetry")]
+        span.record_error(&result);
+
+        result
+    }
+
+    async fn execute_operation_with_details_inner<T: Operation>(
+        &self,
+        op: &mut T,
+        mut session: Option<&mut ClientSession>,
+    ) -> Result<ExecutionDetails<T>> {
+        // Validate inputs that can be checked before server selection and connection
+        // checkout.
         if self.inner.shutdown.executed.load(Ordering::SeqCst) {
             return Err(ErrorKind::Shutdown.into());
         }
@@ -122,7 +145,6 @@ impl Client {
         }
 
         // Validate the session and update its transaction status if needed.
-        let mut session = session.into();
         if let Some(ref mut session) = session {
             if !TrackingArc::ptr_eq(&self.inner, &session.client().inner) {
                 return Err(Error::invalid_argument(
@@ -154,7 +176,13 @@ impl Client {
             }
         }
 
-        Box::pin(async { self.execute_operation_with_retry(op, session).await }).await
+        Box::pin(async {
+            self.execute_operation_with_retry(op, session)
+                .with_current_context()
+                .await
+        })
+        .with_current_context()
+        .await
     }
 
     /// Execute the given operation, returning the cursor created by the operation.
@@ -408,6 +436,7 @@ impl Client {
                     retryability,
                     effective_criteria,
                 )
+                .with_current_context()
                 .await
             {
                 Ok(output) => ExecutionDetails {
@@ -496,9 +525,21 @@ impl Client {
             let should_redact = cmd.should_redact();
             let cmd_name = cmd.name.clone();
             let target_db = cmd.target_db.clone();
+            #[cfg(feature = "opentelemetry")]
+            let cmd_attrs = crate::otel::CommandAttributes::new(&cmd);
 
             let mut message = Message::try_from(cmd)?;
             message.request_id = Some(request_id);
+
+            #[cfg(feature = "opentelemetry")]
+            let span = self.start_command_span(
+                op,
+                &connection_info,
+                connection.stream_description()?,
+                &message,
+                cmd_attrs,
+            );
+
             #[cfg(feature = "in-use-encryption")]
             {
                 let guard = self.inner.csfle.read().await;
@@ -629,6 +670,8 @@ impl Client {
                     }
                 }
             };
+            #[cfg(feature = "opentelemetry")]
+            span.record_command_result::<T>(&result);
 
             if result
                 .as_ref()

src/client/executor.rs

@@ -612,6 +612,10 @@ pub struct ClientOptions {
     /// Limit on the number of mongos connections that may be created for sharded topologies.
     pub srv_max_hosts: Option<u32>,
 
+    /// Configuration for opentelemetry.
+    #[cfg(feature = "opentelemetry")]
+    pub tracing: Option<crate::otel::OpentelemetryOptions>,
+
     /// Information from the SRV URI that generated these client options, if applicable.
     #[builder(setter(skip))]
     #[serde(skip)]

src/client/options.rs

@@ -162,6 +162,8 @@ impl ClientOptions {
             tracing_max_document_length_bytes: None,
             srv_max_hosts: conn_str.srv_max_hosts,
             srv_service_name: conn_str.srv_service_name,
+            #[cfg(feature = "opentelemetry")]
+            tracing: None,
         }
     }
 }