1use crate::client::Client;
4use crate::header::Header;
5use crate::namespace::Namespace;
6use crate::segment::Subsegment;
7
8#[derive(Debug)]
10pub enum SubsegmentSession<C, N>
11where
12 C: Client,
13 N: Namespace + Send + Sync,
14{
15 Entered {
17 client: C,
19 header: Header,
21 subsegment: Subsegment,
23 namespace: N,
25 },
26 Failed,
28}
29
30impl<C, N> SubsegmentSession<C, N>
31where
32 C: Client,
33 N: Namespace + Send + Sync,
34{
35 pub(crate) fn new(client: C, header: &Header, namespace: N, name_prefix: &str) -> Self {
36 let mut subsegment = Subsegment::begin(
37 header.trace_id.clone(),
38 header.parent_id.clone(),
39 namespace.name(name_prefix),
40 );
41 namespace.update_subsegment(&mut subsegment);
42 match client.send(&subsegment) {
43 Ok(_) => Self::Entered {
44 client,
45 header: header.with_parent_id(subsegment.id.clone()),
46 subsegment,
47 namespace,
48 },
49 Err(_) => Self::Failed,
50 }
51 }
52
53 pub(crate) fn failed() -> Self {
54 Self::Failed
55 }
56
57 pub fn x_amzn_trace_id(&self) -> Option<String> {
59 match self {
60 Self::Entered { header, .. } => Some(header.to_string()),
61 Self::Failed => None,
62 }
63 }
64
65 pub fn namespace_mut(&mut self) -> Option<&mut N> {
67 match self {
68 Self::Entered { namespace, .. } => Some(namespace),
69 Self::Failed => None,
70 }
71 }
72}
73
74impl<C, N> Drop for SubsegmentSession<C, N>
75where
76 C: Client,
77 N: Namespace + Send + Sync,
78{
79 fn drop(&mut self) {
80 match self {
81 Self::Entered {
82 client,
83 subsegment,
84 namespace,
85 ..
86 } => {
87 subsegment.end();
88 namespace.update_subsegment(subsegment);
89 let _ = client
90 .send(subsegment)
91 .map_err(|e| eprintln!("failed to end subsegment: {e}"));
92 }
93 Self::Failed => (),
94 }
95 }
96}