telemetry_rust/middleware/aws/
mod.rs

1//! Instrumentation utilities for AWS SDK operations.
2//!
3//! This module provides comprehensive instrumentation for AWS services,
4//! including automatic instrumentation and a low-level API for manual span creation.
5//! It supports both individual AWS SDK operations and streaming/pagination.
6//!
7//! # Features
8//!
9//! - **Span Creation**: Manual span creation with [`AwsSpan`] and [`AwsSpanBuilder`]
10//! - **Instrumentation**: Automatic instrumentation for AWS SDK operations with [`AwsInstrument`] trait
11//! - **Stream Instrumentation**: Automatic instrumentation for AWS [`PaginationStream`](`aws_smithy_async::future::pagination_stream::PaginationStream`) with [`AwsStreamInstrument`] trait
12//!
13//! # Feature Flags
14//!
15//! - `aws-instrumentation`: Enables [`Future`] instrumentation via [`AwsInstrument`] trait
16//! - `aws-stream-instrumentation`: Enables [`Stream`][`futures_util::Stream`] instrumentation via [`AwsStreamInstrument`] trait
17
18use aws_types::request_id::RequestId;
19use opentelemetry::{
20    global::{self, BoxedSpan, BoxedTracer},
21    trace::{Span as _, SpanBuilder, SpanKind, Status, Tracer},
22};
23use std::error::Error;
24use tracing::Span;
25
26use crate::{Context, KeyValue, OpenTelemetrySpanExt, StringValue, semconv};
27
28mod instrumentation;
29mod operations;
30
31pub use instrumentation::*;
32pub use operations::*;
33
34/// A wrapper around an OpenTelemetry span specifically designed for AWS operations.
35///
36/// This struct represents an active span for an AWS SDK operation.
37/// It provides convenient methods for setting span attributes and recording
38/// AWS operation status upon its completion, including AWS request ID and optional error.
39///
40/// # Usage
41///
42/// `AwsSpan` can be used for manual instrumentation when you need fine-grained
43/// control over span lifecycle. But for most use cases, consider using the higher-level
44/// traits like [`AwsInstrument`] or [`AwsBuilderInstrument`] which provide automatic
45/// instrumentation.
46///
47/// Should be constructed using [`AwsSpanBuilder`] by calling [`AwsSpanBuilder::start`].
48///
49/// # Example
50///
51/// ```rust
52/// use aws_sdk_dynamodb::{Client as DynamoClient, types::AttributeValue};
53/// use telemetry_rust::{KeyValue, middleware::aws::DynamodbSpanBuilder, semconv};
54///
55/// async fn query_table() -> Result<i32, Box<dyn std::error::Error>> {
56///     let config = aws_config::load_from_env().await;
57///     let dynamo_client = DynamoClient::new(&config);
58///
59///     // Create and start a span manually
60///     let mut span = DynamodbSpanBuilder::query("table_name")
61///         .attribute(KeyValue::new(semconv::AWS_DYNAMODB_INDEX_NAME, "my_index"))
62///         .start();
63///
64///     let response = dynamo_client
65///         .query()
66///         .table_name("table_name")
67///         .index_name("my_index")
68///         .key_condition_expression("PK = :pk")
69///         .expression_attribute_values(":pk", AttributeValue::S("Test".to_string()))
70///         .send()
71///         .await;
72///
73///     // Add attributes from response
74///     if let Some(output) = response.as_ref().ok() {
75///         let count = output.count() as i64;
76///         let scanned_count = output.scanned_count() as i64;
77///         span.set_attributes([
78///             KeyValue::new(semconv::AWS_DYNAMODB_COUNT, count),
79///             KeyValue::new(semconv::AWS_DYNAMODB_SCANNED_COUNT, scanned_count),
80///         ]);
81///     }
82///
83///     // The span automatically handles success/error and request ID extraction
84///     span.end(&response);
85///
86///     let response = response?;
87///     println!("DynamoDB items: {:#?}", response.items());
88///     Ok(response.count())
89/// }
90/// ```
91pub struct AwsSpan {
92    span: BoxedSpan,
93}
94
95impl AwsSpan {
96    /// Ends the span with AWS response information.
97    ///
98    /// This method finalizes the span by recording the outcome of an AWS operation.
99    /// It automatically extracts request IDs and handles error reporting.
100    ///
101    /// # Arguments
102    ///
103    /// * `aws_response` - The result of the AWS operation, which must implement
104    ///   `RequestId` for both success and error cases
105    ///
106    /// # Behavior
107    ///
108    /// - On success: Sets span status to OK and records the request ID
109    /// - On error: Records the error, sets error status, and records the request ID if available
110    pub fn end<T, E>(self, aws_response: &Result<T, E>)
111    where
112        T: RequestId,
113        E: RequestId + Error,
114    {
115        let mut span = self.span;
116        let (status, request_id) = match aws_response {
117            Ok(resp) => (Status::Ok, resp.request_id()),
118            Err(error) => {
119                span.record_error(&error);
120                (Status::error(error.to_string()), error.request_id())
121            }
122        };
123        if let Some(value) = request_id {
124            span.set_attribute(KeyValue::new(semconv::AWS_REQUEST_ID, value.to_owned()));
125        }
126        span.set_status(status);
127    }
128
129    /// Sets a single attribute on the span.
130    ///
131    /// This method allows you to add custom attributes to the span after it has been created.
132    /// This is useful for adding dynamic attributes that become available during operation execution.
133    ///
134    /// For more information see [`BoxedSpan::set_attribute`]
135    ///
136    /// # Arguments
137    ///
138    /// * `attribute` - The [`KeyValue`] attribute to add to the span
139    ///
140    /// # Example
141    ///
142    /// ```rust
143    /// use telemetry_rust::{KeyValue, middleware::aws::AwsSpanBuilder};
144    ///
145    /// let mut span = AwsSpanBuilder::client("DynamoDB", "GetItem", []).start();
146    /// span.set_attribute(KeyValue::new("custom.attribute", "value"));
147    /// ```
148    pub fn set_attribute(&mut self, attribute: KeyValue) {
149        self.span.set_attribute(attribute);
150    }
151
152    /// Sets multiple attributes on the span.
153    ///
154    /// This method allows you to add multiple custom attributes to the span at once.
155    /// This is more efficient than calling `set_attribute` multiple times.
156    ///
157    /// For more information see [`BoxedSpan::set_attributes`]
158    ///
159    /// # Arguments
160    ///
161    /// * `attributes` - An iterator of [`KeyValue`] attributes to add to the span
162    ///
163    /// # Example
164    ///
165    /// ```rust
166    /// use telemetry_rust::{KeyValue, middleware::aws::AwsSpanBuilder, semconv};
167    ///
168    /// let mut span = AwsSpanBuilder::client("DynamoDB", "GetItem", []).start();
169    /// span.set_attributes([
170    ///     KeyValue::new(semconv::DB_NAMESPACE, "my_table"),
171    ///     KeyValue::new("custom.attribute", "value"),
172    /// ]);
173    /// ```
174    pub fn set_attributes(&mut self, attributes: impl IntoIterator<Item = KeyValue>) {
175        self.span.set_attributes(attributes);
176    }
177}
178
179impl From<BoxedSpan> for AwsSpan {
180    #[inline]
181    fn from(span: BoxedSpan) -> Self {
182        Self { span }
183    }
184}
185
186/// Builder for creating AWS-specific OpenTelemetry spans.
187///
188/// This builder provides a fluent interface for constructing [`AwsSpan`] with
189/// required attributes and proper span kinds for different types of AWS operations.
190/// It automatically sets standard RPC attributes following OpenTelemetry semantic
191/// conventions for AWS services.
192///
193/// # Usage
194///
195/// This builder can be used with [`AwsInstrument`] trait to instrument any AWS operation,
196/// or to manually create [`AwsSpan`] if you need control over span lifecycle.
197/// For automatic instrumentation, use [`AwsBuilderInstrument`] trait.
198pub struct AwsSpanBuilder<'a> {
199    inner: SpanBuilder,
200    tracer: BoxedTracer,
201    context: Option<&'a Context>,
202}
203
204impl<'a> AwsSpanBuilder<'a> {
205    fn new(
206        span_kind: SpanKind,
207        service: impl Into<StringValue>,
208        method: impl Into<StringValue>,
209        custom_attributes: impl IntoIterator<Item = KeyValue>,
210    ) -> Self {
211        let service: StringValue = service.into();
212        let method: StringValue = method.into();
213        let tracer = global::tracer("aws_sdk");
214        let span_name = format!("{service}.{method}");
215        let mut attributes = vec![
216            KeyValue::new(semconv::RPC_METHOD, method),
217            KeyValue::new(semconv::RPC_SYSTEM, "aws-api"),
218            KeyValue::new(semconv::RPC_SERVICE, service),
219        ];
220        attributes.extend(custom_attributes);
221        let inner = tracer
222            .span_builder(span_name)
223            .with_attributes(attributes)
224            .with_kind(span_kind);
225
226        Self {
227            inner,
228            tracer,
229            context: None,
230        }
231    }
232
233    /// Creates a client span builder for AWS operations.
234    ///
235    /// Client spans represent outbound calls to AWS services from your application.
236    ///
237    /// # Arguments
238    ///
239    /// * `service` - The AWS service name (e.g., "S3", "DynamoDB")
240    /// * `method` - The operation name (e.g., "GetObject", "PutItem")
241    /// * `attributes` - Additional custom attributes for the span
242    pub fn client(
243        service: impl Into<StringValue>,
244        method: impl Into<StringValue>,
245        attributes: impl IntoIterator<Item = KeyValue>,
246    ) -> Self {
247        Self::new(SpanKind::Client, service, method, attributes)
248    }
249
250    /// Creates a producer span builder for AWS operations.
251    ///
252    /// Producer spans represent operations that send messages or data to AWS services.
253    ///
254    /// # Arguments
255    ///
256    /// * `service` - The AWS service name (e.g., "SQS", "SNS")
257    /// * `method` - The operation name (e.g., "SendMessage", "Publish")
258    /// * `attributes` - Additional custom attributes for the span
259    pub fn producer(
260        service: impl Into<StringValue>,
261        method: impl Into<StringValue>,
262        attributes: impl IntoIterator<Item = KeyValue>,
263    ) -> Self {
264        Self::new(SpanKind::Producer, service, method, attributes)
265    }
266
267    /// Creates a consumer span builder for AWS operations.
268    ///
269    /// Consumer spans represent operations that receive messages or data from AWS services.
270    ///
271    /// # Arguments
272    ///
273    /// * `service` - The AWS service name (e.g., "SQS", "Kinesis")
274    /// * `method` - The operation name (e.g., "ReceiveMessage", "GetRecords")
275    /// * `attributes` - Additional custom attributes for the span
276    pub fn consumer(
277        service: impl Into<StringValue>,
278        method: impl Into<StringValue>,
279        attributes: impl IntoIterator<Item = KeyValue>,
280    ) -> Self {
281        Self::new(SpanKind::Consumer, service, method, attributes)
282    }
283
284    /// Adds multiple attributes to the span being built.
285    ///
286    /// # Arguments
287    ///
288    /// * `iter` - An iterator of [`KeyValue`] attributes to add to the span
289    pub fn attributes(mut self, iter: impl IntoIterator<Item = KeyValue>) -> Self {
290        if let Some(attributes) = &mut self.inner.attributes {
291            attributes.extend(iter);
292        }
293        self
294    }
295
296    /// Adds a single attribute to the span being built.
297    ///
298    /// This is a convenience method for adding one attribute at a time.
299    ///
300    /// # Arguments
301    ///
302    /// * `attribute` - The [`KeyValue`] attribute to add to the span
303    #[inline]
304    pub fn attribute(self, attribute: KeyValue) -> Self {
305        self.attributes(std::iter::once(attribute))
306    }
307
308    /// Sets the parent [`Context`] for the span.
309    ///
310    /// # Arguments
311    ///
312    /// * `context` - The OpenTelemetry [`Context`] to use as the parent
313    #[inline]
314    pub fn context(mut self, context: &'a Context) -> Self {
315        self.context = Some(context);
316        self
317    }
318
319    /// Optionally sets the parent [`Context`] for the span.
320    ///
321    /// # Arguments
322    ///
323    /// * `context` - An optional OpenTelemetry [`Context`] to use as the parent
324    #[inline]
325    pub fn set_context(mut self, context: Option<&'a Context>) -> Self {
326        self.context = context;
327        self
328    }
329
330    #[inline(always)]
331    fn start_with_context(self, parent_cx: &Context) -> AwsSpan {
332        self.inner
333            .start_with_context(&self.tracer, parent_cx)
334            .into()
335    }
336
337    /// Starts the span and returns an [`AwsSpan`].
338    ///
339    /// This method creates and starts the span using either the explicitly set context
340    /// or the current tracing span's context as the parent.
341    #[inline]
342    pub fn start(self) -> AwsSpan {
343        match self.context {
344            Some(context) => self.start_with_context(context),
345            None => self.start_with_context(&Span::current().context()),
346        }
347    }
348}
telemetry_rust/middleware/aws/mod.rs

telemetry_rust/middleware/aws/
mod.rs
