telemetry_rust/middleware/aws/instrumentation/
future.rs

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

telemetry_rust/middleware/aws/instrumentation/
future.rs
