use std::collections::BTreeSet;
use std::error::Error;
use std::fmt::{self, Display, Formatter};
use lodestone_wire::{
Cursor, ExportUsageRequest, ExportUsageResult, ListIndexesRequest, ListIndexesResult, SearchRequest, SearchResult,
};
use crate::config::parse_u64;
use crate::{CallOptions, Client, ClientError};
const MAX_PAGES: u32 = 10_000;
pub struct IndexPager {
client: Client,
request: ListIndexesRequest,
options: CallOptions,
seen_cursors: BTreeSet<String>,
pages: u32,
ended: bool,
}
pub struct SearchPager {
client: Client,
request: SearchRequest,
options: CallOptions,
seen_cursors: BTreeSet<String>,
pages: u32,
ended: bool,
}
pub struct UsagePager {
client: Client,
request: ExportUsageRequest,
options: CallOptions,
seen_cursors: BTreeSet<String>,
prior_cumulative_through: Option<u64>,
prior_sequence: Option<u64>,
prior_authority_revision: Option<u64>,
pages: u32,
ended: bool,
}
impl Client {
pub fn index_pager(&self, request: ListIndexesRequest, options: CallOptions) -> Result<IndexPager, PagerError> {
validate_page_size(request.page_size)?;
Ok(IndexPager {
client: self.clone(),
seen_cursors: cursor_history(&request.cursor)?,
request,
options,
pages: 0,
ended: false,
})
}
pub fn search_pager(&self, request: SearchRequest, options: CallOptions) -> Result<SearchPager, PagerError> {
validate_page_size(request.page_size)?;
let valid_binding = match (&request.indexes, &request.pit) {
(Some(indexes), None) => !indexes.is_empty(),
(None, Some(_)) => true,
_ => false,
};
if !valid_binding {
return Err(PagerError::InvalidBinding);
}
Ok(SearchPager {
client: self.clone(),
seen_cursors: cursor_history(&request.cursor)?,
request,
options,
pages: 0,
ended: false,
})
}
pub fn usage_pager(&self, request: ExportUsageRequest, options: CallOptions) -> Result<UsagePager, PagerError> {
validate_page_size(request.page_size)?;
let start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
let end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
if start >= end {
return Err(PagerError::InvalidBinding);
}
Ok(UsagePager {
client: self.clone(),
seen_cursors: cursor_history(&request.cursor)?,
request,
options,
prior_cumulative_through: None,
prior_sequence: None,
prior_authority_revision: None,
pages: 0,
ended: false,
})
}
}
impl IndexPager {
pub async fn next_page(&mut self) -> Result<Option<ListIndexesResult>, PagerError> {
if self.ended {
return Ok(None);
}
consume_budget(&mut self.pages)?;
let result = self
.client
.list_indexes(self.request.clone(), self.options.clone())
.await
.map_err(PagerError::Client)?;
if result.restart_count != self.request.restart_count {
self.ended = true;
return Err(PagerError::StaleCursor);
}
advance_cursor(
&mut self.request.cursor,
&mut self.seen_cursors,
result.next_cursor.clone(),
&mut self.ended,
)?;
Ok(Some(result))
}
}
impl SearchPager {
pub async fn next_page(&mut self) -> Result<Option<SearchResult>, PagerError> {
if self.ended {
return Ok(None);
}
consume_budget(&mut self.pages)?;
let result = self
.client
.search(self.request.clone(), self.options.clone())
.await
.map_err(PagerError::Client)?;
advance_cursor(
&mut self.request.cursor,
&mut self.seen_cursors,
result.next_cursor.clone(),
&mut self.ended,
)?;
Ok(Some(result))
}
}
impl UsagePager {
pub async fn next_page(&mut self) -> Result<Option<ExportUsageResult>, PagerError> {
if self.ended {
return Ok(None);
}
consume_budget(&mut self.pages)?;
let result = self
.client
.export_usage(self.request.clone(), self.options.clone())
.await
.map_err(PagerError::Client)?;
if let Err(error) = validate_usage_page(
&self.request,
&result,
&mut self.prior_cumulative_through,
&mut self.prior_sequence,
&mut self.prior_authority_revision,
) {
self.ended = true;
return Err(error);
}
advance_cursor(
&mut self.request.cursor,
&mut self.seen_cursors,
result.next_cursor.clone(),
&mut self.ended,
)?;
Ok(Some(result))
}
}
fn validate_usage_page(
request: &ExportUsageRequest,
result: &ExportUsageResult,
prior_cumulative_through: &mut Option<u64>,
prior_sequence: &mut Option<u64>,
prior_authority_revision: &mut Option<u64>,
) -> Result<(), PagerError> {
let request_start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
let request_end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
let cumulative = parse_u64(&result.cumulative_through.value).ok_or(PagerError::InvalidCorePage)?;
if prior_cumulative_through.is_some_and(|prior| cumulative < prior) {
return Err(PagerError::AuthorityRegression);
}
let mut rows = Vec::with_capacity(result.units.len().saturating_add(result.corrections.len()));
for row in &result.units {
rows.push(validate_usage_row(
&row.sequence.value,
&row.window_start.value,
&row.window_end.value,
&row.authority_revision.value,
&row.source_digest,
request_start,
request_end,
)?);
}
for row in &result.corrections {
if parse_u64(&row.correction_sequence.value).is_none_or(|value| value == 0) {
return Err(PagerError::InvalidCorePage);
}
rows.push(validate_usage_row(
&row.sequence.value,
&row.window_start.value,
&row.window_end.value,
&row.authority_revision.value,
&row.source_digest,
request_start,
request_end,
)?);
}
rows.sort_unstable_by_key(|row| row.0);
for (sequence, authority_revision) in rows {
if prior_sequence.is_some_and(|prior| sequence <= prior)
|| prior_authority_revision.is_some_and(|prior| authority_revision != prior)
{
return Err(PagerError::AuthorityRegression);
}
*prior_sequence = Some(sequence);
*prior_authority_revision = Some(authority_revision);
}
if prior_sequence.is_some_and(|sequence| cumulative < sequence)
|| result.next_cursor.is_some() && result.units.is_empty() && result.corrections.is_empty()
{
return Err(PagerError::InvalidCorePage);
}
*prior_cumulative_through = Some(cumulative);
Ok(())
}
fn validate_usage_row(
sequence: &str,
window_start: &str,
window_end: &str,
authority_revision: &str,
source_digest: &lodestone_wire::Digest,
request_start: u64,
request_end: u64,
) -> Result<(u64, u64), PagerError> {
let sequence = parse_u64(sequence).ok_or(PagerError::InvalidCorePage)?;
let window_start = parse_u64(window_start).ok_or(PagerError::InvalidCorePage)?;
let window_end = parse_u64(window_end).ok_or(PagerError::InvalidCorePage)?;
let authority_revision = parse_u64(authority_revision).ok_or(PagerError::InvalidCorePage)?;
if sequence == 0
|| authority_revision == 0
|| window_start < request_start
|| window_end > request_end
|| window_start >= window_end
|| source_digest.validate().is_err()
{
return Err(PagerError::InvalidCorePage);
}
Ok((sequence, authority_revision))
}
fn validate_page_size(page_size: u32) -> Result<(), PagerError> {
if page_size == 0 || page_size > 10_000 {
Err(PagerError::InvalidBinding)
} else {
Ok(())
}
}
fn consume_budget(pages: &mut u32) -> Result<(), PagerError> {
if *pages >= MAX_PAGES {
return Err(PagerError::PageBudgetExhausted);
}
*pages = pages.saturating_add(1);
Ok(())
}
fn advance_cursor(
request_cursor: &mut Option<Cursor>,
seen_cursors: &mut BTreeSet<String>,
next_cursor: Option<Cursor>,
ended: &mut bool,
) -> Result<(), PagerError> {
if let Some(next) = &next_cursor {
if !valid_cursor(next) || !seen_cursors.insert(next.value.clone()) {
*ended = true;
return Err(PagerError::CursorCycle);
}
} else {
*ended = true;
}
*request_cursor = next_cursor;
Ok(())
}
fn cursor_history(cursor: &Option<Cursor>) -> Result<BTreeSet<String>, PagerError> {
let mut seen = BTreeSet::new();
if let Some(cursor) = cursor {
if !valid_cursor(cursor) {
return Err(PagerError::InvalidBinding);
}
let _ = seen.insert(cursor.value.clone());
}
Ok(seen)
}
fn valid_cursor(cursor: &Cursor) -> bool {
!cursor.value.is_empty()
&& cursor.value.len() <= 16 * 1024
&& !cursor.value.bytes().any(|byte| byte.is_ascii_control())
}
#[derive(Debug)]
pub enum PagerError {
InvalidBinding,
Client(ClientError),
StaleCursor,
CursorCycle,
AuthorityRegression,
InvalidCorePage,
PageBudgetExhausted,
}
impl Display for PagerError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::InvalidBinding => "pager filter or page binding is invalid",
Self::Client(_) => "Core page request failed",
Self::StaleCursor => "Core catalog cursor became stale",
Self::CursorCycle => "Core pagination cursor repeated or is invalid",
Self::AuthorityRegression => "Core paged authority watermark regressed",
Self::InvalidCorePage => "Core page violates its generated semantic contract",
Self::PageBudgetExhausted => "Core pager reached its bounded page budget",
})
}
}
impl Error for PagerError {}
#[cfg(test)]
mod tests {
use super::*;
fn cursor(value: &str) -> Cursor {
Cursor {
value: value.to_owned(),
}
}
#[test]
fn cursor_history_refuses_non_adjacent_cycles() -> Result<(), PagerError> {
let mut request = None;
let mut seen = BTreeSet::new();
let mut ended = false;
advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended)?;
advance_cursor(&mut request, &mut seen, Some(cursor("b")), &mut ended)?;
assert!(matches!(
advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended),
Err(PagerError::CursorCycle)
));
assert!(ended);
Ok(())
}
#[test]
fn initial_cursor_is_validated_and_participates_in_cycle_detection() -> Result<(), PagerError> {
assert!(matches!(
cursor_history(&Some(cursor("bad\nvalue"))),
Err(PagerError::InvalidBinding)
));
let mut request = Some(cursor("initial"));
let mut seen = cursor_history(&request)?;
let mut ended = false;
assert!(matches!(
advance_cursor(&mut request, &mut seen, Some(cursor("initial")), &mut ended),
Err(PagerError::CursorCycle)
));
Ok(())
}
#[test]
fn usage_sequence_watermark_is_not_compared_to_time_window() -> Result<(), PagerError> {
let request = ExportUsageRequest {
window_start: lodestone_wire::U64 { value: "10".to_owned() },
window_end: lodestone_wire::U64 { value: "20".to_owned() },
page_size: 10,
cursor: None,
};
let result = ExportUsageResult {
units: Vec::new(),
corrections: Vec::new(),
next_cursor: None,
cumulative_through: lodestone_wire::U64 {
value: "1000000".to_owned(),
},
};
validate_usage_page(&request, &result, &mut None, &mut None, &mut None)?;
Ok(())
}
}