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 {:?} without recording invocation",
144                    origin
145                ));
146            }
147
148            if self.latencies.contains_key(origin) {
149                return Err(format!(
150                    "provided latency for origin {:?} without recording invocation",
151                    origin
152                ));
153            }
154        }
155
156        if self.consistent.is_some() && self.invoked.len() != 2 {
157            return Err("provided consistency without recording both invocations".to_string());
158        }
159
160        Ok(())
161    }
162}
163
164#[cfg(test)]
165mod tests {
166
167    use launchdarkly_server_sdk_evaluation::{
168        ContextBuilder, Detail, Flag, MigrationFlagParameters, Reason,
169    };
170    use test_case::test_case;
171
172    use super::{MigrationOpTracker, Operation, Origin, Stage};
173    use crate::test_common::basic_flag;
174
175    fn minimal_tracker(flag: Flag) -> MigrationOpTracker {
176        let mut tracker = MigrationOpTracker::new(
177            flag.key.clone(),
178            Some(flag),
179            ContextBuilder::new("user")
180                .build()
181                .expect("failed to build context"),
182            Detail {
183                value: Some(Stage::Live),
184                variation_index: Some(1),
185                reason: Reason::Fallthrough {
186                    in_experiment: false,
187                },
188            },
189            Stage::Live,
190        );
191        tracker.operation(Operation::Read);
192        tracker.invoked(Origin::Old);
193        tracker.invoked(Origin::New);
194
195        tracker
196    }
197
198    #[test]
199    fn build_minimal_tracker() {
200        let tracker = minimal_tracker(basic_flag("flag-key"));
201        let result = tracker.build();
202
203        assert!(result.is_ok());
204    }
205
206    #[test]
207    fn build_without_flag() {
208        let mut tracker = minimal_tracker(basic_flag("flag-key"));
209        tracker.flag = None;
210        let result = tracker.build();
211
212        assert!(result.is_ok());
213    }
214
215    #[test_case(Origin::Old)]
216    #[test_case(Origin::New)]
217    fn track_invocations_individually(origin: Origin) {
218        let mut tracker = MigrationOpTracker::new(
219            "flag-key".into(),
220            Some(basic_flag("flag-key")),
221            ContextBuilder::new("user")
222                .build()
223                .expect("failed to build context"),
224            Detail {
225                value: Some(Stage::Live),
226                variation_index: Some(1),
227                reason: Reason::Fallthrough {
228                    in_experiment: false,
229                },
230            },
231            Stage::Live,
232        );
233        tracker.operation(Operation::Read);
234        tracker.invoked(origin);
235
236        let event = tracker.build().expect("failed to build event");
237        assert_eq!(event.invoked.len(), 1);
238        assert!(event.invoked.contains(&origin));
239    }
240
241    #[test]
242    fn tracks_both_invocations() {
243        let mut tracker = MigrationOpTracker::new(
244            "flag-key".into(),
245            Some(basic_flag("flag-key")),
246            ContextBuilder::new("user")
247                .build()
248                .expect("failed to build context"),
249            Detail {
250                value: Some(Stage::Live),
251                variation_index: Some(1),
252                reason: Reason::Fallthrough {
253                    in_experiment: false,
254                },
255            },
256            Stage::Live,
257        );
258        tracker.operation(Operation::Read);
259        tracker.invoked(Origin::Old);
260        tracker.invoked(Origin::New);
261
262        let event = tracker.build().expect("failed to build event");
263        assert_eq!(event.invoked.len(), 2);
264        assert!(event.invoked.contains(&Origin::Old));
265        assert!(event.invoked.contains(&Origin::New));
266    }
267
268    #[test_case(false)]
269    #[test_case(true)]
270    fn tracks_consistency(expectation: bool) {
271        let mut tracker = minimal_tracker(basic_flag("flag-key"));
272        tracker.operation(Operation::Read);
273        tracker.consistent(|| expectation);
274
275        let event = tracker.build().expect("failed to build event");
276        assert_eq!(event.consistency_check, Some(expectation));
277        assert_eq!(event.consistency_check_ratio, None);
278    }
279
280    #[test_case(false)]
281    #[test_case(true)]
282    fn consistency_can_be_disabled_through_sampling_ratio(expectation: bool) {
283        let mut flag = basic_flag("flag-key");
284        flag.migration_settings = Some(MigrationFlagParameters {
285            check_ratio: Some(0),
286        });
287
288        let mut tracker = minimal_tracker(flag);
289        tracker.operation(Operation::Read);
290        tracker.consistent(|| expectation);
291
292        let event = tracker.build().expect("failed to build event");
293        assert_eq!(event.consistency_check, None);
294        assert_eq!(event.consistency_check_ratio, Some(0));
295    }
296
297    #[test_case(Origin::Old)]
298    #[test_case(Origin::New)]
299    fn track_errors_individually(origin: Origin) {
300        let mut tracker = minimal_tracker(basic_flag("flag-key"));
301        tracker.error(origin);
302
303        let event = tracker.build().expect("failed to build event");
304        assert_eq!(event.errors.len(), 1);
305        assert!(event.errors.contains(&origin));
306    }
307
308    #[test]
309    fn tracks_both_errors() {
310        let mut tracker = minimal_tracker(basic_flag("flag-key"));
311        tracker.error(Origin::Old);
312        tracker.error(Origin::New);
313
314        let event = tracker.build().expect("failed to build event");
315        assert_eq!(event.errors.len(), 2);
316        assert!(event.errors.contains(&Origin::Old));
317        assert!(event.errors.contains(&Origin::New));
318    }
319
320    #[test_case(Origin::Old)]
321    #[test_case(Origin::New)]
322    fn track_latencies_individually(origin: Origin) {
323        let mut tracker = minimal_tracker(basic_flag("flag-key"));
324        tracker.latency(origin, std::time::Duration::from_millis(100));
325
326        let event = tracker.build().expect("failed to build event");
327        assert_eq!(event.latency.len(), 1);
328        assert_eq!(
329            event.latency.get(&origin),
330            Some(&std::time::Duration::from_millis(100))
331        );
332    }
333
334    #[test]
335    fn track_both_latencies() {
336        let mut tracker = minimal_tracker(basic_flag("flag-key"));
337        tracker.latency(Origin::Old, std::time::Duration::from_millis(100));
338        tracker.latency(Origin::New, std::time::Duration::from_millis(200));
339
340        let event = tracker.build().expect("failed to build event");
341        assert_eq!(event.latency.len(), 2);
342        assert_eq!(
343            event.latency.get(&Origin::Old),
344            Some(&std::time::Duration::from_millis(100))
345        );
346        assert_eq!(
347            event.latency.get(&Origin::New),
348            Some(&std::time::Duration::from_millis(200))
349        );
350    }
351
352    #[test]
353    fn fails_without_calling_invocations() {
354        let mut tracker = MigrationOpTracker::new(
355            "flag-key".into(),
356            Some(basic_flag("flag-key")),
357            ContextBuilder::new("user")
358                .build()
359                .expect("failed to build context"),
360            Detail {
361                value: Some(Stage::Live),
362                variation_index: Some(1),
363                reason: Reason::Fallthrough {
364                    in_experiment: false,
365                },
366            },
367            Stage::Live,
368        );
369        tracker.operation(Operation::Read);
370
371        let failure = tracker
372            .build()
373            .expect_err("tracker should have failed to build event");
374
375        assert_eq!(failure, "no origins were invoked");
376    }
377
378    #[test]
379    fn fails_without_operation() {
380        let mut tracker = MigrationOpTracker::new(
381            "flag-key".into(),
382            Some(basic_flag("flag-key")),
383            ContextBuilder::new("user")
384                .build()
385                .expect("failed to build context"),
386            Detail {
387                value: Some(Stage::Live),
388                variation_index: Some(1),
389                reason: Reason::Fallthrough {
390                    in_experiment: false,
391                },
392            },
393            Stage::Live,
394        );
395        tracker.invoked(Origin::Old);
396        tracker.invoked(Origin::New);
397
398        let failure = tracker
399            .build()
400            .expect_err("tracker should have failed to build event");
401
402        assert_eq!(failure, "operation not provided");
403    }
404}