-
Notifications
You must be signed in to change notification settings - Fork 188
RUST-2138 Add optional support for OpenTelemetry #1495
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
3ee3278
dbff5c5
ddedc25
6668f5c
d43214c
80b3363
d04dcff
a208ef3
cc36313
8f0572c
718c591
71c27aa
983f003
8c7db7a
cede087
4372184
1936b47
4069091
478951e
83390c8
2e22390
3c19ac5
a6e4913
f18f8ff
a291f61
3467389
7f03cdb
604ad7d
6923014
dc984f2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
||
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() | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The utility functions here just got moved out of
tracing
since they're needed for otel as well.