Skip to main content

launchdarkly_server_sdk/migrations/
tracker.rs

1use std::{
2    collections::{HashMap, HashSet},
3    time::Duration,
4};
5
6use launchdarkly_server_sdk_evaluation::{Context, Detail, Flag};
7use rand::rng;
8
9use crate::{
10    events::event::{BaseEvent, EventFactory, MigrationOpEvent},
11    sampler::{Sampler, ThreadRngSampler},
12};
13
14use super::{Operation, Origin, Stage};
15
16/// A MigrationOpTracker is responsible for managing the collection of measurements that a user
17/// might wish to record throughout a migration-assisted operation.
18///
19/// Example measurements include latency, errors, and consistency.
20pub struct MigrationOpTracker {
21    key: String,
22    flag: Option<Flag>,
23    context: Context,
24    detail: Detail<Stage>,
25    default_stage: Stage,
26    operation: Option<Operation>,
27    invoked: HashSet<Origin>,
28    consistent: Option<bool>,
29    consistent_ratio: Option<u32>,
30    errors: HashSet<Origin>,
31    latencies: HashMap<Origin, Duration>,
32}
33
34impl MigrationOpTracker {
35    pub(crate) fn new(
36        key: String,
37        flag: Option<Flag>,
38        context: Context,
39        detail: Detail<Stage>,
40        default_stage: Stage,
41    ) -> Self {
42        let consistent_ratio = match &flag {
43            Some(f) => f
44                .migration_settings
45                .as_ref()
46                .map(|s| s.check_ratio.unwrap_or(1)),
47            None => None,
48        };
49
50        Self {
51            key,
52            flag,
53            context,
54            detail,
55            default_stage,
56            operation: None,
57            invoked: HashSet::new(),
58            consistent: None,
59            consistent_ratio,
60            errors: HashSet::new(),
61            latencies: HashMap::new(),
62        }
63    }
64
65    /// Sets the migration related operation associated with these tracking measurements.
66    pub fn operation(&mut self, operation: Operation) {
67        self.operation = Some(operation);
68    }
69
70    /// Allows recording which origins were called during a migration.
71    pub fn invoked(&mut self, origin: Origin) {
72        self.invoked.insert(origin);
73    }
74
75    /// This method accepts a callable which should take no parameters and return a single boolean
76    /// to represent the consistency check results for a read operation.
77    ///
78    /// A callable is provided in case sampling rules do not require consistency checking to run.
79    /// In this case, we can avoid the overhead of a function by not using the callable.
80    pub fn consistent(&mut self, is_consistent: impl Fn() -> bool) {
81        if ThreadRngSampler::new(rng()).sample(self.consistent_ratio.unwrap_or(1)) {
82            self.consistent = Some(is_consistent());
83        }
84    }
85
86    /// Allows recording which origins were called during a migration.
87    pub fn error(&mut self, origin: Origin) {
88        self.errors.insert(origin);
89    }
90
91    /// Allows tracking the recorded latency for an individual operation.
92    pub fn latency(&mut self, origin: Origin, latency: Duration) {
93        if latency.is_zero() {
94            return;
95        }
96
97        self.latencies.insert(origin, latency);
98    }
99
100    /// Creates an instance of [crate::MigrationOpEvent]. This event data can be
101    /// provided to the [crate::Client::track_migration_op] method to rely this metric
102    /// information upstream to LaunchDarkly services.
103    pub fn build(&self) -> Result<MigrationOpEvent, String> {
104        let operation = self
105            .operation
106            .ok_or_else(|| "operation not provided".to_string())?;
107
108        self.check_invoked_consistency()?;
109
110        if self.key.is_empty() {
111            return Err("operation cannot contain an empty key".to_string());
112        }
113
114        let invoked = self.invoked.clone();
115        if invoked.is_empty() {
116            return Err("no origins were invoked".to_string());
117        }
118
119        Ok(MigrationOpEvent {
120            base: BaseEvent::new(EventFactory::now(), self.context.clone()),
121            key: self.key.clone(),
122            version: self.flag.as_ref().map(|f| f.version),
123            operation,
124            default_stage: self.default_stage,
125            evaluation: self.detail.clone(),
126            invoked,
127            consistency_check_ratio: self.consistent_ratio,
128            consistency_check: self.consistent,
129            errors: self.errors.clone(),
130            latency: self.latencies.clone(),
131            sampling_ratio: self.flag.as_ref().and_then(|f| f.sampling_ratio),
132        })
133    }
134
135    fn check_invoked_consistency(&self) -> Result<(), String> {
136        for origin in [Origin::Old, Origin::New].iter() {
137            if self.invoked.contains(origin) {
138                continue;
139            }
140
141            if self.errors.contains(origin) {
142                return Err(format!(
143                    "provided error for origin {origin:?} without recording invocation"
144                ));
145            }
146
147            if self.latencies.contains_key(origin) {
148                return Err(format!(
149                    "provided latency for origin {origin:?} without recording invocation"
150                ));
151            }
152        }
153
154        if self.consistent.is_some() && self.invoked.len() != 2 {
155            return Err("provided consistency without recording both invocations".to_string());
156        }
157
158        Ok(())
159    }
160}
161
162#[cfg(test)]
163mod tests {
164
165    use launchdarkly_server_sdk_evaluation::{
166        ContextBuilder, Detail, Flag, MigrationFlagParameters, Reason,
167    };
168    use test_case::test_case;
169
170    use super::{MigrationOpTracker, Operation, Origin, Stage};
171    use crate::test_common::basic_flag;
172
173    fn minimal_tracker(flag: Flag) -> MigrationOpTracker {
174        let mut tracker = MigrationOpTracker::new(
175            flag.key.clone(),
176            Some(flag),
177            ContextBuilder::new("user")
178                .build()
179                .expect("failed to build context"),
180            Detail {
181                value: Some(Stage::Live),
182                variation_index: Some(1),
183                reason: Reason::Fallthrough {
184                    in_experiment: false,
185                },
186            },
187            Stage::Live,
188        );
189        tracker.operation(Operation::Read);
190        tracker.invoked(Origin::Old);
191        tracker.invoked(Origin::New);
192
193        tracker
194    }
195
196    #[test]
197    fn build_minimal_tracker() {
198        let tracker = minimal_tracker(basic_flag("flag-key"));
199        let result = tracker.build();
200
201        assert!(result.is_ok());
202    }
203
204    #[test]
205    fn build_without_flag() {
206        let mut tracker = minimal_tracker(basic_flag("flag-key"));
207        tracker.flag = None;
208        let result = tracker.build();
209
210        assert!(result.is_ok());
211    }
212
213    #[test_case(Origin::Old)]
214    #[test_case(Origin::New)]
215    fn track_invocations_individually(origin: Origin) {
216        let mut tracker = MigrationOpTracker::new(
217            "flag-key".into(),
218            Some(basic_flag("flag-key")),
219            ContextBuilder::new("user")
220                .build()
221                .expect("failed to build context"),
222            Detail {
223                value: Some(Stage::Live),
224                variation_index: Some(1),
225                reason: Reason::Fallthrough {
226                    in_experiment: false,
227                },
228            },
229            Stage::Live,
230        );
231        tracker.operation(Operation::Read);
232        tracker.invoked(origin);
233
234        let event = tracker.build().expect("failed to build event");
235        assert_eq!(event.invoked.len(), 1);
236        assert!(event.invoked.contains(&origin));
237    }
238
239    #[test]
240    fn tracks_both_invocations() {
241        let mut tracker = MigrationOpTracker::new(
242            "flag-key".into(),
243            Some(basic_flag("flag-key")),
244            ContextBuilder::new("user")
245                .build()
246                .expect("failed to build context"),
247            Detail {
248                value: Some(Stage::Live),
249                variation_index: Some(1),
250                reason: Reason::Fallthrough {
251                    in_experiment: false,
252                },
253            },
254            Stage::Live,
255        );
256        tracker.operation(Operation::Read);
257        tracker.invoked(Origin::Old);
258        tracker.invoked(Origin::New);
259
260        let event = tracker.build().expect("failed to build event");
261        assert_eq!(event.invoked.len(), 2);
262        assert!(event.invoked.contains(&Origin::Old));
263        assert!(event.invoked.contains(&Origin::New));
264    }
265
266    #[test_case(false)]
267    #[test_case(true)]
268    fn tracks_consistency(expectation: bool) {
269        let mut tracker = minimal_tracker(basic_flag("flag-key"));
270        tracker.operation(Operation::Read);
271        tracker.consistent(|| expectation);
272
273        let event = tracker.build().expect("failed to build event");
274        assert_eq!(event.consistency_check, Some(expectation));
275        assert_eq!(event.consistency_check_ratio, None);
276    }
277
278    #[test_case(false)]
279    #[test_case(true)]
280    fn consistency_can_be_disabled_through_sampling_ratio(expectation: bool) {
281        let mut flag = basic_flag("flag-key");
282        flag.migration_settings = Some(MigrationFlagParameters {
283            check_ratio: Some(0),
284        });
285
286        let mut tracker = minimal_tracker(flag);
287        tracker.operation(Operation::Read);
288        tracker.consistent(|| expectation);
289
290        let event = tracker.build().expect("failed to build event");
291        assert_eq!(event.consistency_check, None);
292        assert_eq!(event.consistency_check_ratio, Some(0));
293    }
294
295    #[test_case(Origin::Old)]
296    #[test_case(Origin::New)]
297    fn track_errors_individually(origin: Origin) {
298        let mut tracker = minimal_tracker(basic_flag("flag-key"));
299        tracker.error(origin);
300
301        let event = tracker.build().expect("failed to build event");
302        assert_eq!(event.errors.len(), 1);
303        assert!(event.errors.contains(&origin));
304    }
305
306    #[test]
307    fn tracks_both_errors() {
308        let mut tracker = minimal_tracker(basic_flag("flag-key"));
309        tracker.error(Origin::Old);
310        tracker.error(Origin::New);
311
312        let event = tracker.build().expect("failed to build event");
313        assert_eq!(event.errors.len(), 2);
314        assert!(event.errors.contains(&Origin::Old));
315        assert!(event.errors.contains(&Origin::New));
316    }
317
318    #[test_case(Origin::Old)]
319    #[test_case(Origin::New)]
320    fn track_latencies_individually(origin: Origin) {
321        let mut tracker = minimal_tracker(basic_flag("flag-key"));
322        tracker.latency(origin, std::time::Duration::from_millis(100));
323
324        let event = tracker.build().expect("failed to build event");
325        assert_eq!(event.latency.len(), 1);
326        assert_eq!(
327            event.latency.get(&origin),
328            Some(&std::time::Duration::from_millis(100))
329        );
330    }
331
332    #[test]
333    fn track_both_latencies() {
334        let mut tracker = minimal_tracker(basic_flag("flag-key"));
335        tracker.latency(Origin::Old, std::time::Duration::from_millis(100));
336        tracker.latency(Origin::New, std::time::Duration::from_millis(200));
337
338        let event = tracker.build().expect("failed to build event");
339        assert_eq!(event.latency.len(), 2);
340        assert_eq!(
341            event.latency.get(&Origin::Old),
342            Some(&std::time::Duration::from_millis(100))
343        );
344        assert_eq!(
345            event.latency.get(&Origin::New),
346            Some(&std::time::Duration::from_millis(200))
347        );
348    }
349
350    #[test]
351    fn fails_without_calling_invocations() {
352        let mut tracker = MigrationOpTracker::new(
353            "flag-key".into(),
354            Some(basic_flag("flag-key")),
355            ContextBuilder::new("user")
356                .build()
357                .expect("failed to build context"),
358            Detail {
359                value: Some(Stage::Live),
360                variation_index: Some(1),
361                reason: Reason::Fallthrough {
362                    in_experiment: false,
363                },
364            },
365            Stage::Live,
366        );
367        tracker.operation(Operation::Read);
368
369        let failure = tracker
370            .build()
371            .expect_err("tracker should have failed to build event");
372
373        assert_eq!(failure, "no origins were invoked");
374    }
375
376    #[test]
377    fn fails_without_operation() {
378        let mut tracker = MigrationOpTracker::new(
379            "flag-key".into(),
380            Some(basic_flag("flag-key")),
381            ContextBuilder::new("user")
382                .build()
383                .expect("failed to build context"),
384            Detail {
385                value: Some(Stage::Live),
386                variation_index: Some(1),
387                reason: Reason::Fallthrough {
388                    in_experiment: false,
389                },
390            },
391            Stage::Live,
392        );
393        tracker.invoked(Origin::Old);
394        tracker.invoked(Origin::New);
395
396        let failure = tracker
397            .build()
398            .expect_err("tracker should have failed to build event");
399
400        assert_eq!(failure, "operation not provided");
401    }
402}