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, StringValue, 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
192impl From<BoxedSpan> for AwsSpan {
193    #[inline]
194    fn from(span: BoxedSpan) -> Self {
195        Self { span }
196    }
197}
198
199/// Builder for creating AWS-specific OpenTelemetry spans.
200///
201/// This builder provides a fluent interface for constructing [`AwsSpan`] with
202/// required attributes and proper span kinds for different types of AWS operations.
203/// It automatically sets standard RPC attributes following OpenTelemetry semantic
204/// conventions for AWS services.
205///
206/// # Usage
207///
208/// This builder can be used with [`AwsInstrument`] trait to instrument any AWS operation,
209/// or to manually create [`AwsSpan`] if you need control over span lifecycle.
210/// For automatic instrumentation, use [`AwsBuilderInstrument`] trait.
211pub struct AwsSpanBuilder<'a> {
212    inner: SpanBuilder,
213    tracer: BoxedTracer,
214    context: Option<&'a Context>,
215}
216
217impl<'a> AwsSpanBuilder<'a> {
218    fn new(
219        span_kind: SpanKind,
220        service: impl Into<StringValue>,
221        method: impl Into<StringValue>,
222        custom_attributes: impl IntoIterator<Item = KeyValue>,
223    ) -> Self {
224        let service: StringValue = service.into();
225        let method: StringValue = method.into();
226        let tracer = global::tracer("aws_sdk");
227        let span_name = format!("{service}.{method}");
228        let mut attributes = vec![
229            KeyValue::new(semconv::RPC_METHOD, method),
230            KeyValue::new(semconv::RPC_SYSTEM, "aws-api"),
231            KeyValue::new(semconv::RPC_SERVICE, service),
232        ];
233        attributes.extend(custom_attributes);
234        let inner = tracer
235            .span_builder(span_name)
236            .with_attributes(attributes)
237            .with_kind(span_kind);
238
239        Self {
240            inner,
241            tracer,
242            context: None,
243        }
244    }
245
246    /// Creates a client span builder for AWS operations.
247    ///
248    /// Client spans represent outbound calls to AWS services from your application.
249    ///
250    /// # Arguments
251    ///
252    /// * `service` - The AWS service name (e.g., "S3", "DynamoDB")
253    /// * `method` - The operation name (e.g., "GetObject", "PutItem")
254    /// * `attributes` - Additional custom attributes for the span
255    pub fn client(
256        service: impl Into<StringValue>,
257        method: impl Into<StringValue>,
258        attributes: impl IntoIterator<Item = KeyValue>,
259    ) -> Self {
260        Self::new(SpanKind::Client, service, method, attributes)
261    }
262
263    /// Creates a producer span builder for AWS operations.
264    ///
265    /// Producer spans represent operations that send messages or data to AWS services.
266    ///
267    /// # Arguments
268    ///
269    /// * `service` - The AWS service name (e.g., "SQS", "SNS")
270    /// * `method` - The operation name (e.g., "SendMessage", "Publish")
271    /// * `attributes` - Additional custom attributes for the span
272    pub fn producer(
273        service: impl Into<StringValue>,
274        method: impl Into<StringValue>,
275        attributes: impl IntoIterator<Item = KeyValue>,
276    ) -> Self {
277        Self::new(SpanKind::Producer, service, method, attributes)
278    }
279
280    /// Creates a consumer span builder for AWS operations.
281    ///
282    /// Consumer spans represent operations that receive messages or data from AWS services.
283    ///
284    /// # Arguments
285    ///
286    /// * `service` - The AWS service name (e.g., "SQS", "Kinesis")
287    /// * `method` - The operation name (e.g., "ReceiveMessage", "GetRecords")
288    /// * `attributes` - Additional custom attributes for the span
289    pub fn consumer(
290        service: impl Into<StringValue>,
291        method: impl Into<StringValue>,
292        attributes: impl IntoIterator<Item = KeyValue>,
293    ) -> Self {
294        Self::new(SpanKind::Consumer, service, method, attributes)
295    }
296
297    /// Adds multiple attributes to the span being built.
298    ///
299    /// # Arguments
300    ///
301    /// * `iter` - An iterator of [`KeyValue`] attributes to add to the span
302    pub fn attributes(mut self, iter: impl IntoIterator<Item = KeyValue>) -> Self {
303        if let Some(attributes) = &mut self.inner.attributes {
304            attributes.extend(iter);
305        }
306        self
307    }
308
309    /// Adds a single attribute to the span being built.
310    ///
311    /// This is a convenience method for adding one attribute at a time.
312    ///
313    /// # Arguments
314    ///
315    /// * `attribute` - The [`KeyValue`] attribute to add to the span
316    #[inline]
317    pub fn attribute(self, attribute: KeyValue) -> Self {
318        self.attributes(std::iter::once(attribute))
319    }
320
321    /// Sets the parent [`Context`] for the span.
322    ///
323    /// # Arguments
324    ///
325    /// * `context` - The OpenTelemetry [`Context`] to use as the parent
326    #[inline]
327    pub fn context(mut self, context: &'a Context) -> Self {
328        self.context = Some(context);
329        self
330    }
331
332    /// Optionally sets the parent [`Context`] for the span.
333    ///
334    /// # Arguments
335    ///
336    /// * `context` - An optional OpenTelemetry [`Context`] to use as the parent
337    #[inline]
338    pub fn set_context(mut self, context: Option<&'a Context>) -> Self {
339        self.context = context;
340        self
341    }
342
343    #[inline(always)]
344    fn start_with_context(self, parent_cx: &Context) -> AwsSpan {
345        self.inner
346            .start_with_context(&self.tracer, parent_cx)
347            .into()
348    }
349
350    /// Starts the span and returns an [`AwsSpan`].
351    ///
352    /// This method creates and starts the span using either the explicitly set context
353    /// or the current tracing span's context as the parent.
354    #[inline]
355    pub fn start(self) -> AwsSpan {
356        match self.context {
357            Some(context) => self.start_with_context(context),
358            None => self.start_with_context(&Span::current().context()),
359        }
360    }
361}
telemetry_rust/middleware/aws/mod.rs

telemetry_rust/middleware/aws/
mod.rs
