telemetry_rust/middleware/aws/instrumentation/
future.rs

1use aws_types::request_id::RequestId;
2use std::{error::Error, future::Future};
3
4use crate::{
5    future::{InstrumentedFuture, InstrumentedFutureContext},
6    middleware::aws::{AwsSpan, AwsSpanBuilder},
7};
8
9impl<T, E> InstrumentedFutureContext<Result<T, E>> for AwsSpan
10where
11    T: RequestId,
12    E: RequestId + Error,
13{
14    fn on_result(self, result: &Result<T, E>) {
15        self.end(result);
16    }
17}
18
19/// Trait for instrumenting AWS futures with automatic span management.
20///
21/// This trait provides a convenient way to wrap AWS SDK futures with OpenTelemetry
22/// instrumentation, automatically handling span creation, error recording, and cleanup.
23///
24/// For automatic span attributes extraction, use [`AwsBuilderInstrument`][`super::AwsBuilderInstrument`] trait.
25///
26/// # Example
27///
28/// ```rust
29/// use aws_sdk_dynamodb::{Client as DynamoClient, types::AttributeValue};
30/// use telemetry_rust::{
31///     KeyValue,
32///     middleware::aws::{AwsInstrument, DynamodbSpanBuilder},
33///     semconv,
34/// };
35///
36/// async fn query_table() -> Result<i32, Box<dyn std::error::Error>> {
37///     let config = aws_config::load_from_env().await;
38///     let dynamo_client = DynamoClient::new(&config);
39///     let resp =
40///         dynamo_client
41///             .query()
42///             .table_name("table_name")
43///             .index_name("my_index")
44///             .key_condition_expression("PK = :pk")
45///             .expression_attribute_values(":pk", AttributeValue::S("Test".to_string()))
46///             .send()
47///             .instrument(DynamodbSpanBuilder::query("table_name").attribute(
48///                 KeyValue::new(semconv::AWS_DYNAMODB_INDEX_NAME, "my_index"),
49///             ))
50///             .await?;
51///     println!("DynamoDB items: {:#?}", resp.items());
52///     Ok(resp.count())
53/// }
54/// ```
55pub trait AwsInstrument<T, E, F>
56where
57    T: RequestId,
58    E: RequestId + Error,
59    F: Future<Output = Result<T, E>>,
60{
61    /// Instruments the future with an AWS span.
62    ///
63    /// Creates an instrumented future that will automatically start a span when polled
64    /// and properly handle the result when the future completes.
65    ///
66    /// # Arguments
67    ///
68    /// * `span` - A span builder or span to use for instrumentation
69    ///
70    /// # Returns
71    ///
72    /// An instrumented future that will record AWS operation details
73    fn instrument<'a>(
74        self,
75        span: impl Into<AwsSpanBuilder<'a>>,
76    ) -> InstrumentedFuture<F, AwsSpan>;
77}
78
79impl<T, E, F> AwsInstrument<T, E, F> for F
80where
81    T: RequestId,
82    E: RequestId + Error,
83    F: Future<Output = Result<T, E>>,
84{
85    fn instrument<'a>(
86        self,
87        span: impl Into<AwsSpanBuilder<'a>>,
88    ) -> InstrumentedFuture<F, AwsSpan> {
89        let span = span.into().start();
90        InstrumentedFuture::new(self, span)
91    }
92}
telemetry_rust/middleware/aws/instrumentation/future.rs

telemetry_rust/middleware/aws/instrumentation/
future.rs
