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

telemetry_rust/middleware/aws/
mod.rs
