Skip to main content

snap_control/pg_wap/
segments.rs

1// Copyright 2026 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Segment management.
16//!
17//! [`SegmentManager`] keeps the public segments of every (src, dst) pair that is in use available
18//! and fresh, and evicts the pairs that stop being used.
19
20use std::{
21    collections::HashMap,
22    sync::{
23        Arc,
24        atomic::{AtomicU64, AtomicUsize, Ordering},
25    },
26    time::{Duration, SystemTime, UNIX_EPOCH},
27};
28
29use anyhow::Context;
30use endhost_api_models::SegmentsDiscovery;
31use scc::hash_map::OccupiedEntry;
32use sciparse::{
33    identifier::isd_asn::IsdAsn,
34    segment::{SegmentFp, Segments, SignedPathSegment},
35};
36use serde::{Deserialize, Serialize};
37use serde_with::{DurationSeconds, serde_as};
38use tokio::select;
39
40/// Number of segments to request per segment fetch.
41const SEGMENT_FETCH_PAGE_SIZE: i32 = 250;
42
43// Note: This can be improved a lot, instead of storing all segments per src, dst, we can split
44// up CORE and NON_CORE segments, and manage them separately. This would deduplicate
45// segments, and allow us to do smaller more precise fetches.
46
47/// Keeps the public segments of every (src, dst) pair that is in use available and fresh.
48///
49/// A pair starts being managed the first time it is asked for, is refetched before its
50/// segments expire, and is evicted once it has not been used for
51/// [`SegmentManagerConfig::idle_eviction_time`].
52///
53/// A consumer that keeps using a pair longer than that must hold a [`PairGuard`] from
54/// [`Self::hold_pair`] on it, which exempts it from eviction.
55///
56/// Cheap to clone; all clones share the same state.
57#[derive(Clone)]
58pub struct SegmentManager(Arc<SegmentManagerInner>);
59
60/// The state shared by all clones of a [`SegmentManager`].
61struct SegmentManagerInner {
62    segments: scc::HashMap<(IsdAsn, IsdAsn), SegmentStore>,
63
64    segment_fetcher: Box<dyn SegmentsDiscovery>,
65
66    config: SegmentManagerConfig,
67
68    next_maintain: tokio::sync::watch::Sender<SystemTime>,
69}
70
71/// Configuration for a [`SegmentManager`].
72#[serde_as]
73#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
74#[serde(deny_unknown_fields)]
75pub struct SegmentManagerConfig {
76    /// The minimum duration between two fetches of segments for the same (src, dst) pair.
77    #[serde_as(as = "DurationSeconds<u64>")]
78    #[serde(default = "SegmentManagerConfig::default_minimum_segment_fetch_interval")]
79    pub minimum_segment_fetch_interval: Duration,
80    /// The maximum duration between two fetches of segments for the same (src, dst) pair.
81    #[serde_as(as = "DurationSeconds<u64>")]
82    #[serde(default = "SegmentManagerConfig::default_maximum_segment_fetch_interval")]
83    pub maximum_segment_fetch_interval: Duration,
84    /// The duration after which an unused (src, dst) pair is removed from the store.
85    #[serde_as(as = "DurationSeconds<u64>")]
86    #[serde(default = "SegmentManagerConfig::default_idle_eviction_time")]
87    pub idle_eviction_time: Duration,
88    /// Minimum duration a segment must have left to be kept in the store.
89    #[serde_as(as = "DurationSeconds<u64>")]
90    #[serde(default = "SegmentManagerConfig::default_min_segment_lifetime")]
91    pub min_segment_lifetime: Duration,
92    /// The duration before the segments expiry is reached to start fetching new segments.
93    #[serde_as(as = "DurationSeconds<u64>")]
94    #[serde(default = "SegmentManagerConfig::default_segment_lifetime_buffer")]
95    pub segment_lifetime_buffer: Duration,
96}
97
98impl SegmentManagerConfig {
99    fn default_minimum_segment_fetch_interval() -> Duration {
100        // With the default segment lifetime buffer, we have 8 chances to refetch a segment before
101        // it expires. This should be enough to avoid hitting issues.
102        Duration::from_secs(30)
103    }
104
105    fn default_maximum_segment_fetch_interval() -> Duration {
106        Duration::from_secs(60 * 60)
107    }
108
109    fn default_idle_eviction_time() -> Duration {
110        Duration::from_secs(60)
111    }
112
113    fn default_min_segment_lifetime() -> Duration {
114        Duration::from_secs(60)
115    }
116
117    fn default_segment_lifetime_buffer() -> Duration {
118        Duration::from_secs(5 * 60)
119    }
120
121    /// Checks that the timings are usable.
122    ///
123    /// Returns an error if a duration is zero, or if the minimum fetch interval is greater than
124    /// the maximum fetch interval.
125    pub fn validate(&self) -> anyhow::Result<()> {
126        if self.minimum_segment_fetch_interval > self.maximum_segment_fetch_interval {
127            anyhow::bail!(
128                "minimum_segment_fetch_interval ({:?}) is greater than \
129                 maximum_segment_fetch_interval ({:?})",
130                self.minimum_segment_fetch_interval,
131                self.maximum_segment_fetch_interval
132            );
133        }
134
135        if self.minimum_segment_fetch_interval == Duration::ZERO {
136            anyhow::bail!("minimum_segment_fetch_interval must be greater than 0");
137        }
138
139        if self.maximum_segment_fetch_interval == Duration::ZERO {
140            anyhow::bail!("maximum_segment_fetch_interval must be greater than 0");
141        }
142
143        if self.idle_eviction_time == Duration::ZERO {
144            anyhow::bail!("idle_eviction_time must be greater than 0");
145        }
146
147        Ok(())
148    }
149}
150
151impl Default for SegmentManagerConfig {
152    fn default() -> Self {
153        Self {
154            minimum_segment_fetch_interval: Self::default_minimum_segment_fetch_interval(),
155            maximum_segment_fetch_interval: Self::default_maximum_segment_fetch_interval(),
156            idle_eviction_time: Self::default_idle_eviction_time(),
157            min_segment_lifetime: Self::default_min_segment_lifetime(),
158            segment_lifetime_buffer: Self::default_segment_lifetime_buffer(),
159        }
160    }
161}
162
163impl SegmentManager {
164    /// Creates a new, empty segment manager with the timings of `config`.
165    ///
166    /// Returns an error if `config` does not [`SegmentManagerConfig::validate`].
167    pub fn new(
168        config: SegmentManagerConfig,
169        segment_fetcher: Box<dyn SegmentsDiscovery>,
170    ) -> anyhow::Result<Self> {
171        config.validate()?;
172
173        Ok(Self(Arc::new(SegmentManagerInner {
174            segments: scc::HashMap::new(),
175            segment_fetcher,
176            config,
177            next_maintain: tokio::sync::watch::channel(SystemTime::now()).0,
178        })))
179    }
180
181    /// Returns the segment with the given fingerprint, if it is in the store.
182    ///
183    /// Counts as a use of the segment's (src, dst) pair.
184    pub async fn segment(&self, id: SegmentStoreId, now: SystemTime) -> Option<SignedPathSegment> {
185        let store = self.0.segments.get_async(&(id.src(), id.dst())).await?;
186        store.get().touch(now);
187
188        let segments = match id {
189            SegmentStoreId::Core { .. } => &store.get().core_segments,
190            SegmentStoreId::NonCore { .. } => &store.get().non_core_segments,
191        };
192
193        segments
194            .get(&id.fp())
195            .map(|segment_entry| segment_entry.segment.clone())
196    }
197
198    /// Keeps the given (src, dst) pair managed for as long as the returned guard is held,
199    /// fetching its segments if it is not managed yet.
200    ///
201    /// A guarded pair is exempt from idle eviction, so a long lived consumer cannot have the
202    /// segments it depends on evicted out from under it. Refreshing is unaffected: a guarded
203    /// pair keeps being refetched as its segments approach expiry.
204    pub async fn hold_pair(
205        &self,
206        src: IsdAsn,
207        dst: IsdAsn,
208        now: SystemTime,
209    ) -> anyhow::Result<PairGuard> {
210        let store = self.ensure_managed(src, dst, now).await?;
211        let ref_count = store.get().acquire();
212
213        tracing::trace!(%src, %dst, "Holding a pair against eviction");
214
215        Ok(PairGuard {
216            ref_count,
217            src,
218            dst,
219        })
220    }
221
222    /// Returns a [`SegmentsIter`] for the given (src, dst) pair, fetching segments if
223    /// necessary.
224    ///
225    /// Allows iterating over all core and non-core segments for the given (src, dst) pair.
226    pub async fn segments<'this>(
227        &'this self,
228        src: IsdAsn,
229        dst: IsdAsn,
230        now: SystemTime,
231    ) -> anyhow::Result<SegmentsIter<'this>> {
232        let store = self.ensure_managed(src, dst, now).await?;
233        store.get().touch(now);
234
235        Ok(SegmentsIter { store })
236    }
237
238    async fn wait_to_next_maintain(&self) {
239        let mut rx = self.0.next_maintain.subscribe();
240
241        loop {
242            let next = *rx.borrow_and_update();
243            let duration = next
244                .duration_since(SystemTime::now())
245                .unwrap_or(Duration::ZERO);
246
247            tracing::trace!("Next segment manager maintenance in {:?}", duration);
248
249            select! {
250                // Due time was updated, check again.
251                _ = rx.changed() => {
252                    tracing::trace!("Next segment manager maintenance time changed");
253                },
254                // Wait until the next maintenance time
255                _ = tokio::time::sleep(duration) => return,
256            }
257        }
258    }
259
260    /// Runs the maintenance loop for the segment manager, which periodically evicts unused
261    /// (src, dst) pairs and refetches the ones that are due for a refresh.
262    ///
263    /// Reads the wall clock itself, unlike every other entry point. Never returns.
264    pub async fn run(&self) {
265        loop {
266            let now = SystemTime::now();
267
268            // Reset the next maintenance time to the maximum fetch interval from now
269            self.0
270                .next_maintain
271                .send_replace(now + self.0.config.maximum_segment_fetch_interval);
272            let next = self.maintain(now).await;
273            // Update the next maintenance time with the earlier time
274            self.update_next_maintain(next);
275
276            // Wait until the next maintenance time, adapting to changes in the next maintenance
277            // time that happen while we wait.
278            self.wait_to_next_maintain().await;
279        }
280    }
281
282    /// Updates the next maintenance time to `next`, if it is earlier than the current value.
283    fn update_next_maintain(&self, next: SystemTime) {
284        self.0.next_maintain.send_if_modified(|current| {
285            if next < *current {
286                *current = next;
287                true
288            } else {
289                false
290            }
291        });
292    }
293
294    /// Ensures that the segment store for the given (src, dst) pair is managed, fetching
295    /// segments if necessary.
296    async fn ensure_managed(
297        &self,
298        src: IsdAsn,
299        dst: IsdAsn,
300        now: SystemTime,
301    ) -> anyhow::Result<OccupiedEntry<'_, (IsdAsn, IsdAsn), SegmentStore>> {
302        // Try to get the store for the (src, dst) pair. If it exists, touch it and return it.
303        if let Some(entry) = self.0.segments.get_async(&(src, dst)).await {
304            entry.get().touch(now);
305            return Ok(entry);
306        }
307
308        // Otherwise, fetch the segments for the (src, dst) pair and create a new store.
309        tracing::info!(%src, %dst, "Fetching segments for a new (src, dst) pair");
310
311        // XXX: This is racy, multiple callers can fetch the same (src, dst) pair concurrently. We
312        // could add a fetch-in-progress lock to avoid this. ATM we just throw out extra fetches.
313        // TODO: Paging is not supported yet.
314        let fetched = self
315            .0
316            .segment_fetcher
317            .list_segments(src, dst, SEGMENT_FETCH_PAGE_SIZE, String::new())
318            .await
319            .with_context(|| format!("Failed to fetch segments for ({src}, {dst})"))?;
320
321        // Prepare the store for the (src, dst) pair with the fetched segments.
322        let mut store = SegmentStore::new((src, dst), now);
323        store.update_segments(fetched.segments, now, &self.0.config);
324
325        match self.0.segments.entry_async((src, dst)).await {
326            // Someone else already created a store for the (src, dst) pair, return it and discard
327            // the one we just created.
328            scc::hash_map::Entry::Occupied(occupied_entry) => Ok(occupied_entry),
329            // No entry exists for the (src, dst) pair, fetch segments and create a new store.
330            scc::hash_map::Entry::Vacant(vacant_entry) => {
331                // Update the next maintenance time with the new store's next refresh time, if it is
332                // earlier than the current value.
333                self.update_next_maintain(store.next_refresh);
334                Ok(vacant_entry.insert_entry(store))
335            }
336        }
337    }
338
339    /// Maintains the segment store by removing unused entries and fetching new segments for
340    /// (src, dst) pairs that are due for refresh.
341    ///
342    /// Returns the next scheduled maintenance time, which is the earliest next refresh time of
343    /// all (src, dst) pairs in the store.
344    ///
345    /// If no (src, dst) pairs are due for refresh, the next maintenance time is the maximum
346    /// segment fetch interval from now.
347    pub async fn maintain(&self, now: SystemTime) -> SystemTime {
348        let mut update_keys = Vec::new();
349        let mut next_maintain = now + self.0.config.maximum_segment_fetch_interval;
350
351        // Evict all unused entries from the store and plan to update segments for (src, dst)
352        // pairs that are due for refresh.
353        self.0
354            .segments
355            .retain_async(|key, entry| {
356                // A referenced pair is still in use, however long ago it was last looked at.
357                if !entry.has_references()
358                    && entry.last_use() + self.0.config.idle_eviction_time < now
359                {
360                    tracing::debug!(
361                        src = %key.0, dst = %key.1,
362                        "Evicting unused segments from the store"
363                    );
364                    return false;
365                }
366
367                // Plan to update segments for the (src, dst) pair if the next refresh is due.
368                if entry.next_refresh < now {
369                    update_keys.push(*key);
370                } else {
371                    // If we don't update now, track the next refresh time for scheduling the
372                    // next maintenance run.
373                    next_maintain = next_maintain.min(entry.next_refresh);
374                }
375
376                true
377            })
378            .await;
379
380        for (src, dst) in update_keys {
381            // Fetch new segments for the (src, dst) pair and update the store.
382            // TODO: Paging is not supported yet.
383            let fetched = self
384                .0
385                .segment_fetcher
386                .list_segments(src, dst, SEGMENT_FETCH_PAGE_SIZE, String::new())
387                .await;
388
389            let fetched = match fetched {
390                Ok(fetched) => fetched,
391                Err(e) => {
392                    // The entry stays due for a refresh; retry no earlier than the minimum
393                    // fetch interval from now.
394                    next_maintain =
395                        next_maintain.min(now + self.0.config.minimum_segment_fetch_interval);
396                    tracing::error!(%src, %dst, "Failed to fetch segments: {e}");
397                    continue;
398                }
399            };
400
401            let Some(mut entry) = self.0.segments.get_async(&(src, dst)).await else {
402                tracing::debug!(
403                    %src, %dst,
404                    "Segments were evicted from the store before they could be updated"
405                );
406                continue;
407            };
408
409            entry
410                .get_mut()
411                .update_segments(fetched.segments, now, &self.0.config);
412
413            next_maintain = next_maintain.min(entry.get().next_refresh);
414        }
415
416        next_maintain
417    }
418}
419
420/// A store of segments for a given (src, dst) pair, with expiration and last-use timestamps.
421struct SegmentStore {
422    /// The query used to fetch the segments.
423    query: (IsdAsn, IsdAsn),
424    /// When the next refresh of the segments is due, based on the earliest expiration of the
425    /// segments.
426    next_refresh: SystemTime,
427    /// When the segments were last used, as seconds since the UNIX epoch.
428    last_use: AtomicU64,
429    /// Number of live [`PairGuard`]s on this pair. If this is greater than zero, the pair is
430    /// exempt from idle eviction.
431    refs: Arc<AtomicUsize>,
432    /// All core segments for the (src ISD, dst ISD) pair, indexed by their fingerprint.
433    core_segments: HashMap<SegmentFp, SegmentEntry>,
434    /// All non-core segments for the (src, dst) pair, indexed by their fingerprint.
435    non_core_segments: HashMap<SegmentFp, SegmentEntry>,
436}
437
438struct SegmentEntry {
439    segment: SignedPathSegment,
440    expiration: SystemTime,
441}
442
443impl SegmentStore {
444    /// Creates a new empty segment store for the given (src, dst) pair which needs to be
445    /// updated with new segments.
446    fn new(query: (IsdAsn, IsdAsn), now: SystemTime) -> Self {
447        let this = Self {
448            query,
449            next_refresh: now,
450            last_use: AtomicU64::new(0),
451            refs: Arc::new(AtomicUsize::new(0)),
452            core_segments: HashMap::new(),
453            non_core_segments: HashMap::new(),
454        };
455        this.touch(now);
456        this
457    }
458
459    /// Takes one more reference on this pair, returning the count to be held by its guard.
460    fn acquire(&self) -> Arc<AtomicUsize> {
461        self.refs.fetch_add(1, Ordering::Relaxed);
462        self.refs.clone()
463    }
464
465    /// Whether anything holds a [`PairGuard`] on this pair.
466    fn has_references(&self) -> bool {
467        self.refs.load(Ordering::Relaxed) > 0
468    }
469
470    /// Records that the store was used at `now`.
471    fn touch(&self, now: SystemTime) {
472        let secs = now
473            .duration_since(UNIX_EPOCH)
474            .unwrap_or(Duration::ZERO)
475            .as_secs();
476
477        self.last_use.fetch_max(secs, Ordering::Relaxed);
478    }
479
480    /// When the store was last used. A store counts as used when it is created.
481    fn last_use(&self) -> SystemTime {
482        UNIX_EPOCH + Duration::from_secs(self.last_use.load(Ordering::Relaxed))
483    }
484
485    /// Updates the segments store.
486    ///
487    /// Does the following steps:
488    /// - Merges `new_segments` into the store
489    /// - Removes segments that have less than [`SegmentManagerConfig::min_segment_lifetime`] left
490    /// - Updates the next refresh time based on the earliest expiration of the remaining segments,
491    ///   or the maximum fetch interval, whichever is earlier
492    ///
493    /// The next refresh time is between now + the minimum fetch interval and now + the maximum
494    /// fetch interval, trying to capture the earliest expiration of the segments.
495    fn update_segments(
496        &mut self,
497        new_segments: Segments,
498        now: SystemTime,
499        config: &SegmentManagerConfig,
500    ) {
501        let start_core_count = self.core_segments.len();
502        let start_non_core_count = self.non_core_segments.len();
503
504        let (new_core, new_non_core) = new_segments.split_parts();
505
506        for (segments, store) in [
507            (new_core, &mut self.core_segments),
508            (new_non_core, &mut self.non_core_segments),
509        ] {
510            for segment in segments {
511                // We don't compare the new segment with the existing one, we just replace it.
512                // New segments should always have the same or later expiration than the existing
513                // ones.
514                let expiration = segment.expires_earliest();
515                store.insert(
516                    segment.fingerprint(),
517                    SegmentEntry {
518                        segment,
519                        expiration,
520                    },
521                );
522            }
523        }
524
525        // Never refetch later than max_refetch_delay from now, even if all segments live
526        // longer than that.
527        let mut next_refetch = now + config.maximum_segment_fetch_interval;
528        for store in [&mut self.core_segments, &mut self.non_core_segments] {
529            store.retain(|_, segment| {
530                // Drop segments that are expired, or so close to expiry that a path built
531                // from them would be useless.
532                if segment.expiration < now + config.min_segment_lifetime {
533                    return false;
534                }
535
536                // Try to have fetched a replacement before the earliest segment expires.
537                let refetch_at = segment
538                    .expiration
539                    .checked_sub(config.segment_lifetime_buffer)
540                    .unwrap_or(now);
541
542                next_refetch = next_refetch.min(refetch_at);
543
544                true
545            });
546        }
547
548        // If we have neither core nor non-core segments, force a early refetch
549        if self.core_segments.is_empty() && self.non_core_segments.is_empty() {
550            tracing::debug!(
551                src = %self.query.0, dst = %self.query.1,
552                "Segment store is empty after update, forcing a refetch in {}s",
553                config.minimum_segment_fetch_interval.as_secs()
554            );
555
556            next_refetch = now + config.minimum_segment_fetch_interval;
557        }
558
559        // Ensure that the next refresh time is at least now + the minimum fetch interval, to
560        // avoid fetching too frequently.
561        next_refetch = next_refetch.max(now + config.minimum_segment_fetch_interval);
562
563        tracing::info!(
564            "Updated segments for ({}, {}): core: {} -> {}, non-core: {} -> {}, next refresh: {:?}",
565            self.query.0,
566            self.query.1,
567            start_core_count,
568            self.core_segments.len(),
569            start_non_core_count,
570            self.non_core_segments.len(),
571            next_refetch
572        );
573
574        self.next_refresh = next_refetch;
575    }
576}
577
578/// Keeps the segments of one (src, dst) pair managed for as long as it is held.
579///
580/// Handed out by [`SegmentManager::hold_pair`]. Dropping it makes the pair subject to idle
581/// eviction again.
582pub struct PairGuard {
583    /// The reference count of the guarded pair, shared with its entry in the store.
584    ref_count: Arc<AtomicUsize>,
585    src: IsdAsn,
586    dst: IsdAsn,
587}
588
589impl Drop for PairGuard {
590    fn drop(&mut self) {
591        let previous = self.ref_count.fetch_sub(1, Ordering::Relaxed);
592        debug_assert!(
593            previous > 0,
594            "released a pair guard that was never acquired"
595        );
596
597        tracing::trace!(src = %self.src, dst = %self.dst, "Released a pair guard");
598    }
599}
600
601/// A unique identifier for a segment in the store, based on the (src, dst) pair it was fetched
602/// for and its fingerprint.
603#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
604pub enum SegmentStoreId {
605    /// A core segment.
606    Core {
607        /// The source of the query the segment was fetched for.
608        src: IsdAsn,
609        /// The destination of the query the segment was fetched for.
610        dst: IsdAsn,
611        /// The fingerprint of the segment.
612        fp: SegmentFp,
613    },
614    /// An up or down segment.
615    NonCore {
616        /// The source of the query the segment was fetched for.
617        src: IsdAsn,
618        /// The destination of the query the segment was fetched for.
619        dst: IsdAsn,
620        /// The fingerprint of the segment.
621        fp: SegmentFp,
622    },
623}
624
625impl SegmentStoreId {
626    /// The source of the query the segment was fetched for.
627    pub fn src(&self) -> IsdAsn {
628        match self {
629            Self::Core { src, .. } | Self::NonCore { src, .. } => *src,
630        }
631    }
632
633    /// The destination of the query the segment was fetched for.
634    pub fn dst(&self) -> IsdAsn {
635        match self {
636            Self::Core { dst, .. } | Self::NonCore { dst, .. } => *dst,
637        }
638    }
639
640    /// The fingerprint of the segment.
641    pub fn fp(&self) -> SegmentFp {
642        match self {
643            Self::Core { fp, .. } | Self::NonCore { fp, .. } => *fp,
644        }
645    }
646
647    /// Returns true if this identifies a core segment.
648    pub fn is_core(&self) -> bool {
649        matches!(self, Self::Core { .. })
650    }
651}
652
653/// Borrowed view of the segments of one (src, dst) pair.
654///
655/// This holds a lock on the store entry, so it must not be held across an `.await` that can
656/// touch the [`SegmentManager`] again.
657pub struct SegmentsIter<'store> {
658    store: OccupiedEntry<'store, (IsdAsn, IsdAsn), SegmentStore>,
659}
660
661impl SegmentsIter<'_> {
662    /// Iterates over the core segments of the pair.
663    pub fn iter_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
664        self.store
665            .get()
666            .core_segments
667            .values()
668            .map(|entry| &entry.segment)
669    }
670
671    /// Iterates over the up and down segments of the pair.
672    pub fn iter_non_core_segments(&self) -> impl Iterator<Item = &SignedPathSegment> {
673        self.store
674            .get()
675            .non_core_segments
676            .values()
677            .map(|entry| &entry.segment)
678    }
679
680    /// Returns whether the pair holds a core segment with the given fingerprint.
681    pub fn has_core_segment(&self, fp: SegmentFp) -> bool {
682        self.store.get().core_segments.contains_key(&fp)
683    }
684
685    /// Returns whether the pair holds an up or down segment with the given fingerprint.
686    pub fn has_non_core_segment(&self, fp: SegmentFp) -> bool {
687        self.store.get().non_core_segments.contains_key(&fp)
688    }
689}
690
691#[cfg(test)]
692mod tests {
693    use super::*;
694    use crate::pg_wap::test_util::{
695        Fixture, IDLE_EVICTION_TIME, MAX_FETCH_INTERVAL, MIN_FETCH_INTERVAL, MIN_SEGMENT_LIFETIME,
696        MockFetcher, SEGMENT_LIFETIME_BUFFER, at, core_ia, leaf_ia, other_leaf_ia,
697        secs_since_epoch, short_lived_up_segment, store_id, up_segment,
698    };
699
700    #[tokio::test]
701    async fn segments_are_fetched_once_per_pair_and_queryable_by_fingerprint() {
702        let up = up_segment(0);
703        let fixture = Fixture::new(
704            MockFetcher::with_up_segments(vec![up.clone()]),
705            Duration::from_secs(100),
706        );
707
708        let fingerprints = {
709            let segments = fixture
710                .segments
711                .segments(leaf_ia(), core_ia(), at(0))
712                .await
713                .expect("segments are fetched");
714            segments
715                .iter_non_core_segments()
716                .map(SignedPathSegment::fingerprint)
717                .collect::<Vec<_>>()
718        };
719        assert_eq!(fingerprints, vec![up.fingerprint()]);
720        assert_eq!(fixture.fetcher.calls(), 1);
721
722        // A second request for the same pair is served from the store.
723        fixture
724            .segments
725            .segments(leaf_ia(), core_ia(), at(0))
726            .await
727            .expect("segments are served from the store");
728        assert_eq!(fixture.fetcher.calls(), 1);
729
730        let id = store_id(&up);
731        assert!(fixture.segments.segment(id, at(0)).await.is_some());
732        assert!(
733            fixture
734                .segments
735                .segment(
736                    SegmentStoreId::NonCore {
737                        src: leaf_ia(),
738                        dst: core_ia(),
739                        fp: SegmentFp::default(),
740                    },
741                    at(0)
742                )
743                .await
744                .is_none(),
745            "an unknown fingerprint is not found"
746        );
747    }
748
749    #[tokio::test]
750    async fn segments_close_to_expiry_are_not_stored() {
751        let up = up_segment(0);
752        let fixture = Fixture::new(
753            MockFetcher::with_up_segments(vec![up.clone()]),
754            Duration::from_secs(100),
755        );
756
757        // The fixture keeps segments with at least 60s left; leave 30s.
758        let nearly_expired = at(secs_since_epoch(up.expires_earliest()) - 30);
759
760        let segments = fixture
761            .segments
762            .segments(leaf_ia(), core_ia(), nearly_expired)
763            .await
764            .expect("segments are fetched");
765        assert_eq!(segments.iter_non_core_segments().count(), 0);
766    }
767
768    #[tokio::test]
769    async fn maintenance_is_scheduled_no_later_than_the_fetch_interval_cap() {
770        let fixture = Fixture::new(
771            // Lives for 24h, i.e. far longer than the cap.
772            MockFetcher::with_up_segments(vec![up_segment(0)]),
773            Duration::from_secs(100),
774        );
775
776        assert_eq!(
777            fixture.segments.maintain(at(0)).await,
778            at(0) + MAX_FETCH_INTERVAL,
779            "with nothing managed there is nothing to schedule around"
780        );
781
782        fixture
783            .segments
784            .segments(leaf_ia(), core_ia(), at(0))
785            .await
786            .expect("segments are fetched");
787
788        assert_eq!(
789            fixture.segments.maintain(at(0)).await,
790            at(0) + MAX_FETCH_INTERVAL,
791            "a long lived pair is still refetched at the cap"
792        );
793        assert_eq!(
794            fixture.fetcher.calls(),
795            1,
796            "nothing was due, so nothing was refetched"
797        );
798    }
799
800    #[tokio::test]
801    async fn maintenance_is_scheduled_a_buffer_before_the_earliest_expiry() {
802        // Lives ~28min, so its replacement is due well inside the fetch interval cap.
803        let up = short_lived_up_segment(0, 4);
804        let fixture = Fixture::new(
805            MockFetcher::with_up_segments(vec![up.clone()]),
806            Duration::from_secs(100),
807        );
808
809        fixture
810            .segments
811            .segments(leaf_ia(), core_ia(), at(0))
812            .await
813            .expect("segments are fetched");
814
815        assert_eq!(
816            fixture.segments.maintain(at(0)).await,
817            up.expires_earliest() - SEGMENT_LIFETIME_BUFFER,
818            "the replacement has to be fetched before the segment expires"
819        );
820    }
821
822    #[tokio::test]
823    async fn maintenance_is_scheduled_for_the_earliest_of_all_pairs() {
824        let short = short_lived_up_segment(0, 4);
825        let fixture = Fixture::new(
826            MockFetcher::with_up_segments(vec![up_segment(0)]),
827            Duration::from_secs(100),
828        );
829
830        fixture
831            .segments
832            .segments(leaf_ia(), core_ia(), at(0))
833            .await
834            .expect("the long lived pair is fetched");
835        fixture.fetcher.set_segments(Segments {
836            up_segments: vec![short.clone()],
837            ..Segments::default()
838        });
839        fixture
840            .segments
841            .segments(leaf_ia(), other_leaf_ia(), at(0))
842            .await
843            .expect("the short lived pair is fetched");
844
845        assert_eq!(
846            fixture.segments.maintain(at(0)).await,
847            short.expires_earliest() - SEGMENT_LIFETIME_BUFFER,
848            "the pair that expires first decides when maintenance runs again"
849        );
850    }
851
852    #[tokio::test]
853    async fn a_pair_without_segments_is_retried_at_the_minimum_interval() {
854        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
855
856        fixture
857            .segments
858            .segments(leaf_ia(), core_ia(), at(0))
859            .await
860            .expect("the empty pair is managed");
861        assert_eq!(fixture.fetcher.calls(), 1);
862
863        assert_eq!(
864            fixture.segments.maintain(at(0)).await,
865            at(0) + MIN_FETCH_INTERVAL,
866            "a pair we know nothing about is retried as soon as we are allowed to"
867        );
868
869        // And that retry actually happens, rather than the pair sitting there empty.
870        fixture
871            .segments
872            .maintain(at(MIN_FETCH_INTERVAL.as_secs() + 1))
873            .await;
874        assert_eq!(fixture.fetcher.calls(), 2);
875    }
876
877    #[tokio::test]
878    async fn maintenance_reports_the_schedule_of_the_pairs_it_just_refetched() {
879        let fixture = Fixture::new(
880            MockFetcher::with_up_segments(vec![up_segment(0)]),
881            Duration::from_secs(100),
882        );
883
884        // Held, so the pair survives long enough to become due for a refetch.
885        let _guard = fixture
886            .segments
887            .hold_pair(leaf_ia(), core_ia(), at(0))
888            .await
889            .expect("the pair is fetched and held");
890        assert_eq!(fixture.fetcher.calls(), 1);
891
892        let due_at = MAX_FETCH_INTERVAL.as_secs() + 1;
893
894        assert_eq!(
895            fixture.segments.maintain(at(due_at)).await,
896            at(due_at) + MAX_FETCH_INTERVAL,
897            "the returned time must be the refetched pair's new schedule, not the one it was due at"
898        );
899        assert_eq!(fixture.fetcher.calls(), 2, "the pair was due and refetched");
900    }
901
902    #[tokio::test]
903    async fn a_failed_refetch_is_retried_at_the_minimum_interval() {
904        let fixture = Fixture::new(
905            MockFetcher::with_up_segments(vec![up_segment(0)]),
906            Duration::from_secs(100),
907        );
908
909        let _guard = fixture
910            .segments
911            .hold_pair(leaf_ia(), core_ia(), at(0))
912            .await
913            .expect("the pair is fetched and held");
914
915        let due_at = MAX_FETCH_INTERVAL.as_secs() + 1;
916        fixture.fetcher.set_failing(true);
917
918        assert_eq!(
919            fixture.segments.maintain(at(due_at)).await,
920            at(due_at) + MIN_FETCH_INTERVAL,
921            "a failed fetch is retried as soon as we are allowed to"
922        );
923        assert_eq!(
924            fixture.fetcher.calls(),
925            2,
926            "the failing fetch was attempted"
927        );
928
929        // The pair stayed due, so the next run tries again instead of waiting for the cap.
930        fixture.fetcher.set_failing(false);
931        let retried_at = at(due_at + MIN_FETCH_INTERVAL.as_secs() + 1);
932        assert_eq!(
933            fixture.segments.maintain(retried_at).await,
934            retried_at + MAX_FETCH_INTERVAL
935        );
936        assert_eq!(fixture.fetcher.calls(), 3);
937        assert!(
938            fixture
939                .segments
940                .segment(store_id(&up_segment(0)), retried_at)
941                .await
942                .is_some(),
943            "the retry restored the segments of the pair"
944        );
945    }
946
947    #[tokio::test]
948    async fn unused_pairs_are_evicted_and_used_ones_kept() {
949        let fixture = Fixture::new(
950            MockFetcher::with_up_segments(vec![up_segment(0)]),
951            Duration::from_secs(100),
952        );
953
954        fixture
955            .segments
956            .segments(leaf_ia(), core_ia(), at(0))
957            .await
958            .expect("segments are fetched");
959        assert_eq!(fixture.fetcher.calls(), 1);
960
961        // The fixture evicts after 60s of no use.
962        fixture.segments.maintain(at(30)).await;
963        fixture
964            .segments
965            .segments(leaf_ia(), core_ia(), at(30))
966            .await
967            .expect("segments are still in the store");
968        assert_eq!(fixture.fetcher.calls(), 1, "the pair was still fresh");
969
970        // The last use is the request above, at t=30.
971        fixture.segments.maintain(at(91)).await;
972        fixture
973            .segments
974            .segments(leaf_ia(), core_ia(), at(91))
975            .await
976            .expect("segments are fetched again");
977        assert_eq!(
978            fixture.fetcher.calls(),
979            2,
980            "the evicted pair has to be fetched again"
981        );
982    }
983    #[tokio::test]
984    async fn the_wake_up_only_moves_earlier() {
985        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
986        let next_maintain = &fixture.segments.0.next_maintain;
987
988        next_maintain.send_replace(at(1000));
989
990        fixture.segments.update_next_maintain(at(500));
991        assert_eq!(*next_maintain.borrow(), at(500), "an earlier wake up wins");
992
993        fixture.segments.update_next_maintain(at(900));
994        assert_eq!(
995            *next_maintain.borrow(),
996            at(500),
997            "a later wake up does not push the earlier one back"
998        );
999    }
1000
1001    #[tokio::test]
1002    async fn a_new_pair_moves_the_wake_up_to_its_own_refresh() {
1003        // An empty fetch leaves the store empty, so the pair is due again after the minimum
1004        // fetch interval.
1005        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1006        let next_maintain = &fixture.segments.0.next_maintain;
1007
1008        // What the maintenance loop leaves behind after maintaining an empty manager.
1009        next_maintain.send_replace(at(0) + MAX_FETCH_INTERVAL);
1010
1011        fixture
1012            .segments
1013            .segments(leaf_ia(), core_ia(), at(0))
1014            .await
1015            .expect("segments are fetched");
1016        assert_eq!(
1017            *next_maintain.borrow(),
1018            at(0) + MIN_FETCH_INTERVAL,
1019            "a pair added after the last maintenance run has to be waited for"
1020        );
1021
1022        next_maintain.send_replace(at(0) + MAX_FETCH_INTERVAL);
1023        fixture
1024            .segments
1025            .segments(leaf_ia(), core_ia(), at(0))
1026            .await
1027            .expect("segments are served from the store");
1028        assert_eq!(
1029            *next_maintain.borrow(),
1030            at(0) + MAX_FETCH_INTERVAL,
1031            "a pair that is already managed is already scheduled for"
1032        );
1033    }
1034
1035    // Test using `start_paused` run in virtual time, once all tasks have yielded, sleeps advance
1036    // the clock to the next scheduled wake up.
1037    #[tokio::test(start_paused = true)]
1038    async fn the_wait_ends_at_the_scheduled_wake_up() {
1039        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1040        fixture
1041            .segments
1042            .0
1043            .next_maintain
1044            .send_replace(SystemTime::now() + Duration::from_secs(30));
1045
1046        let started = tokio::time::Instant::now();
1047        fixture.segments.wait_to_next_maintain().await;
1048
1049        let waited = started.elapsed();
1050        assert!(
1051            waited > Duration::from_secs(29) && waited <= Duration::from_secs(31),
1052            "waited {waited:?} instead of the scheduled 30s"
1053        );
1054    }
1055
1056    // Test using `start_paused` run in virtual time, once all tasks have yielded, sleeps advance
1057    // the clock to the next scheduled wake up.
1058    #[tokio::test(start_paused = true)]
1059    async fn an_earlier_wake_up_cuts_the_wait_short() {
1060        let fixture = Fixture::new(MockFetcher::empty(), Duration::from_secs(100));
1061        fixture
1062            .segments
1063            .0
1064            .next_maintain
1065            .send_replace(SystemTime::now() + MAX_FETCH_INTERVAL);
1066
1067        let started = tokio::time::Instant::now();
1068        let segments = fixture.segments.clone();
1069        let waiting = tokio::spawn(async move { segments.wait_to_next_maintain().await });
1070
1071        // Let the waiting task arm its timer on the far away wake up before moving it.
1072        tokio::time::sleep(Duration::from_secs(10)).await;
1073        fixture
1074            .segments
1075            .update_next_maintain(SystemTime::now() + Duration::from_secs(5));
1076
1077        tokio::time::timeout(MAX_FETCH_INTERVAL / 2, waiting)
1078            .await
1079            .expect("the wait ends before the wake up it started with")
1080            .expect("the waiting task does not panic");
1081
1082        let waited = started.elapsed();
1083
1084        assert!(
1085            waited < Duration::from_secs(20),
1086            "waited {waited:?} instead of following the earlier wake up"
1087        );
1088    }
1089
1090    /// The maintenance loop reads the wall clock, which a test cannot move, so this one runs on
1091    /// real time with a manager configured in milliseconds.
1092    #[tokio::test]
1093    async fn the_maintenance_loop_wakes_for_a_pair_added_while_it_waits() {
1094        let fetcher = MockFetcher::empty();
1095        let segments = SegmentManager::new(
1096            SegmentManagerConfig {
1097                minimum_segment_fetch_interval: Duration::from_millis(50),
1098                // Far longer than the test runs, so reaching it fails the test.
1099                maximum_segment_fetch_interval: Duration::from_secs(60),
1100                idle_eviction_time: IDLE_EVICTION_TIME,
1101                min_segment_lifetime: MIN_SEGMENT_LIFETIME,
1102                segment_lifetime_buffer: SEGMENT_LIFETIME_BUFFER,
1103            },
1104            Box::new(fetcher.clone()),
1105        )
1106        .expect("a valid SegmentManagerConfig");
1107
1108        let maintained = segments.clone();
1109        tokio::spawn(async move { maintained.run().await });
1110
1111        // The loop maintains the empty manager and settles on the fetch interval cap.
1112        tokio::time::sleep(Duration::from_millis(20)).await;
1113        assert_eq!(fetcher.calls(), 0, "nothing asked for a pair yet");
1114
1115        segments
1116            .segments(leaf_ia(), core_ia(), SystemTime::now())
1117            .await
1118            .expect("segments are fetched");
1119        assert_eq!(fetcher.calls(), 1);
1120
1121        // The empty store is due again after the minimum fetch interval, so the loop has to
1122        // refetch long before the fetch interval cap.
1123        let mut calls = fetcher.calls();
1124        for _ in 0..50 {
1125            if calls > 1 {
1126                break;
1127            }
1128            tokio::time::sleep(Duration::from_millis(100)).await;
1129            calls = fetcher.calls();
1130        }
1131        assert!(
1132            calls > 1,
1133            "the loop kept waiting instead of maintaining the pair added while it waited"
1134        );
1135    }
1136}