msgpack_tracing/
string_cache.rs

1use crate::tape::{
2    FieldValue, Instruction, InstructionId, InstructionSet, InstructionSetTrait, InstructionTrait,
3    TapeMachine, Value,
4};
5use chrono::{DateTime, Utc};
6use std::{collections::HashMap, num::NonZeroU64};
7use tracing::Level;
8
9#[derive(Clone, Copy, Debug)]
10pub enum CacheInstruction<'a> {
11    Restart,
12    NewString(&'a str),
13    NewSpan {
14        parent: Option<NonZeroU64>,
15        span: NonZeroU64,
16        name: CacheString<'a>,
17    },
18    FinishedSpan,
19    NewRecord(NonZeroU64),
20    FinishedRecord,
21    StartEvent {
22        time: DateTime<Utc>,
23        span: Option<NonZeroU64>,
24        target: CacheString<'a>,
25        priority: Level,
26    },
27    FinishedEvent,
28    AddValue(FieldValue<'a, CacheString<'a>>),
29    DeleteSpan(NonZeroU64),
30}
31impl InstructionTrait for CacheInstruction<'_> {
32    fn id(self) -> InstructionId {
33        match self {
34            CacheInstruction::Restart => InstructionId::Restart,
35            CacheInstruction::NewString(..) => InstructionId::NewString,
36            CacheInstruction::NewSpan { .. } => InstructionId::NewSpan,
37            CacheInstruction::FinishedSpan => InstructionId::FinishedSpan,
38            CacheInstruction::NewRecord(..) => InstructionId::NewRecord,
39            CacheInstruction::FinishedRecord => InstructionId::FinishedRecord,
40            CacheInstruction::StartEvent { .. } => InstructionId::StartEvent,
41            CacheInstruction::FinishedEvent => InstructionId::FinishedEvent,
42            CacheInstruction::AddValue(..) => InstructionId::AddValue,
43            CacheInstruction::DeleteSpan(..) => InstructionId::DeleteSpan,
44        }
45    }
46}
47
48pub struct CacheInstructionSet;
49impl InstructionSetTrait for CacheInstructionSet {
50    type Instruction<'a> = CacheInstruction<'a>;
51}
52
53#[derive(Clone, Copy, Debug)]
54pub enum CacheString<'a> {
55    Present(&'a str),
56    Cached(u64),
57}
58
59pub struct StringCache<T> {
60    forward: T,
61    strings: HashMap<String, u64>,
62}
63impl<T> StringCache<T>
64where
65    T: TapeMachine<CacheInstructionSet>,
66{
67    pub fn new(forward: T) -> Self {
68        Self {
69            forward,
70            strings: Default::default(),
71        }
72    }
73
74    fn cache_value<'a>(&mut self, value: Value<'a, &'a str>) -> Value<'a, CacheString<'a>> {
75        match value {
76            Value::Debug(string) => Value::Debug(self.cache_string(string)),
77            Value::String(string) => Value::String(self.cache_string(string)),
78            Value::Float(value) => Value::Float(value),
79            Value::Integer(value) => Value::Integer(value),
80            Value::Unsigned(value) => Value::Unsigned(value),
81            Value::Bool(value) => Value::Bool(value),
82            Value::ByteArray(value) => Value::ByteArray(value),
83        }
84    }
85
86    fn cache_string<'a>(&mut self, string: &'a str) -> CacheString<'a> {
87        if let Some(id) = self.strings.get(string) {
88            return CacheString::Cached(*id);
89        }
90
91        let id = self.strings.len() as u64;
92        let present = !matches!(
93            (id, string.len()),
94            (0..=0xffff, 4..=255)
95                | (0x1_0000..=0xff_ffff, 5..=255)
96                | (0x100_0000..=0xff_ffff_ffff, 7..=255)
97                | (_, 11..)
98        );
99
100        if present {
101            CacheString::Present(string)
102        } else {
103            self.forward.handle(CacheInstruction::NewString(string));
104            self.strings.insert(string.to_owned(), id);
105            CacheString::Cached(id)
106        }
107    }
108}
109impl<T> TapeMachine<InstructionSet> for StringCache<T>
110where
111    T: TapeMachine<CacheInstructionSet>,
112{
113    fn needs_restart(&mut self) -> bool {
114        self.forward.needs_restart()
115    }
116
117    fn handle(&mut self, instruction: Instruction) {
118        match instruction {
119            Instruction::Restart => {
120                self.strings.clear();
121                self.forward.handle(CacheInstruction::Restart);
122            }
123            Instruction::NewSpan { parent, span, name } => {
124                let name = self.cache_string(name);
125                self.forward
126                    .handle(CacheInstruction::NewSpan { parent, span, name });
127            }
128            Instruction::FinishedSpan => {
129                self.forward.handle(CacheInstruction::FinishedSpan);
130            }
131            Instruction::NewRecord(span) => {
132                self.forward.handle(CacheInstruction::NewRecord(span));
133            }
134            Instruction::FinishedRecord => {
135                self.forward.handle(CacheInstruction::FinishedRecord);
136            }
137            Instruction::StartEvent {
138                time,
139                span,
140                target,
141                priority,
142            } => {
143                let target = self.cache_string(target);
144                self.forward.handle(CacheInstruction::StartEvent {
145                    time,
146                    span,
147                    target,
148                    priority,
149                });
150            }
151            Instruction::FinishedEvent => {
152                self.forward.handle(CacheInstruction::FinishedEvent);
153            }
154            Instruction::AddValue(FieldValue { name, value }) => {
155                let name = self.cache_string(name);
156                let value = self.cache_value(value);
157                self.forward
158                    .handle(CacheInstruction::AddValue(FieldValue { name, value }));
159            }
160            Instruction::DeleteSpan(span) => {
161                self.forward.handle(CacheInstruction::DeleteSpan(span));
162            }
163        }
164    }
165}
166
167pub struct StringUncache<T> {
168    forward: T,
169    strings: Vec<String>,
170}
171impl<T> StringUncache<T>
172where
173    T: TapeMachine<InstructionSet>,
174{
175    pub fn new(forward: T) -> Self {
176        Self {
177            forward,
178            strings: Default::default(),
179        }
180    }
181
182    fn uncache<'a>(strings: &'a [String], string: CacheString<'a>) -> &'a str {
183        match string {
184            CacheString::Present(str) => str,
185            CacheString::Cached(index) => strings[index as usize].as_str(),
186        }
187    }
188
189    fn uncache_value<'a>(
190        strings: &'a [String],
191        value: Value<'a, CacheString<'a>>,
192    ) -> Value<'a, &'a str> {
193        match value {
194            Value::Debug(string) => Value::Debug(Self::uncache(strings, string)),
195            Value::String(string) => Value::String(Self::uncache(strings, string)),
196            Value::Float(value) => Value::Float(value),
197            Value::Integer(value) => Value::Integer(value),
198            Value::Unsigned(value) => Value::Unsigned(value),
199            Value::Bool(value) => Value::Bool(value),
200            Value::ByteArray(items) => Value::ByteArray(items),
201        }
202    }
203}
204impl<T> TapeMachine<CacheInstructionSet> for StringUncache<T>
205where
206    T: TapeMachine<InstructionSet>,
207{
208    fn needs_restart(&mut self) -> bool {
209        self.forward.needs_restart()
210    }
211
212    fn handle(&mut self, instruction: CacheInstruction) {
213        match instruction {
214            CacheInstruction::Restart => {
215                self.forward.handle(Instruction::Restart);
216            }
217            CacheInstruction::NewString(str) => {
218                self.strings.push(str.to_owned());
219            }
220            CacheInstruction::NewSpan { parent, span, name } => {
221                let name = Self::uncache(&self.strings, name);
222                self.forward
223                    .handle(Instruction::NewSpan { parent, span, name });
224            }
225            CacheInstruction::FinishedSpan => {
226                self.forward.handle(Instruction::FinishedSpan);
227            }
228            CacheInstruction::NewRecord(span) => {
229                self.forward.handle(Instruction::NewRecord(span));
230            }
231            CacheInstruction::FinishedRecord => {
232                self.forward.handle(Instruction::FinishedRecord);
233            }
234            CacheInstruction::StartEvent {
235                time,
236                span,
237                target,
238                priority,
239            } => {
240                let target = Self::uncache(&self.strings, target);
241
242                self.forward.handle(Instruction::StartEvent {
243                    time,
244                    span,
245                    target,
246                    priority,
247                });
248            }
249            CacheInstruction::FinishedEvent => {
250                self.forward.handle(Instruction::FinishedEvent);
251            }
252            CacheInstruction::AddValue(FieldValue { name, value }) => {
253                let name = Self::uncache(&self.strings, name);
254                let value = Self::uncache_value(&self.strings, value);
255                self.forward
256                    .handle(Instruction::AddValue(FieldValue { name, value }));
257            }
258            CacheInstruction::DeleteSpan(span) => {
259                self.forward.handle(Instruction::DeleteSpan(span));
260            }
261        }
262    }
263}
msgpack_tracing/string_cache.rs

msgpack_tracing/
string_cache.rs
