use super::rest::AlpacaRestError;
use crate::exchange::http::truncate_str;
use std::collections::HashSet;
use std::collections::hash_map::RandomState;
use std::hash::BuildHasher;
const TOKEN_PREFIX_BYTES: usize = 200;
#[derive(Debug)]
pub(super) struct PaginationGuard {
limit: usize,
pages: usize,
visited: HashSet<u64>,
hasher: RandomState,
}
impl PaginationGuard {
pub(super) fn new(limit: usize) -> Self {
Self {
limit,
pages: 0,
visited: HashSet::with_capacity(limit),
hasher: RandomState::new(),
}
}
pub(super) fn observe(&mut self, page_token: Option<&str>) -> Result<(), AlpacaRestError> {
self.pages += 1;
if self.pages > self.limit {
return Err(AlpacaRestError::PaginationLimitExceeded {
pages: self.pages,
limit: self.limit,
});
}
if let Some(token) = page_token
&& !self.visited.insert(self.hasher.hash_one(token))
{
return Err(AlpacaRestError::CyclicPagination {
page_token: truncate_str(token, TOKEN_PREFIX_BYTES),
});
}
Ok(())
}
pub(super) fn pages(&self) -> usize {
self.pages
}
}
#[cfg(test)]
#[allow(clippy::expect_used)] mod tests {
use super::*;
#[test]
fn distinct_tokens_under_the_cap_are_accepted() {
let mut guard = PaginationGuard::new(5);
guard.observe(None).expect("first page has no token");
for i in 1..5 {
guard
.observe(Some(&format!("token-{i}")))
.unwrap_or_else(|_| panic!("page {} within the cap should be accepted", i + 1));
}
}
#[test]
fn exceeding_the_page_cap_is_a_terminal_error() {
let mut guard = PaginationGuard::new(3);
guard.observe(None).expect("page 1 within cap");
guard.observe(Some("a")).expect("page 2 within cap");
guard.observe(Some("b")).expect("page 3 within cap");
let err = guard
.observe(Some("c"))
.expect_err("page limit + 1 must be rejected");
assert!(
matches!(
err,
AlpacaRestError::PaginationLimitExceeded { pages: 4, limit: 3 }
),
"expected PaginationLimitExceeded {{ pages: 4, limit: 3 }}, got {err:?}"
);
}
#[test]
fn reusing_a_token_is_detected_as_a_cycle() {
let mut guard = PaginationGuard::new(10);
guard.observe(None).expect("first page");
guard
.observe(Some("cursor-abc"))
.expect("first use is fine");
let err = guard
.observe(Some("cursor-abc"))
.expect_err("second use is a cycle");
assert!(
matches!(
&err,
AlpacaRestError::CyclicPagination { page_token } if page_token == "cursor-abc"
),
"expected CyclicPagination for cursor-abc, got {err:?}"
);
}
#[test]
fn cycle_detection_is_token_exact() {
let mut guard = PaginationGuard::new(10);
guard.observe(Some("cursor=a")).expect("page 1");
guard
.observe(Some("cursor=b"))
.expect("page 2 is a different token");
}
#[test]
fn tokenless_pages_never_cycle() {
let mut guard = PaginationGuard::new(10);
guard.observe(None).expect("page 1");
guard
.observe(None)
.expect("a second tokenless page is not a cycle");
}
#[test]
fn long_tokens_sharing_a_prefix_are_not_a_false_cycle() {
let mut guard = PaginationGuard::new(10);
let base = "a".repeat(2 * TOKEN_PREFIX_BYTES);
guard.observe(Some(&format!("{base}1"))).expect("page 1");
guard
.observe(Some(&format!("{base}2")))
.expect("page 2 differs only in its last byte but is a different token");
}
#[test]
fn a_cycling_oversized_token_is_reported_with_a_bounded_prefix() {
let mut guard = PaginationGuard::new(10);
let token = "a".repeat(2 * TOKEN_PREFIX_BYTES);
guard.observe(Some(&token)).expect("first use is fine");
let err = guard
.observe(Some(&token))
.expect_err("second use is a cycle");
match err {
AlpacaRestError::CyclicPagination { page_token } => {
assert_eq!(page_token.len(), TOKEN_PREFIX_BYTES);
assert!(token.starts_with(&page_token));
}
other => panic!("expected CyclicPagination, got {other:?}"),
}
}
}