use std::sync::Arc;
use std::time::Duration;
use crate::journal::{JournalError, PartitionJournal};
use tokio_util::sync::CancellationToken;
use super::marks::{DirtySet, MAX_TRACKED_PARTITIONS, Pending, is_conversation_partition};
use super::project;
use super::store::{
self, CoverageState, ExcisionScan, JournalState, SearchProjection, SegmentStats, StoreError,
};
use super::terms::TermKey;
use crate::metrics;
const REPLAY_BUDGET_BYTES: u64 = 32 * 1024 * 1024;
const COMPACT_SEGMENT_THRESHOLD: u32 = 16;
const COMPACT_ROW_CEILING: u64 = 500_000;
const RECONCILE_PARTITION_LIMIT: usize = 65_536;
const DRAIN_INTERVAL: Duration = Duration::from_secs(2);
const COVERAGE_SWEEP_INTERVAL: Duration = Duration::from_mins(30);
const RECONCILE_RETRY_INTERVAL: Duration = Duration::from_mins(1);
const EXCISION_SCAN_BUDGET_BYTES: u64 = 4 * 1024 * 1024;
const BARRIER_LAG_WARN_EVENTS: u64 = 10_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Outcome {
Published {
indexed_through: u64,
},
Unavailable {
reason: UnavailableReason,
},
Removed,
Destroyed,
AlreadyCurrent,
BarrierHeld {
open_turn_at: u64,
journal_lag: u64,
},
Deferred,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum UnavailableReason {
ReplayBudgetExceeded,
RecordTooLarge,
ReplayFailed,
SourceUnreadable,
SourceEmpty,
StoreFailed,
}
impl UnavailableReason {
const fn is_transient(self) -> bool {
match self {
Self::ReplayFailed | Self::StoreFailed => true,
Self::ReplayBudgetExceeded
| Self::RecordTooLarge
| Self::SourceUnreadable
| Self::SourceEmpty => false,
}
}
const fn label(self) -> &'static str {
match self {
Self::ReplayBudgetExceeded => "replay_budget_exceeded",
Self::RecordTooLarge => "record_too_large",
Self::ReplayFailed => "replay_failed",
Self::SourceUnreadable => "source_unreadable",
Self::SourceEmpty => "source_empty",
Self::StoreFailed => "store_failed",
}
}
}
const fn replay_reason(error: &JournalError) -> UnavailableReason {
match error {
JournalError::Unreadable(_) => UnavailableReason::SourceUnreadable,
JournalError::Unreachable(_) => UnavailableReason::ReplayFailed,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RebuildCause {
Excision,
SourceReplacement,
UnavailableRecovery,
}
impl RebuildCause {
const fn label(self) -> &'static str {
match self {
Self::Excision => "excision",
Self::SourceReplacement => "source_replacement",
Self::UnavailableRecovery => "unavailable_recovery",
}
}
}
enum Pass {
Done(Outcome),
RebuildInstead {
cause: RebuildCause,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ReconcileSummary {
pub(crate) visited: usize,
pub(crate) refused: usize,
pub(crate) complete: bool,
}
#[derive(Clone)]
pub(crate) struct LiveJournal {
journal: Arc<dyn PartitionJournal>,
}
impl LiveJournal {
pub(crate) const fn new(journal: Arc<dyn PartitionJournal>) -> Self {
Self { journal }
}
}
impl JournalState for LiveJournal {
async fn source_incarnation(
&self,
partition: &str,
) -> Option<polyc_state::revision::PartitionIncarnation> {
self.journal
.partition_incarnation(partition.to_owned())
.await
.ok()
.flatten()
}
async fn excision_since(&self, partition: &str, scanned_through: u64) -> ExcisionScan {
let Ok(tail) = self
.journal
.partition_event_count(partition.to_owned())
.await
else {
return ExcisionScan::Unknown;
};
if scanned_through >= tail {
return ExcisionScan::Clear;
}
let Ok(replay) = self
.journal
.replay_range_with_positions_bounded(
partition.to_owned(),
scanned_through,
tail,
EXCISION_SCAN_BUDGET_BYTES,
)
.await
else {
return ExcisionScan::Unknown;
};
if replay.events.iter().any(|(_, event)| {
polyc_proto::kinds::base(&event.kind) == polyc_proto::kinds::TAINT_EXCISION
}) {
return ExcisionScan::Pending;
}
if replay.budget_exceeded {
ExcisionScan::Unknown
} else {
ExcisionScan::Clear
}
}
}
pub(crate) struct SearchIndexWorker {
projection: SearchProjection,
journal: Arc<dyn PartitionJournal>,
dirty: Arc<DirtySet>,
term_key: TermKey,
key_id: String,
replay_budget: u64,
row_ceiling: u64,
unrecorded_refusal: bool,
last_reconcile: Option<std::time::Instant>,
last_sweep: std::time::Instant,
recovery_blocked: std::collections::BTreeSet<String>,
}
impl SearchIndexWorker {
pub(crate) fn new(
projection: SearchProjection,
journal: Arc<dyn PartitionJournal>,
dirty: Arc<DirtySet>,
term_key: TermKey,
) -> Self {
let key_id = term_key.key_id();
Self {
projection,
journal,
dirty,
term_key,
key_id,
replay_budget: REPLAY_BUDGET_BYTES,
row_ceiling: COMPACT_ROW_CEILING,
unrecorded_refusal: false,
last_reconcile: None,
last_sweep: std::time::Instant::now(),
recovery_blocked: std::collections::BTreeSet::new(),
}
}
pub(crate) async fn run(mut self, shutdown: CancellationToken) {
loop {
if tokio::time::timeout(DRAIN_INTERVAL, shutdown.cancelled())
.await
.is_ok()
{
break;
}
if self.sweep_is_due() {
let degraded = self.dirty.degraded();
match self.reconcile(&shutdown).await {
Ok(summary) => tracing::info!(
visited = summary.visited,
refused = summary.refused,
complete = summary.complete,
degraded,
"search index swept every conversation to re-establish coverage"
),
Err(error) => tracing::error!(
%error,
"search index could not enumerate partitions to reconcile; coverage \
stays whatever the last sweep left"
),
}
self.last_sweep = std::time::Instant::now();
self.last_reconcile = self.dirty.degraded().then(std::time::Instant::now);
}
for (partition, outcome) in self.drain_once(&shutdown).await {
match outcome {
Outcome::Unavailable { reason } => tracing::warn!(
partition,
?reason,
"search index left a conversation unsearchable"
),
Outcome::BarrierHeld {
open_turn_at,
journal_lag,
} if journal_lag >= BARRIER_LAG_WARN_EVENTS => tracing::warn!(
partition,
open_turn_at,
journal_lag,
"search index cannot advance a conversation's watermark: a turn opened \
and never completed, so every committed turn behind it stays \
unindexable until that turn completes or fails"
),
Outcome::BarrierHeld { .. }
| Outcome::Published { .. }
| Outcome::Removed
| Outcome::Destroyed
| Outcome::AlreadyCurrent
| Outcome::Deferred => {}
}
}
}
}
fn sweep_is_due(&self) -> bool {
if self.dirty.degraded() {
return self.reconcile_is_due();
}
self.last_sweep.elapsed() >= COVERAGE_SWEEP_INTERVAL
}
fn reconcile_is_due(&self) -> bool {
self.last_reconcile
.is_none_or(|last| last.elapsed() >= RECONCILE_RETRY_INTERVAL)
}
pub(crate) async fn drain_once(
&mut self,
shutdown: &CancellationToken,
) -> Vec<(String, Outcome)> {
let pending = self.dirty.drain();
let mut outcomes = Vec::with_capacity(pending.len());
let mut remaining = pending.into_iter();
for (partition, pending) in remaining.by_ref() {
if shutdown.is_cancelled() {
self.dirty.mark(&partition, pending);
break;
}
let outcome = self.apply(&partition, pending).await;
if outcome == Outcome::Deferred {
self.dirty.mark(&partition, pending);
}
outcomes.push((partition, outcome));
}
for (partition, pending) in remaining {
self.dirty.mark(&partition, pending);
}
metrics::record_search_index_queue(self.dirty.pending_len(), self.dirty.degraded());
outcomes
}
pub(crate) async fn reconcile(
&mut self,
shutdown: &CancellationToken,
) -> Result<ReconcileSummary, JournalError> {
let degrade_mark = self.dirty.degrade_count();
let partitions = self.journal.list_partitions().await?;
let journal = LiveJournal::new(Arc::clone(&self.journal));
self.unrecorded_refusal = false;
self.recovery_blocked.clear();
let mut visited = 0usize;
let mut refused = 0usize;
let mut complete = true;
for partition in partitions
.iter()
.filter(|partition| is_conversation_partition(partition))
{
if shutdown.is_cancelled() {
complete = false;
break;
}
if visited >= RECONCILE_PARTITION_LIMIT {
complete = false;
tracing::error!(
limit = RECONCILE_PARTITION_LIMIT,
"search index reconcile stopped at its partition bound; coverage is \
re-established for a prefix of the deployment only, so the index stays \
degraded"
);
break;
}
visited += 1;
match self.reestablish(partition, &journal).await {
Outcome::Unavailable { .. } => refused += 1,
Outcome::Deferred => {
complete = false;
break;
}
Outcome::Published { .. }
| Outcome::Removed
| Outcome::Destroyed
| Outcome::BarrierHeld { .. }
| Outcome::AlreadyCurrent => {}
}
}
let cleared = complete
&& !self.unrecorded_refusal
&& !self.dirty.has_pending_removal()
&& self.dirty.clear_degraded(degrade_mark);
metrics::record_search_index_reconcile(cleared, refused);
metrics::record_search_index_queue(self.dirty.pending_len(), self.dirty.degraded());
Ok(ReconcileSummary {
visited,
refused,
complete,
})
}
async fn reestablish(&mut self, partition: &str, journal: &LiveJournal) -> Outcome {
let state = match self
.projection
.verified_coverage(partition, &self.key_id, journal)
.await
{
Ok(state) => state,
Err(err) if err.is_unreadable() => CoverageState::NeverIndexed,
Err(err) => return self.store_failed(partition, &err).await,
};
match state {
CoverageState::Destroyed => Outcome::Destroyed,
CoverageState::Indexed(coverage) if coverage.available => {
match self
.journal
.partition_event_count(partition.to_owned())
.await
{
Ok(end) => self.index_partition(partition, Some(end)).await,
Err(error) => {
self.replay_failed(
partition,
&error,
"search index reconcile could not read a partition's length",
)
.await
}
}
}
CoverageState::Indexed(_) | CoverageState::NeverIndexed | CoverageState::Stale => {
self.index_partition(partition, None).await
}
}
}
async fn apply(&mut self, partition: &str, pending: Pending) -> Outcome {
match pending {
Pending::Destroy => self.destroy(partition).await,
Pending::Replace => self.replace(partition).await,
Pending::Remove => self.remove(partition).await,
Pending::Rebuild => self.index_partition(partition, None).await,
Pending::IndexThrough(boundary) => {
self.index_partition(partition, Some(boundary)).await
}
}
}
async fn destroy(&self, partition: &str) -> Outcome {
match self.projection.destroy(partition, &self.key_id).await {
Ok(()) => Outcome::Destroyed,
Err(err) => {
tracing::error!(
partition,
error = %err,
"search index could not record a destroyed conversation; its rows may still \
be on disk"
);
self.dirty.mark(partition, Pending::Destroy);
Outcome::Unavailable {
reason: UnavailableReason::StoreFailed,
}
}
}
}
async fn remove(&self, partition: &str) -> Outcome {
match self.projection.remove(partition).await {
Ok(()) => Outcome::Removed,
Err(err) => {
tracing::error!(
partition,
error = %err,
"search index could not remove a migrated conversation's segments"
);
self.dirty.mark(partition, Pending::Remove);
Outcome::Unavailable {
reason: UnavailableReason::StoreFailed,
}
}
}
}
async fn replace(&mut self, partition: &str) -> Outcome {
if let Err(err) = self
.projection
.mark_unavailable(partition, &self.key_id)
.await
{
tracing::error!(
partition,
error = %err,
"search index could not revoke stale-source read authority before replacement"
);
self.unrecorded_refusal = true;
self.dirty.mark(partition, Pending::Replace);
return Outcome::Unavailable {
reason: UnavailableReason::StoreFailed,
};
}
if let Err(err) = self.projection.remove(partition).await {
tracing::error!(
partition,
error = %err,
"search index could not remove a stale source before replacement rebuild"
);
self.dirty.mark(partition, Pending::Replace);
return Outcome::Unavailable {
reason: UnavailableReason::StoreFailed,
};
}
self.index_partition(partition, None).await
}
pub(crate) async fn index_partition(
&mut self,
partition: &str,
boundary: Option<u64>,
) -> Outcome {
let mut boundary = boundary;
let mut attempted_recovery = false;
loop {
match self.index_pass(partition, boundary).await {
Pass::Done(outcome) => {
self.note_recovery(partition, attempted_recovery, &outcome);
return outcome;
}
Pass::RebuildInstead { cause } => {
tracing::info!(
partition,
cause = cause.label(),
"search index is rebuilding a conversation a forward pass cannot repair"
);
metrics::record_search_index_rebuild(cause.label());
attempted_recovery |= cause == RebuildCause::UnavailableRecovery;
boundary = None;
}
}
}
}
fn note_recovery(&mut self, partition: &str, attempted_recovery: bool, outcome: &Outcome) {
match outcome {
Outcome::Published { .. } | Outcome::Destroyed | Outcome::Removed => {
self.recovery_blocked.remove(partition);
}
Outcome::Unavailable { reason } if attempted_recovery && !reason.is_transient() => {
if self.recovery_blocked.len() < MAX_TRACKED_PARTITIONS {
self.recovery_blocked.insert(partition.to_owned());
}
}
Outcome::Unavailable { .. }
| Outcome::BarrierHeld { .. }
| Outcome::AlreadyCurrent
| Outcome::Deferred => {}
}
}
async fn index_pass(&mut self, partition: &str, boundary: Option<u64>) -> Pass {
let partition = partition.to_owned();
let state = match self.projection.coverage(&partition, &self.key_id).await {
Ok(state) => state,
Err(err) if err.is_unreadable() => CoverageState::NeverIndexed,
Err(err) => return Pass::Done(self.store_failed(&partition, &err).await),
};
if matches!(state, CoverageState::Destroyed) {
return Pass::Done(Outcome::Destroyed);
}
let source_incarnation = match self.journal.partition_incarnation(partition.clone()).await {
Ok(Some(incarnation)) => incarnation,
Ok(None) => {
return Pass::Done(
self.fail_unavailable(&partition, UnavailableReason::SourceEmpty)
.await,
);
}
Err(error) => {
return Pass::Done(
self.replay_failed(
&partition,
&error,
"search index could not resolve the partition's exact source",
)
.await,
);
}
};
if let Some(cause) = self.rebuild_instead(&partition, boundary, &state, source_incarnation)
{
return Pass::RebuildInstead { cause };
}
let start = match self.resolve_start(&partition, boundary, &state).await {
Ok(start) => start,
Err(outcome) => return Pass::Done(outcome),
};
let end = match self.resolve_end(&partition, boundary).await {
Ok(end) => end,
Err(outcome) => return Pass::Done(outcome),
};
if boundary.is_some() && end <= start {
return Pass::Done(Outcome::AlreadyCurrent);
}
if end == 0 {
tracing::warn!(
partition,
"search index found no events to rebuild a conversation from; leaving it \
unsearchable"
);
return Pass::Done(
self.fail_unavailable(&partition, UnavailableReason::SourceEmpty)
.await,
);
}
let replay = match self
.journal
.replay_range_with_positions_bounded(partition.clone(), start, end, self.replay_budget)
.await
{
Ok(replay) => replay,
Err(error) => {
return Pass::Done(
self.replay_failed(
&partition,
&error,
"search index replay failed; leaving the conversation unsearchable",
)
.await,
);
}
};
if replay.budget_exceeded {
tracing::warn!(
partition,
bytes_read = replay.bytes_read,
"search index replay exceeded its byte budget; leaving the conversation \
unsearchable"
);
return Pass::Done(
self.fail_unavailable(&partition, UnavailableReason::ReplayBudgetExceeded)
.await,
);
}
let built = project::project(
&self.term_key,
&partition,
source_incarnation,
&replay.events,
end,
);
self.settle_pass(&partition, source_incarnation, boundary, start, end, &built)
.await
}
fn rebuild_instead(
&self,
partition: &str,
boundary: Option<u64>,
state: &CoverageState,
source_incarnation: polyc_state::revision::PartitionIncarnation,
) -> Option<RebuildCause> {
boundary?;
if matches!(
state,
CoverageState::Indexed(coverage)
if coverage.source_incarnation != source_incarnation
) {
return Some(RebuildCause::SourceReplacement);
}
if matches!(state, CoverageState::Indexed(coverage) if !coverage.available)
&& !self.recovery_blocked.contains(partition)
{
return Some(RebuildCause::UnavailableRecovery);
}
None
}
async fn settle_pass(
&mut self,
partition: &str,
source_incarnation: polyc_state::revision::PartitionIncarnation,
boundary: Option<u64>,
start: u64,
end: u64,
built: &project::Projected,
) -> Pass {
let partition = partition.to_owned();
let current_incarnation = self.journal.partition_incarnation(partition.clone()).await;
if !matches!(current_incarnation, Ok(Some(current)) if current == source_incarnation) {
tracing::warn!(
partition,
"search index source changed while replaying; refusing the stale publication"
);
let outcome = self
.fail_unavailable(&partition, UnavailableReason::ReplayFailed)
.await;
self.dirty.mark(&partition, Pending::Replace);
return Pass::Done(outcome);
}
if boundary.is_some() && built.excision_in_range {
return Pass::RebuildInstead {
cause: RebuildCause::Excision,
};
}
if boundary.is_some() && built.coverage.indexed_through <= start {
return Pass::Done(built.open_turn_at.map_or(Outcome::AlreadyCurrent, |at| {
metrics::record_search_index_barrier_held();
Outcome::BarrierHeld {
open_turn_at: at,
journal_lag: end.saturating_sub(at),
}
}));
}
let outcome = self.publish(&partition, start, built).await;
if let (Some(open_turn_at), Outcome::Published { .. }) = (built.open_turn_at, &outcome) {
metrics::record_search_index_barrier_held();
Pass::Done(Outcome::BarrierHeld {
open_turn_at,
journal_lag: end.saturating_sub(open_turn_at),
})
} else {
Pass::Done(outcome)
}
}
async fn publish(
&mut self,
partition: &str,
start: u64,
built: &project::Projected,
) -> Outcome {
let published = if start == 0 {
self.projection
.rebuild(partition, &built.messages, &built.coverage, &self.key_id)
.await
.map(|stats| store::Appended {
stats,
clamped_unavailable: false,
})
} else {
self.projection
.append(partition, &built.messages, &built.coverage, &self.key_id)
.await
};
match published {
Ok(appended) => {
if appended.clamped_unavailable {
tracing::warn!(
partition,
"search index indexed a conversation that stays refused: its recovery \
rebuild already failed the same way, so it waits for the next sweep or \
for a fix outside this worker"
);
}
self.compact_if_due(partition, appended.stats).await;
Outcome::Published {
indexed_through: built.coverage.indexed_through,
}
}
Err(StoreError::Destroyed) => Outcome::Destroyed,
Err(StoreError::TooLarge) => {
self.fail_unavailable(partition, UnavailableReason::RecordTooLarge)
.await
}
Err(err) => self.store_failed(partition, &err).await,
}
}
async fn compact_if_due(&self, partition: &str, stats: SegmentStats) {
if !should_compact(stats, self.row_ceiling) {
return;
}
if let Err(err) = self.projection.compact(partition, &self.key_id).await {
tracing::warn!(
partition,
error = %err,
"search index could not compact a conversation's segments; its rows are intact \
and reads stay correct, only slower"
);
}
}
async fn resolve_start(
&mut self,
partition: &str,
boundary: Option<u64>,
state: &CoverageState,
) -> Result<u64, Outcome> {
let (CoverageState::Indexed(coverage), Some(_)) = (state, boundary) else {
return Ok(0);
};
if coverage.indexed_through == 0 {
return Ok(0);
}
match self.projection.segment_stats(partition, &self.key_id).await {
Ok(on_disk) if on_disk.segments > 0 => Ok(coverage.indexed_through),
Ok(_) => Ok(0),
Err(err) if err.is_unreadable() => Ok(0),
Err(StoreError::TooLarge) => {
tracing::warn!(
partition,
"search index cannot decode a conversation's own rows; leaving it unsearchable"
);
Err(self
.fail_unavailable(partition, UnavailableReason::RecordTooLarge)
.await)
}
Err(err) => Err(self.store_failed(partition, &err).await),
}
}
async fn resolve_end(
&mut self,
partition: &str,
boundary: Option<u64>,
) -> Result<u64, Outcome> {
match boundary {
Some(boundary) => Ok(boundary),
None => match self
.journal
.partition_event_count(partition.to_owned())
.await
{
Ok(count) => Ok(count),
Err(error) => Err(self
.replay_failed(
partition,
&error,
"search index could not read a partition's length; leaving the \
conversation unsearchable",
)
.await),
},
}
}
async fn replay_failed(
&mut self,
partition: &str,
error: &JournalError,
message: &'static str,
) -> Outcome {
if matches!(error, JournalError::Unreachable(_)) && self.journal.is_stopping() {
tracing::debug!(
partition,
"search index stopped mid-conversation because the event log has shut down; \
nothing published and the watermark is untouched"
);
return Outcome::Deferred;
}
tracing::warn!(partition, %error, "{message}");
self.fail_unavailable(partition, replay_reason(error)).await
}
async fn store_failed(&mut self, partition: &str, err: &StoreError) -> Outcome {
tracing::error!(
partition,
error = %err,
"search index projection operation failed; leaving the conversation unsearchable"
);
self.fail_unavailable(partition, UnavailableReason::StoreFailed)
.await
}
async fn fail_unavailable(&mut self, partition: &str, reason: UnavailableReason) -> Outcome {
if let Err(err) = self
.projection
.mark_unavailable(partition, &self.key_id)
.await
{
tracing::error!(
partition,
error = %err,
"search index could not even record that a conversation is unsearchable"
);
self.unrecorded_refusal = true;
}
if reason.is_transient() {
self.dirty.mark(partition, Pending::Rebuild);
}
metrics::record_search_index_unavailable(reason.label());
Outcome::Unavailable { reason }
}
#[cfg(test)]
pub(crate) const fn with_replay_budget(mut self, bytes: u64) -> Self {
self.replay_budget = bytes;
self
}
#[cfg(test)]
pub(crate) const fn with_row_ceiling(mut self, rows: u64) -> Self {
self.row_ceiling = rows;
self
}
}
const fn should_compact(stats: SegmentStats, row_ceiling: u64) -> bool {
stats.segments >= COMPACT_SEGMENT_THRESHOLD || stats.unmerged_rows >= row_ceiling
}
#[cfg(test)]
mod tests;