use crate::error::{LaserError, decode_managed_reply};
use crate::filters::group::{ASSIGNMENT_REFRESH, Membership};
use crate::filters::guard::LocalGuard;
use crate::filters::progress::{Blocked, PartitionProgress};
use crate::filters::route::{RouteTo, Routes};
use crate::iggy::prelude::{
Client, IggyClient, IggyClientBuilder, IggyError, IggyMessage, PolledMessages, StreamClient,
TopicClient,
};
use crate::laser::Laser;
use bytes::Bytes;
use iggy_binary_protocol::batch::{BatchIntegrity, decode_batch_slice_with};
use iggy_common::Identifier;
use laser_wire::codes::{AGDX_FILTERED_ACK_CODE, AGDX_FILTERED_POLL_CODE, FILTER_OP_VERSION};
use laser_wire::filter::{
AppliedPolicy, ConsumerFilter, FilterConsumer, FilterError, FilterErrorReason, FilterOutcome,
FilterRef, FilterReply, FilterSource, FilteredAck, FilteredPage, FilteredPollRequest,
FilteredStart, ReadMode, SourceGeneration, StopReason,
};
use laser_wire::framing::encode_named;
use laser_wire::limits::MAX_FILTER_CATALOG_PAGE;
use laser_wire::schema::Digest32;
use laser_wire::validate::Validate;
use serde::de::DeserializeOwned;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::{Instant, sleep};
const DEFAULT_COUNT: u32 = 100;
const DEFAULT_MAX_REPLY_BYTES: u32 = 1024 * 1024;
const DEFAULT_IDLE_INTERVAL: Duration = Duration::from_millis(250);
const DEFAULT_MAX_UNACKED_PAGES: usize = 1024;
pub struct FilteredReaderBuilder<'a> {
laser: &'a Laser,
source: FilterSource,
consumer: Option<FilterConsumer>,
partitions: Vec<u32>,
filter: Option<FilterRef>,
start: FilteredStart,
count: u32,
max_reply_bytes: u32,
read_mode: ReadMode,
local_guard: bool,
idle_interval: Duration,
max_unacked_pages: usize,
}
pub struct FilteredReader {
owner: Arc<()>,
laser: Laser,
coordinator: Arc<IggyClient>,
consumer: FilterConsumer,
request: FilteredPollRequest,
inline_digest: Option<Digest32>,
idle_interval: Duration,
max_unacked_pages: usize,
guard: Option<LocalGuard>,
guard_enabled: bool,
membership: Option<Membership>,
listed_at: Option<Instant>,
partitions: BTreeMap<u32, PartitionProgress>,
revoked: BTreeSet<u32>,
routes: Routes,
retired_connections: u64,
cursor: usize,
sequence: u64,
examined_in_round: u64,
lost_routes: BTreeSet<u32>,
buffered: VecDeque<MatchedRecord>,
deferred_error: Option<LaserError>,
pinned_source: Option<SourceIncarnation>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct SourceIncarnation {
stream_id: u32,
stream_created_at_micros: u64,
topic_id: u32,
topic_created_at_micros: u64,
}
impl SourceIncarnation {
pub(crate) async fn read(
client: &IggyClient,
source: &FilterSource,
) -> Result<Self, LaserError> {
let stream = client
.get_stream(&Identifier::named(&source.stream)?)
.await?
.ok_or_else(|| {
FilterError::new(FilterErrorReason::NotFound, "source stream does not exist")
})?;
let topic = client
.get_topic(
&Identifier::numeric(stream.id)?,
&Identifier::named(&source.topic)?,
)
.await?
.ok_or_else(|| {
FilterError::new(FilterErrorReason::NotFound, "source topic does not exist")
})?;
Ok(Self {
stream_id: stream.id,
stream_created_at_micros: stream.created_at.as_micros(),
topic_id: topic.id,
topic_created_at_micros: topic.created_at.as_micros(),
})
}
}
impl From<SourceGeneration> for SourceIncarnation {
fn from(generation: SourceGeneration) -> Self {
Self {
stream_id: generation.stream_id,
stream_created_at_micros: generation.stream_created_at_micros,
topic_id: generation.topic_id,
topic_created_at_micros: generation.topic_created_at_micros,
}
}
}
#[derive(Debug)]
pub struct MatchedPage {
pub partition_id: u32,
pub records: Vec<MatchedRecord>,
pub policy: AppliedPolicy,
pub generation: SourceGeneration,
pub stop: StopReason,
pub examined: u32,
pub frontier: u64,
pub safe_ack_offset: Option<u64>,
sequence: u64,
owner: Arc<()>,
}
#[derive(Debug)]
pub struct MatchedRecord {
pub partition_id: u32,
pub offset: u64,
pub frontier: u64,
pub evaluated: bool,
pub message: IggyMessage,
sequence: u64,
owner: Arc<()>,
}
impl MatchedRecord {
pub fn json<T: DeserializeOwned>(&self) -> Result<T, LaserError> {
serde_json::from_slice(&self.message.payload)
.map_err(|error| LaserError::Codec(error.to_string()))
}
pub fn headers_malformed(&self) -> bool {
self.message.user_headers_map().is_err()
}
}
enum Read {
Page(MatchedPage),
Empty {
more: bool,
},
Saturated,
}
enum Round {
Page(MatchedPage),
Busy,
Idle,
}
impl<'a> FilteredReaderBuilder<'a> {
pub(crate) fn new(laser: &'a Laser, source: FilterSource) -> Self {
Self {
laser,
source,
consumer: None,
partitions: Vec::new(),
filter: None,
start: FilteredStart::Next,
count: DEFAULT_COUNT,
max_reply_bytes: DEFAULT_MAX_REPLY_BYTES,
read_mode: ReadMode::Primary,
local_guard: false,
idle_interval: DEFAULT_IDLE_INTERVAL,
max_unacked_pages: DEFAULT_MAX_UNACKED_PAGES,
}
}
#[must_use]
pub fn consumer(mut self, name: impl Into<String>) -> Self {
self.consumer = Some(FilterConsumer::Consumer(name.into()));
self
}
#[must_use]
pub fn group(mut self, name: impl Into<String>) -> Self {
self.consumer = Some(FilterConsumer::Group(name.into()));
self
}
#[must_use]
pub fn group_id(mut self, id: u64) -> Self {
self.consumer = Some(FilterConsumer::GroupId(id));
self
}
#[must_use]
pub fn partition(mut self, partition_id: u32) -> Self {
self.partitions.push(partition_id);
self
}
#[must_use]
pub fn filter(mut self, filter: FilterRef) -> Self {
self.filter = Some(filter);
self
}
#[must_use]
pub fn inline(self, filter: ConsumerFilter) -> Self {
self.filter(FilterRef::Inline(filter))
}
#[must_use]
pub fn revision(self, filter_id: u32, revision: u32) -> Self {
self.filter(FilterRef::Revision {
filter_id,
revision,
})
}
#[must_use]
pub fn start(mut self, start: FilteredStart) -> Self {
self.start = start;
self
}
#[must_use]
pub fn count(mut self, count: u32) -> Self {
self.count = count;
self
}
#[must_use]
pub fn max_reply_bytes(mut self, max_reply_bytes: u32) -> Self {
self.max_reply_bytes = max_reply_bytes;
self
}
#[must_use]
pub fn read_mode(mut self, read_mode: ReadMode) -> Self {
self.read_mode = read_mode;
self
}
#[must_use]
pub fn local_guard(mut self, enabled: bool) -> Self {
self.local_guard = enabled;
self
}
#[must_use]
pub fn idle_interval(mut self, interval: Duration) -> Self {
self.idle_interval = interval;
self
}
#[must_use]
pub fn max_unacked_pages(mut self, pages: usize) -> Self {
self.max_unacked_pages = pages;
self
}
pub async fn build(self) -> Result<FilteredReader, LaserError> {
if self.max_unacked_pages == 0 {
return Err(LaserError::Config("max_unacked_pages must be positive"));
}
let capabilities = self.laser.capabilities().await;
if !capabilities.filters.native {
return Err(LaserError::unsupported(
"filters",
"consumer filters are not served by this server",
));
}
if let Some(FilterRef::Inline(inline)) = &self.filter
&& !capabilities
.filters
.evaluates(inline.evaluator_version, inline.codec)
{
return Err(LaserError::unsupported(
"filters",
format!(
"this server does not evaluate {:?} filters of evaluator version {}",
inline.codec, inline.evaluator_version
),
));
}
if self.read_mode == ReadMode::Primary && self.laser.connection_string().is_none() {
return Err(LaserError::Config(
"primary filtered reads open data connections, so they need a Laser built from a connection string. Read in ReadMode::Local for a bring-your-own client",
));
}
let consumer = self.consumer.ok_or(LaserError::Config(
"a filtered reader needs .consumer(name), .group(name), or .group_id(id)",
))?;
let filter = match (self.filter, &consumer) {
(Some(filter), _) => filter,
(None, FilterConsumer::Group(_) | FilterConsumer::GroupId(_)) => FilterRef::Bound,
(None, FilterConsumer::Consumer(_)) => {
return Err(LaserError::Config(
"an independent filtered reader needs .inline(..), .revision(..), or .filter(..)",
));
}
};
let request = FilteredPollRequest {
v: FILTER_OP_VERSION,
source: self.source,
partition_id: 0,
consumer: consumer.clone(),
filter,
start: self.start.clone(),
count: self.count,
max_reply_bytes: self.max_reply_bytes,
read_mode: self.read_mode,
};
request.validate()?;
let inline_digest = match &request.filter {
FilterRef::Inline(filter) => Some(filter.digest()),
_ => None,
};
let guard = if self.local_guard && !matches!(request.filter, FilterRef::Bound) {
let definition = definition(self.laser, &request.filter).await?;
Some(LocalGuard::load(self.laser, &definition).await?)
} else {
None
};
let pinned_source = if matches!(consumer, FilterConsumer::GroupId(_)) {
Some(SourceIncarnation::read(&self.laser.client(), &request.source).await?)
} else {
None
};
let coordinator = match (&consumer, self.laser.connection_string()) {
(FilterConsumer::Group(_) | FilterConsumer::GroupId(_), Some(connection_string)) => {
let client =
IggyClientBuilder::from_connection_string(connection_string)?.build()?;
client.connect().await?;
Arc::new(client)
}
_ => self.laser.client(),
};
let (membership, listed_at, partition_ids) = match &consumer {
FilterConsumer::Group(_) | FilterConsumer::GroupId(_) => {
if !self.partitions.is_empty() {
return Err(LaserError::Config(
"a group reader reads the partitions the group assigns it",
));
}
let membership =
Membership::join(&coordinator, &request.source, &consumer, pinned_source)
.await?;
let assigned = membership.partitions().to_vec();
(Some(membership), None, assigned)
}
FilterConsumer::Consumer(_) if self.partitions.is_empty() => (
None,
Some(Instant::now()),
topic_partitions(self.laser, &request.source).await?,
),
FilterConsumer::Consumer(_) => (None, None, self.partitions),
};
let partitions = partition_ids
.into_iter()
.map(|partition_id| (partition_id, PartitionProgress::new(self.start.clone())))
.collect();
Ok(FilteredReader {
owner: Arc::new(()),
laser: self.laser.clone(),
coordinator,
consumer,
request,
inline_digest,
idle_interval: self.idle_interval,
max_unacked_pages: self.max_unacked_pages,
guard,
guard_enabled: self.local_guard,
membership,
listed_at,
partitions,
revoked: BTreeSet::new(),
routes: Routes::default(),
retired_connections: 0,
cursor: 0,
sequence: 0,
examined_in_round: 0,
lost_routes: BTreeSet::new(),
buffered: VecDeque::new(),
deferred_error: None,
pinned_source,
})
}
}
impl FilteredReader {
pub async fn next_page(&mut self) -> Result<MatchedPage, LaserError> {
loop {
match self.round().await? {
Round::Page(page) => return Ok(page),
Round::Busy => {}
Round::Idle => sleep(self.idle_interval).await,
}
}
}
pub async fn try_next_page(&mut self) -> Result<Option<MatchedPage>, LaserError> {
loop {
match self.round().await? {
Round::Page(page) => return Ok(Some(page)),
Round::Busy => {}
Round::Idle => return Ok(None),
}
}
}
pub async fn read_round(&mut self) -> Result<(Option<MatchedPage>, bool), LaserError> {
Ok(match self.round().await? {
Round::Page(page) => (Some(page), false),
Round::Busy => (None, true),
Round::Idle => (None, false),
})
}
pub const fn examined_in_round(&self) -> u64 {
self.examined_in_round
}
pub const fn idle_interval(&self) -> Duration {
self.idle_interval
}
#[doc(hidden)]
pub fn coordinator(&self) -> Arc<IggyClient> {
Arc::clone(&self.coordinator)
}
pub fn owns(&self, page: &MatchedPage) -> bool {
Arc::ptr_eq(&self.owner, &page.owner)
&& self
.partitions
.get(&page.partition_id)
.is_some_and(|progress| progress.knows(page.sequence))
}
pub async fn next_record(&mut self) -> Result<MatchedRecord, LaserError> {
loop {
if let Some(record) = self.buffered.pop_front() {
return Ok(record);
}
let page = self.next_page().await?;
self.buffered.extend(page.records);
}
}
pub async fn ack(&mut self, record: &MatchedRecord) -> Result<(), LaserError> {
self.require_acknowledgments()?;
self.require_owner(&record.owner)?;
self.require_sequence(record.partition_id, record.sequence)?;
self.progress(record.partition_id)?
.complete_record(record.sequence, record.offset);
self.flush(record.partition_id).await
}
pub async fn ack_through(&mut self, record: &MatchedRecord) -> Result<(), LaserError> {
self.require_acknowledgments()?;
self.require_owner(&record.owner)?;
self.require_sequence(record.partition_id, record.sequence)?;
if !self
.progress(record.partition_id)?
.complete_through(record.sequence, record.offset)
{
return Err(LaserError::Invalid(
"the record is past the page acknowledgment boundary".to_owned(),
));
}
self.flush(record.partition_id).await
}
pub async fn ack_page(&mut self, page: &MatchedPage) -> Result<(), LaserError> {
self.require_acknowledgments()?;
self.require_owner(&page.owner)?;
self.require_sequence(page.partition_id, page.sequence)?;
self.progress(page.partition_id)?
.complete_page(page.sequence);
self.flush(page.partition_id).await
}
pub fn data_connections_opened(&self) -> u64 {
self.retired_connections + self.routes.opened()
}
pub fn partitions(&self) -> Vec<u32> {
self.partitions
.keys()
.filter(|partition_id| !self.revoked.contains(partition_id))
.copied()
.collect()
}
#[doc(hidden)]
pub async fn flush_completed(&mut self) -> Result<(), LaserError> {
let mut outcome = Ok(());
if self.request.read_mode == ReadMode::Primary {
for partition_id in self.partitions.keys().copied().collect::<Vec<_>>() {
let flushed = self.flush(partition_id).await;
outcome = outcome.and(flushed);
}
}
outcome
}
pub async fn close(mut self) -> Result<(), LaserError> {
let mut outcome = self.flush_completed().await;
if let Some(membership) = &self.membership {
outcome = outcome.and(membership.leave(&self.coordinator).await);
}
std::mem::take(&mut self.routes).close().await;
if self.membership.is_some() && self.laser.connection_string().is_some() {
let _ = self.coordinator.shutdown().await;
}
outcome
}
async fn round(&mut self) -> Result<Round, LaserError> {
self.examined_in_round = 0;
if let Some(error) = self.deferred_error.take() {
return Err(error);
}
if self.membership.as_ref().is_some_and(Membership::is_due) {
self.refresh_membership().await?;
}
if self
.listed_at
.is_some_and(|listed_at| listed_at.elapsed() >= ASSIGNMENT_REFRESH)
{
self.refresh_partitions().await?;
}
let readable = self.partitions();
if readable.is_empty() {
return Ok(Round::Idle);
}
let now = Instant::now();
let mut busy = false;
let mut failed = None;
let mut saturated = Vec::new();
for step in 0..readable.len() {
let index = (self.cursor + step) % readable.len();
let partition_id = readable[index];
if self.progress(partition_id)?.waiting(now) {
continue;
}
match self.read(partition_id).await {
Ok(Read::Page(page)) => {
self.cursor = index + 1;
self.deferred_error = failed;
return Ok(Round::Page(page));
}
Ok(Read::Empty { more }) => busy |= more,
Ok(Read::Saturated) => saturated.push(partition_id),
Err(error) => {
let retry_at = now + self.idle_interval;
if let Ok(progress) = self.progress(partition_id) {
progress.defer(retry_at);
}
failed.get_or_insert(error);
}
}
}
if let Some(error) = failed {
return Err(error);
}
if !saturated.is_empty() && !busy {
return Err(LaserError::Invalid(format!(
"partitions {saturated:?} each have {} unacknowledged pages. Acknowledge before reading more",
self.max_unacked_pages
)));
}
Ok(if busy { Round::Busy } else { Round::Idle })
}
async fn read(&mut self, partition_id: u32) -> Result<Read, LaserError> {
let maximum = self.max_unacked_pages;
let progress = self.progress(partition_id)?;
if progress.in_flight() >= maximum {
return Ok(Read::Saturated);
}
let start = progress.start().clone();
self.request.partition_id = partition_id;
self.request.start = start;
let payload = encode_named(&self.request)
.map_err(|error| LaserError::Codec(format!("encode filter request: {error}")))?;
let outcome = self
.exchange(partition_id, AGDX_FILTERED_POLL_CODE, payload)
.await;
let mut page = match outcome {
Ok(FilterOutcome::Page(page)) => {
self.lost_routes.remove(&partition_id);
page
}
Ok(_) => return Err(unexpected("filtered poll")),
Err(error) if is_route_lost(&error) => {
if !self.lost_routes.insert(partition_id) {
return Err(error);
}
if let Some(membership) = self.membership.as_mut() {
membership.expire();
}
return Ok(Read::Empty { more: false });
}
Err(error) => return Err(self.restart_after(partition_id, error)),
};
if page.records.len() > self.request.max_reply_bytes as usize + 16 {
return Err(LaserError::Protocol(
"filtered page exceeds the requested record byte limit".to_owned(),
));
}
validate_record_body(&self.request, &page)?;
let polled = PolledMessages::from_bytes(Bytes::from(std::mem::take(&mut page.records)))?;
validate_page(&self.request, self.inline_digest.as_ref(), &page, &polled)?;
self.pin_source(&page)?;
let messages = polled.messages;
if self.guard_enabled && self.guard.is_none() {
self.guard = Some(
LocalGuard::load(
&self.laser,
&bound_definition(&self.laser, &page, &self.request.source).await?,
)
.await?,
);
}
if let Some(guard) = &self.guard
&& let Err(error) = guard.check(&page, &messages)
{
return Err(self.restart_after(partition_id, error));
}
self.sequence += 1;
let sequence = self.sequence;
let primary = self.request.read_mode == ReadMode::Primary;
let retry_at = Instant::now() + self.idle_interval;
let progress = self.progress(partition_id)?;
progress.record(
sequence,
&page,
messages
.iter()
.filter(|_| primary)
.map(|message| message.header.offset),
);
let blocked = progress.blocked();
if blocked.is_some() {
progress.defer(retry_at);
}
self.examined_in_round += u64::from(page.examined);
if messages.is_empty() {
if primary && let Err(error) = self.flush(partition_id).await {
tracing::debug!(target: "laser", partition_id, %error, "storing filtered progress failed, retried with the next acknowledgment");
}
if let Some(blocked) = blocked {
return Err(blocked_error(partition_id, blocked));
}
return Ok(Read::Empty {
more: matches!(page.stop, StopReason::Filled | StopReason::Budget),
});
}
let unevaluated: BTreeSet<u64> = page.unevaluated.iter().copied().collect();
let records = messages
.into_iter()
.map(|message| MatchedRecord {
partition_id,
offset: message.header.offset,
frontier: page.frontier,
evaluated: !unevaluated.contains(&message.header.offset),
message,
sequence,
owner: Arc::clone(&self.owner),
})
.collect();
Ok(Read::Page(MatchedPage {
partition_id,
records,
policy: page.policy,
generation: page.generation,
stop: page.stop,
examined: page.examined,
frontier: page.frontier,
safe_ack_offset: page.safe_ack_offset,
sequence,
owner: Arc::clone(&self.owner),
}))
}
fn restart_after(&mut self, partition_id: u32, error: LaserError) -> LaserError {
let reason = error.filter_reason();
if !matches!(
reason,
Some(FilterErrorReason::SourceChanged | FilterErrorReason::Conflict)
) {
return error;
}
if let Ok(progress) = self.progress(partition_id) {
*progress = PartitionProgress::new(FilteredStart::Next);
}
self.buffered
.retain(|record| record.partition_id != partition_id);
if reason == Some(FilterErrorReason::SourceChanged)
&& let Some(membership) = self.membership.as_mut()
{
membership.expire();
}
if matches!(self.request.filter, FilterRef::Bound) {
self.guard = None;
}
error
}
async fn flush(&mut self, partition_id: u32) -> Result<(), LaserError> {
let Some(target) = self.progress(partition_id)?.unstored().cloned() else {
self.retire(partition_id);
return Ok(());
};
let ack = FilteredAck {
group_id: target.group_id,
v: FILTER_OP_VERSION,
source: self.request.source.clone(),
partition_id,
consumer: self.consumer.clone(),
generation: target.generation,
digest: target.digest.clone(),
offset: target.offset,
};
let payload = encode_named(&ack)
.map_err(|error| LaserError::Codec(format!("encode filter request: {error}")))?;
let mut rerouted = false;
loop {
match self
.exchange(partition_id, AGDX_FILTERED_ACK_CODE, payload.clone())
.await
{
Ok(FilterOutcome::Acknowledged(receipt)) => {
if receipt.partition_id != partition_id
|| receipt.offset != target.offset
|| receipt.generation != target.generation
{
return Err(LaserError::Protocol(
"filtered acknowledgment does not match its request".to_owned(),
));
}
self.progress(partition_id)?.stored(&target);
let revoked = self.revoked.contains(&partition_id);
self.retire(partition_id);
if revoked {
self.routes.release(partition_id).await;
}
return Ok(());
}
Ok(_) => return Err(unexpected("filtered acknowledgment")),
Err(error) if is_route_lost(&error) && !rerouted => rerouted = true,
Err(error) => return Err(error),
}
}
}
async fn exchange(
&mut self,
partition_id: u32,
code: u32,
payload: Vec<u8>,
) -> Result<FilterOutcome, LaserError> {
let payload = Bytes::from(payload);
let reply = match self.request.read_mode {
ReadMode::Local => self.coordinator.send_binary_request(code, payload).await?,
ReadMode::Primary => {
let route_to = RouteTo {
source: &self.request.source,
consumer: &self.consumer,
};
let connection = self
.routes
.connection(
&self.coordinator,
self.laser.connection_string(),
&route_to,
partition_id,
code == AGDX_FILTERED_ACK_CODE,
)
.await?;
match connection.send_binary_request(code, payload).await {
Ok(reply) => reply,
Err(error) => {
self.routes.invalidate(partition_id).await;
return Err(error.into());
}
}
}
};
match decode_managed_reply::<FilterReply>(&reply)? {
FilterReply::Ok(outcome) => Ok(outcome),
FilterReply::Err(error) => {
if matches!(
error.reason,
FilterErrorReason::NotPrimary | FilterErrorReason::MembershipStale
) {
self.routes.invalidate(partition_id).await;
}
Err(error.into())
}
}
}
async fn refresh_membership(&mut self) -> Result<(), LaserError> {
let Some(membership) = self.membership.as_mut() else {
return Ok(());
};
if !membership.sync(&self.coordinator).await? {
return Ok(());
}
let assigned: BTreeSet<u32> = membership.partitions().iter().copied().collect();
let stale_routes = if membership.take_rejoined() {
self.owner = Arc::new(());
self.partitions.clear();
self.revoked.clear();
self.buffered.clear();
Some(std::mem::take(&mut self.routes))
} else {
None
};
let left: Vec<u32> = self
.partitions
.keys()
.filter(|partition_id| !assigned.contains(partition_id))
.copied()
.collect();
for partition_id in &left {
self.revoked.insert(*partition_id);
self.retire(*partition_id);
}
for partition_id in assigned {
if self.revoked.remove(&partition_id) || !self.partitions.contains_key(&partition_id) {
self.buffered
.retain(|record| record.partition_id != partition_id);
self.partitions
.insert(partition_id, PartitionProgress::new(FilteredStart::Next));
}
}
if stale_routes.is_some() && matches!(self.request.filter, FilterRef::Bound) {
self.guard = None;
}
let retired: Vec<_> = left
.into_iter()
.filter_map(|partition| self.routes.forget(partition))
.collect();
if let Some(routes) = stale_routes {
self.retired_connections += routes.opened();
routes.close().await;
}
for client in retired {
let _ = client.shutdown().await;
}
Ok(())
}
async fn refresh_partitions(&mut self) -> Result<(), LaserError> {
let listed = topic_partitions(&self.laser, &self.request.source).await?;
self.listed_at = Some(Instant::now());
for partition_id in listed {
self.partitions
.entry(partition_id)
.or_insert_with(|| PartitionProgress::new(FilteredStart::Next));
}
Ok(())
}
fn pin_source(&mut self, page: &FilteredPage) -> Result<(), LaserError> {
let FilterConsumer::GroupId(id) = self.consumer else {
return Ok(());
};
let incarnation = SourceIncarnation::from(page.generation);
match self.pinned_source {
None => {
self.pinned_source = Some(incarnation);
Ok(())
}
Some(pinned) if pinned == incarnation => Ok(()),
Some(_) => Err(FilterError::new(
FilterErrorReason::NotFound,
format!(
"the numeric consumer group {id} belongs to a replaced stream or topic, build a new reader"
),
)
.into()),
}
}
fn retire(&mut self, partition_id: u32) {
if self.revoked.contains(&partition_id)
&& self
.partitions
.get(&partition_id)
.is_some_and(|progress| progress.in_flight() == 0 && progress.unstored().is_none())
{
self.partitions.remove(&partition_id);
self.revoked.remove(&partition_id);
}
}
fn progress(&mut self, partition_id: u32) -> Result<&mut PartitionProgress, LaserError> {
self.partitions.get_mut(&partition_id).ok_or_else(|| {
FilterError::new(
FilterErrorReason::MembershipStale,
format!(
"partition {partition_id} is no longer read by this reader. Its unacknowledged records are read again"
),
)
.into()
})
}
fn require_sequence(&self, partition: u32, sequence: u64) -> Result<(), LaserError> {
if self
.partitions
.get(&partition)
.is_some_and(|progress| progress.knows(sequence))
{
return Ok(());
}
Err(FilterError::new(
FilterErrorReason::MembershipStale,
"the acknowledgment belongs to a retired partition read",
)
.into())
}
fn require_owner(&self, owner: &Arc<()>) -> Result<(), LaserError> {
if !Arc::ptr_eq(&self.owner, owner) {
return Err(LaserError::Invalid(
"the page or record belongs to another filtered reader, or was read before this reader rejoined its group"
.to_owned(),
));
}
Ok(())
}
fn require_acknowledgments(&self) -> Result<(), LaserError> {
if self.request.read_mode == ReadMode::Local {
return Err(LaserError::Invalid(
"a ReadMode::Local reader cannot acknowledge: only the partition primary proves the history a page came from".to_owned(),
));
}
Ok(())
}
}
async fn bound_definition(
laser: &Laser,
page: &FilteredPage,
names: &FilterSource,
) -> Result<ConsumerFilter, LaserError> {
let source = page.generation;
for index in 0_u32.. {
let bindings = laser
.filters()
.bindings(
None,
Some(&names.stream),
Some(&names.topic),
index,
MAX_FILTER_CATALOG_PAGE,
)
.await?;
if let Some(binding) = bindings.items.into_iter().find(|binding| {
let id = binding.identity;
Some(id.group_id) == page.policy.group_id
&& id.stream_id == source.stream_id
&& id.stream_created_at_micros == source.stream_created_at_micros
&& id.topic_id == source.topic_id
&& id.topic_created_at_micros == source.topic_created_at_micros
}) {
if binding.digest != page.policy.digest {
return Err(LaserError::Protocol(
"served policy differs from the group's catalog binding".to_owned(),
));
}
return definition(
laser,
&FilterRef::Revision {
filter_id: binding.filter_id,
revision: binding.revision,
},
)
.await;
}
if (index + 1).saturating_mul(MAX_FILTER_CATALOG_PAGE) >= bindings.total {
break;
}
}
Err(FilterError::new(
FilterErrorReason::NotFound,
"no readable catalog binding matches this group identity",
)
.into())
}
async fn definition(laser: &Laser, filter: &FilterRef) -> Result<ConsumerFilter, LaserError> {
let (filter_id, revision) = match filter {
FilterRef::Inline(filter) => return Ok(filter.clone()),
FilterRef::Revision {
filter_id,
revision,
} => (*filter_id, *revision),
FilterRef::Bound => {
return Err(LaserError::Config(
"resolve a bound guard from the served page",
));
}
};
let mut page = 0;
loop {
let revisions = laser
.filters()
.revisions(filter_id, page, MAX_FILTER_CATALOG_PAGE)
.await?;
if let Some(found) = revisions
.items
.into_iter()
.find(|info| info.revision == revision)
{
return Ok(found.filter);
}
page += 1;
if page.saturating_mul(MAX_FILTER_CATALOG_PAGE) >= revisions.total {
return Err(FilterError::new(
FilterErrorReason::NotFound,
format!("filter {filter_id} has no revision {revision}"),
)
.into());
}
}
}
async fn topic_partitions(laser: &Laser, source: &FilterSource) -> Result<Vec<u32>, LaserError> {
let stream = Identifier::named(&source.stream)?;
let topic = Identifier::named(&source.topic)?;
let details = laser
.client()
.get_topic(&stream, &topic)
.await?
.ok_or_else(|| IggyError::TopicNameNotFound(source.topic.clone(), source.stream.clone()))?;
Ok((0..details.partitions_count).collect())
}
fn is_route_lost(error: &LaserError) -> bool {
matches!(
error.filter_reason(),
Some(FilterErrorReason::NotPrimary | FilterErrorReason::MembershipStale)
) || matches!(
error,
LaserError::Iggy(IggyError::ConsumerGroupPartitionNotOwned(..))
)
}
const fn blocked_error(partition_id: u32, blocked: Blocked) -> LaserError {
match blocked {
Blocked::Fault { offset, reason } => LaserError::FilterFault {
partition_id,
offset,
reason,
},
Blocked::Oversized { offset } => LaserError::FilterOversizedRecord {
partition_id,
offset,
},
}
}
fn unexpected(verb: &str) -> LaserError {
LaserError::Protocol(format!("{verb}: unexpected reply variant"))
}
fn validate_record_body(
request: &FilteredPollRequest,
page: &FilteredPage,
) -> Result<(), LaserError> {
let invalid =
|| LaserError::Protocol("filtered record body has invalid framing or counts".to_owned());
let head = page.records.get(..16).ok_or_else(invalid)?;
let count = u32::from_le_bytes(head[12..16].try_into().map_err(|_| invalid())?);
if count != page.matched || count > request.count {
return Err(invalid());
}
let mut rest = &page.records[16..];
let mut decoded = 0u32;
while !rest.is_empty() {
let batch =
decode_batch_slice_with(rest, BatchIntegrity::LayoutOnly).map_err(|_| invalid())?;
decoded = decoded
.checked_add(batch.header.message_count)
.ok_or_else(invalid)?;
if decoded > count {
return Err(invalid());
}
for frame in batch.iter() {
batch
.header
.base_offset
.checked_add(u64::from(frame.header.offset_delta))
.ok_or_else(invalid)?;
batch
.header
.origin_timestamp
.checked_add(u64::from(frame.header.timestamp_delta))
.ok_or_else(invalid)?;
}
rest = rest.get(batch.header.total_size()..).ok_or_else(invalid)?;
}
if decoded != count {
return Err(invalid());
}
Ok(())
}
fn validate_page(
request: &FilteredPollRequest,
inline_digest: Option<&Digest32>,
page: &FilteredPage,
polled: &PolledMessages,
) -> Result<(), LaserError> {
let invalid = || {
LaserError::Protocol("filtered page does not match its request or scan boundary".to_owned())
};
if page.v != FILTER_OP_VERSION
|| page.partition_id != request.partition_id
|| page.generation.partition_id != request.partition_id
|| polled.partition_id != request.partition_id
|| page.read_mode != request.read_mode
|| polled.current_offset != page.frontier
|| page.matched as usize != polled.messages.len()
|| page.matched > page.examined
|| page.policy.digest.validate().is_err()
{
return Err(invalid());
}
if let FilterConsumer::GroupId(id) = request.consumer
&& page.policy.group_id != Some(id)
{
return Err(LaserError::Protocol(
"filtered page identifies another consumer group".to_owned(),
));
}
if request.consumer.is_group() != page.policy.group_id.is_some() {
return Err(invalid());
}
if inline_digest.is_some_and(|digest| *digest != page.policy.digest) {
return Err(invalid());
}
if let FilterRef::Revision {
filter_id,
revision,
} = &request.filter
&& (page.policy.filter_id != Some(*filter_id) || page.policy.revision != Some(*revision))
{
return Err(invalid());
}
let start = match &request.start {
FilteredStart::Continue(cursor) => {
if cursor.generation != page.generation
|| cursor.digest != page.policy.digest
|| cursor.group_id != page.policy.group_id
|| cursor.read_mode != request.read_mode
{
return Err(invalid());
}
Some(cursor.next_scan_offset)
}
FilteredStart::Offset(offset) => Some(*offset),
FilteredStart::First => Some(0),
_ => None,
};
let scanned_to = page
.next_scan_offset
.and_then(|offset| offset.checked_sub(1));
if page.examined == 0 {
if page.safe_ack_offset.is_some() || page.next_scan_offset != start {
return Err(invalid());
}
} else if scanned_to.is_none()
|| match request.read_mode {
ReadMode::Primary => page.safe_ack_offset != scanned_to,
ReadMode::Local => page.safe_ack_offset.is_some(),
}
{
return Err(invalid());
}
if scanned_to.is_some_and(|scanned| scanned > page.frontier) {
return Err(invalid());
}
if start
.zip(page.next_scan_offset)
.is_some_and(|(start, next)| next < start)
{
return Err(invalid());
}
let mut previous = None;
for message in &polled.messages {
let offset = message.header.offset;
if previous.is_some_and(|last| offset <= last)
|| start.is_some_and(|start| offset < start)
|| scanned_to.is_none_or(|scanned| offset > scanned)
|| offset > page.frontier
{
return Err(invalid());
}
previous = Some(offset);
}
let unevaluated: BTreeSet<_> = page.unevaluated.iter().copied().collect();
if unevaluated.len() != page.unevaluated.len()
|| unevaluated.iter().any(|offset| {
polled
.messages
.binary_search_by_key(offset, |message| message.header.offset)
.is_err()
})
{
return Err(invalid());
}
if matches!(page.stop, StopReason::Fault | StopReason::OversizedRecord) != page.fault.is_some()
{
return Err(invalid());
}
if let Some(fault) = page.fault
&& scanned_to.is_some_and(|scanned| scanned >= fault.offset)
{
return Err(invalid());
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use laser_wire::filter::{FilterExpr, FilteredPage};
fn sample() -> (FilteredPollRequest, FilteredPage, PolledMessages) {
let filter = ConsumerFilter::json(FilterExpr::present("mode"));
let request = FilteredPollRequest {
v: FILTER_OP_VERSION,
source: FilterSource {
stream: "orbit".to_owned(),
topic: "fleet_changes".to_owned(),
},
partition_id: 0,
consumer: FilterConsumer::Consumer("reader".to_owned()),
filter: FilterRef::Inline(filter.clone()),
start: FilteredStart::First,
count: 10,
max_reply_bytes: 1024,
read_mode: ReadMode::Primary,
};
let page = FilteredPage {
v: FILTER_OP_VERSION,
partition_id: 0,
policy: AppliedPolicy {
group_id: None,
digest: filter.digest(),
filter_id: None,
revision: None,
},
generation: SourceGeneration {
stream_id: 1,
stream_created_at_micros: 1,
topic_id: 1,
topic_created_at_micros: 1,
partition_id: 0,
partition_created_revision: 1,
purge_generation: 0,
},
read_mode: ReadMode::Primary,
next_scan_offset: Some(10),
safe_ack_offset: Some(9),
frontier: 9,
examined: 10,
matched: 0,
stop: StopReason::EndOfVisible,
fault: None,
unevaluated: Vec::new(),
evaluation_limits: None,
records: Vec::new(),
};
let polled = PolledMessages {
partition_id: 0,
current_offset: 9,
count: 0,
messages: Vec::new(),
};
(request, page, polled)
}
#[test]
fn given_empty_scanned_pages_when_validated_then_should_accept_only_consistent_progress() {
let (request, page, polled) = sample();
let digest = page.policy.digest.clone();
validate_page(&request, Some(&digest), &page, &polled).expect("valid scanned range");
let mut malformed = page.clone();
malformed.safe_ack_offset = Some(10);
assert!(validate_page(&request, Some(&digest), &malformed, &polled).is_err());
malformed = page.clone();
malformed.read_mode = ReadMode::Local;
assert!(validate_page(&request, Some(&digest), &malformed, &polled).is_err());
malformed = page.clone();
malformed.policy.digest = ConsumerFilter::json(FilterExpr::present("other")).digest();
assert!(validate_page(&request, Some(&digest), &malformed, &polled).is_err());
malformed = page.clone();
malformed.partition_id = 1;
assert!(validate_page(&request, Some(&digest), &malformed, &polled).is_err());
malformed = page;
malformed.examined = 0;
assert!(validate_page(&request, Some(&digest), &malformed, &polled).is_err());
}
#[test]
fn given_a_continuation_when_the_reply_changes_history_then_should_reject() {
let (mut request, mut page, polled) = sample();
request.start = FilteredStart::Continue(laser_wire::filter::Continuation {
group_id: None,
next_scan_offset: 0,
generation: page.generation,
digest: page.policy.digest.clone(),
read_mode: ReadMode::Primary,
});
page.generation.purge_generation = 1;
assert!(validate_page(&request, None, &page, &polled).is_err());
}
#[test]
fn given_a_forged_inner_count_when_decoding_then_should_reject_before_allocation() {
let (request, mut page, _) = sample();
page.records = vec![0; 16];
page.records[12..16].copy_from_slice(&u32::MAX.to_le_bytes());
assert!(validate_record_body(&request, &page).is_err());
}
#[test]
fn given_a_local_page_when_validated_then_should_carry_no_acknowledgment_offset() {
let (mut request, mut page, polled) = sample();
request.read_mode = ReadMode::Local;
page.read_mode = ReadMode::Local;
page.safe_ack_offset = None;
validate_page(&request, None, &page, &polled).expect("a scanned local page");
page.safe_ack_offset = Some(9);
assert!(
validate_page(&request, None, &page, &polled).is_err(),
"a local page never offers an acknowledgment"
);
}
#[test]
fn given_a_page_that_examined_nothing_when_its_scan_offset_moved_then_should_reject() {
let (request, mut page, polled) = sample();
page.examined = 0;
page.safe_ack_offset = None;
page.next_scan_offset = Some(0);
validate_page(&request, None, &page, &polled).expect("an empty scan repeats its start");
page.next_scan_offset = Some(5);
assert!(validate_page(&request, None, &page, &polled).is_err());
}
}