From 547fa1ba2e8be1ace8839dc50a8ea6c11784769e Mon Sep 17 00:00:00 2001 From: "nadav.govari" Date: Mon, 5 Oct 2026 16:03:22 -0400 Subject: [PATCH] Pick shards based on minimum queue size; refactor persist --- .../quickwit-ingest/src/ingest_v2/ingester.rs | 342 ++---------------- .../src/ingest_v2/mrecordlog_utils.rs | 23 +- .../quickwit-ingest/src/ingest_v2/router.rs | 3 +- .../quickwit-ingest/src/ingest_v2/state.rs | 287 ++++++++++++++- .../protos/quickwit/ingester.proto | 3 +- .../quickwit/quickwit.ingest.ingester.rs | 11 +- quickwit/quickwit-proto/src/ingest/mod.rs | 3 +- 7 files changed, 337 insertions(+), 335 deletions(-) diff --git a/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs b/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs index ee7ec5fc001..9b214153db1 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs @@ -27,18 +27,16 @@ use quickwit_common::metrics::IN_FLIGHT_INGESTER_PERSIST; use quickwit_common::pretty::PrettyDisplay; use quickwit_common::pubsub::{EventBroker, EventSubscriber}; use quickwit_common::rate_limiter::{RateLimiter, RateLimiterSettings}; -use quickwit_common::{ServiceStream, rate_limited_error, rate_limited_warn}; +use quickwit_common::{ServiceStream, rate_limited_error}; use quickwit_metrics::{GaugeGuard, counter, label_values}; use quickwit_proto::control_plane::{ AdviseResetShardsRequest, ControlPlaneService, ControlPlaneServiceClient, }; use quickwit_proto::indexing::ShardPositionsUpdate; use quickwit_proto::ingest::ingester::*; -use quickwit_proto::ingest::{ - CommitTypeV2, DocBatchV2, IngestV2Error, IngestV2Result, ParseFailure, Shard, ShardIds, -}; +use quickwit_proto::ingest::{CommitTypeV2, IngestV2Error, IngestV2Result, Shard, ShardIds}; use quickwit_proto::types::{ - IndexUid, NodeId, Position, QueueId, ShardId, SourceId, SubrequestId, queue_id, split_queue_id, + IndexUid, NodeId, Position, QueueId, ShardId, SourceId, queue_id, split_queue_id, }; use serde_json::{Value as JsonValue, json}; use tokio::sync::{Semaphore, watch}; @@ -46,25 +44,19 @@ use tokio::time::{sleep, timeout}; use tracing::{Span, debug, error, info, instrument, warn}; use super::broadcast::BroadcastIngesterCapacityScoreTask; -use super::doc_mapper::validate_doc_batch; use super::fetch::FetchStreamTask; use super::idle::CloseIdleShardsTask; use super::local_shards::ShardThroughputReadings; use super::metrics::report_local_shards_metrics; use super::models::IngesterShard; -use super::mrecordlog_utils::{ - AppendDocBatchError, append_non_empty_doc_batch, check_enough_capacity, doc_batch_size, - wal_stats, -}; +use super::mrecordlog_utils::wal_stats; use super::rate_meter::RateMeter; use super::state::{IngesterState, InnerIngesterState, WeakIngesterState}; -use crate::estimate_size; use crate::ingest_v2::doc_mapper::get_or_try_build_doc_mapper; use crate::ingest_v2::metrics::{ DECOMMISSION_FAILED, DECOMMISSION_SUCCEEDED, RESET_SHARDS_OPERATIONS_TOTAL, STATUS, report_wal_limits, report_wal_usage, }; -use crate::metrics::{DOCS_BYTES_TOTAL, DOCS_TOTAL, VALIDITY}; use crate::mrecordlog_async::MultiRecordLogAsync; /// Minimum interval between two reset shards operations. @@ -401,316 +393,52 @@ impl Ingester { } let mut persist_successes = Vec::with_capacity(persist_request.subrequests.len()); let mut persist_failures = Vec::new(); - let mut pending_persist_subrequests: HashMap = - HashMap::with_capacity(persist_request.subrequests.len()); - - // Keep track of the shards that need to be closed following an IO error. - let mut shards_to_close: HashSet = HashSet::new(); - - // Keep track of dangling shards, i.e., shards for which there is no longer a corresponding - // queue in the WAL and should be deleted. - let mut shards_to_delete: HashSet = HashSet::new(); - - let commit_type = persist_request.commit_type(); - let force_commit = commit_type == CommitTypeV2::Force; - - let mut state_guard = self.state.lock_fully("persist").await?; - let status = state_guard.status(); + let mut staged_requests = Vec::with_capacity(persist_request.subrequests.len()); + let force_commit = persist_request.commit_type() == CommitTypeV2::Force; + let status = self.status(); if !status.accepts_write_requests() { - persist_failures.reserve_exact(persist_request.subrequests.len()); - - for subrequest in persist_request.subrequests { - let persist_failure = PersistFailure { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - reason: PersistFailureReason::NodeUnavailable as i32, - }; - persist_failures.push(persist_failure); - } - let persist_response = PersistResponse { - ingester_id: persist_request.ingester_id, - successes: Vec::new(), - failures: persist_failures, - routing_update: None, - }; - return Ok(persist_response); + return Err(IngestV2Error::Unavailable(format!( + "ingester {} is not accepting write requests. Status: {}", + self.self_node_id, status + ))); } - // first verify if we would locally accept each subrequest - { - let mut total_requested_capacity = ByteSize::b(0); - - for subrequest in persist_request.subrequests { - let Some(shard) = state_guard - .inner - .find_most_capacity_shard_mut(subrequest.index_uid(), &subrequest.source_id) - else { - warn!( - index_uid=%subrequest.index_uid(), - source_id=%subrequest.source_id, - "no open shard found on ingester" - ); - let persist_failure = PersistFailure { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - reason: PersistFailureReason::NoShardsAvailable as i32, - }; - persist_failures.push(persist_failure); - continue; - }; - let shard_id = shard.shard_id.clone(); - - // A router can only know about a newly opened shard if it has been informed by the - // control plane, which confirms that the shard was correctly opened in the - // metastore. - shard.is_advertisable = true; - let doc_mapper = shard.doc_mapper_opt.clone().expect("shard should be open"); - let validate_docs = shard.validate_docs; - let from_position_exclusive = shard.replication_position_inclusive.clone(); - - let doc_batch = match subrequest.doc_batch { - Some(doc_batch) if !doc_batch.is_empty() => doc_batch, - _ => { - warn!("received empty persist request"); - DocBatchV2::default() - } - }; - let requested_capacity = estimate_size(&doc_batch); - - if let Err(error) = check_enough_capacity( - &state_guard.mrecordlog, - self.disk_capacity, - self.memory_capacity, - requested_capacity + total_requested_capacity, - ) { - rate_limited_warn!( - limit_per_min = 10, - "failed to persist records to ingester `{}`: {error}", - self.self_node_id - ); - let persist_failure = PersistFailure { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - reason: PersistFailureReason::WalFull as i32, - }; - persist_failures.push(persist_failure); - continue; - }; - // Because we return the shard with the most available capacity, if this hits, it - // means that no shard can receive this request, and it should be retried. - if !shard.rate_limiter.acquire_bytes(requested_capacity) { - debug!( - "failed to persist records to shard `{}`: rate limited", - shard.queue_id() - ); - - let persist_failure = PersistFailure { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - reason: PersistFailureReason::NoShardsAvailable as i32, - }; - persist_failures.push(persist_failure); - continue; - } - - // Total number of bytes (valid and invalid documents) - let original_batch_num_bytes = doc_batch.num_bytes() as u64; - - let (valid_doc_batch, parse_failures) = if validate_docs { - validate_doc_batch(doc_batch, doc_mapper).await? - } else { - (doc_batch, Vec::new()) - }; - - if valid_doc_batch.is_empty() { - counter!( - parent: DOCS_TOTAL, - labels: [label_values!(VALIDITY => "invalid")], - ) - .inc_by(parse_failures.len() as u64); - counter!( - parent: DOCS_BYTES_TOTAL, - labels: [label_values!(VALIDITY => "invalid")], - ) - .inc_by(original_batch_num_bytes); - let persist_success = PersistSuccess { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - shard_id: Some(shard_id), - replication_position_inclusive: Some(from_position_exclusive), - num_persisted_docs: 0, - parse_failures, - }; - persist_successes.push(persist_success); - continue; - }; - - counter!( - parent: DOCS_TOTAL, - labels: [label_values!(VALIDITY => "valid")], - ) - .inc_by(valid_doc_batch.num_docs() as u64); - counter!( - parent: DOCS_BYTES_TOTAL, - labels: [label_values!(VALIDITY => "valid")], - ) - .inc_by(valid_doc_batch.num_bytes() as u64); - if !parse_failures.is_empty() { - counter!( - parent: DOCS_TOTAL, - labels: [label_values!(VALIDITY => "invalid")], - ) - .inc_by(parse_failures.len() as u64); - counter!( - parent: DOCS_BYTES_TOTAL, - labels: [label_values!(VALIDITY => "invalid")], - ) - .inc_by(original_batch_num_bytes - valid_doc_batch.num_bytes() as u64); - } - let valid_batch_num_bytes = valid_doc_batch.num_bytes() as u64; - shard.rate_meter.update(valid_batch_num_bytes); - total_requested_capacity += requested_capacity; - - let pending_persist_subrequest = PendingPersistSubrequest { - queue_id: shard.queue_id(), - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - shard_id: Some(shard_id), - doc_batch: valid_doc_batch, - parse_failures, - }; - pending_persist_subrequests.insert( - pending_persist_subrequest.subrequest_id, - pending_persist_subrequest, - ); + let mut state_guard = self.state.lock_fully("persist").await?; + let mut context = state_guard.begin_persist(force_commit); + for subrequest in persist_request.subrequests { + match state_guard + .stage_persist_subrequest(subrequest, &mut context) + .await + { + Ok(staged) => staged_requests.push(staged), + Err(failure) => persist_failures.push(failure), } } - // finally write locally - { - let now = Instant::now(); - for subrequest in pending_persist_subrequests.into_values() { - let queue_id = subrequest.queue_id; - - let batch_num_docs = subrequest.doc_batch.num_docs() as u64; - let batch_size = doc_batch_size(&subrequest.doc_batch, force_commit); - - let append_result = append_non_empty_doc_batch( - &mut state_guard.mrecordlog, - &queue_id, - subrequest.doc_batch, - force_commit, - ) - .await; - - let current_position_inclusive = match append_result { - Ok(current_position_inclusive) => current_position_inclusive, - Err(append_error) => { - let reason = match &append_error { - AppendDocBatchError::Io(io_error) => { - error!( - "failed to persist records to shard `{queue_id}`: {io_error}" - ); - shards_to_close.insert(queue_id); - PersistFailureReason::NodeUnavailable - } - AppendDocBatchError::QueueNotFound(_) => { - error!( - "failed to persist records to shard `{queue_id}`: WAL queue \ - not found" - ); - shards_to_delete.insert(queue_id); - PersistFailureReason::NodeUnavailable - } - }; - let persist_failure = PersistFailure { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - reason: reason as i32, - }; - persist_failures.push(persist_failure); - continue; - } - }; - - let shard = state_guard - .shards - .get_mut(&queue_id) - .expect("shard should exist"); - shard.set_replication_position_inclusive(current_position_inclusive.clone(), now); - shard.queue_size += batch_size; - - let persist_success = PersistSuccess { - subrequest_id: subrequest.subrequest_id, - index_uid: subrequest.index_uid, - source_id: subrequest.source_id, - shard_id: subrequest.shard_id, - replication_position_inclusive: Some(current_position_inclusive), - num_persisted_docs: batch_num_docs as u32, - parse_failures: subrequest.parse_failures, - }; - persist_successes.push(persist_success); + for staged_request in staged_requests { + match state_guard + .persist_subrequest(staged_request, &mut context) + .await + { + Ok(success) => persist_successes.push(success), + Err(failure) => persist_failures.push(failure), } } - if !shards_to_close.is_empty() { - for queue_id in &shards_to_close { - let shard = state_guard - .shards - .get_mut(queue_id) - .expect("shard should exist"); + state_guard.finish_persist(context); - shard.close(); - warn!("closed shard `{queue_id}` following IO error"); - } - } - if !shards_to_delete.is_empty() { - for queue_id in &shards_to_delete { - state_guard.shards.remove(queue_id); - warn!("deleted dangling shard `{queue_id}`"); - } - } let wal_usage = state_guard.mrecordlog.resource_usage(); let disk_used = wal_usage.disk_used_bytes as u64; - let memory_used = wal_usage.memory_used_bytes as u64; - let (open_shard_counts, closed_shards) = state_guard.get_shard_snapshot(); - let capacity_score = state_guard - .wal_capacity_tracker - .score(ByteSize::b(disk_used), ByteSize::b(memory_used)) - as u32; - // Take advantage of the state guard already being held to potentially update the shard - // throughput readings. + let routing_update = state_guard.routing_update(&wal_usage); let shard_throughput_readings_opt = state_guard.harvest_shard_throughput_readings(); drop(state_guard); - if let Some(snapshot) = shard_throughput_readings_opt { report_local_shards_metrics(&snapshot); } + if disk_used >= self.disk_capacity.as_u64() * 90 / 100 { self.background_reset_shards(); } report_wal_usage(wal_usage, self.disk_capacity, self.memory_capacity); - let source_shard_updates = open_shard_counts - .into_iter() - .map(|(index_uid, source_id, count)| SourceShardUpdate { - index_uid: Some(index_uid), - source_id, - open_shard_count: count as u32, - }) - .collect(); - - let routing_update = RoutingUpdate { - capacity_score, - source_shard_updates, - closed_shards, - }; - #[cfg(test)] { persist_successes.sort_by_key(|success| success.subrequest_id); @@ -1162,16 +890,6 @@ async fn apply_local_shard_updates(state: &IngesterState, local_updates: Vec<(Qu ); } -struct PendingPersistSubrequest { - queue_id: QueueId, - subrequest_id: u32, - index_uid: Option, - source_id: SourceId, - shard_id: Option, - doc_batch: DocBatchV2, - parse_failures: Vec, -} - #[cfg(test)] mod tests { #![allow(clippy::mutable_key_type)] diff --git a/quickwit/quickwit-ingest/src/ingest_v2/mrecordlog_utils.rs b/quickwit/quickwit-ingest/src/ingest_v2/mrecordlog_utils.rs index 8a7b039053a..83fef2aa325 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/mrecordlog_utils.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/mrecordlog_utils.rs @@ -24,8 +24,8 @@ use quickwit_proto::types::{Position, QueueId}; use tracing::instrument; use super::mrecord::MRECORD_HEADER_LEN; +use crate::MRecord; use crate::mrecordlog_async::MultiRecordLogAsync; -use crate::{MRecord, estimate_size}; #[derive(Debug, thiserror::Error)] pub(super) enum AppendDocBatchError { @@ -100,12 +100,23 @@ pub(super) async fn append_non_empty_doc_batch( } pub(super) fn doc_batch_size(doc_batch: &DocBatchV2, force_commit: bool) -> ByteSize { - let estimated_size = estimate_size(doc_batch); - if force_commit { - estimated_size + ByteSize::b(MRECORD_HEADER_LEN as u64) - } else { - estimated_size + encoded_doc_batch_size( + doc_batch.num_bytes() as u64, + doc_batch.num_docs(), + force_commit, + ) +} + +pub(super) fn encoded_doc_batch_size( + num_bytes: u64, + num_docs: usize, + force_commit: bool, +) -> ByteSize { + if num_docs == 0 { + return ByteSize::default(); } + let num_records = num_docs as u64 + u64::from(force_commit); + ByteSize::b(num_bytes + num_records * MRECORD_HEADER_LEN as u64) } pub(super) fn read_queue_size(mrecordlog: &MultiRecordLogAsync, queue_id: &QueueId) -> ByteSize { diff --git a/quickwit/quickwit-ingest/src/ingest_v2/router.rs b/quickwit/quickwit-ingest/src/ingest_v2/router.rs index 52f73cc1707..ef741180a91 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/router.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/router.rs @@ -292,12 +292,13 @@ impl IngestRouter { workbench.record_persist_failure(&persist_failure); match persist_failure.reason() { - PersistFailureReason::NoShardsAvailable => { + PersistFailureReason::NoShardsForSource => { // For non-critical failures, we don't mark the nodes unavailable; // a routing update is piggybacked on PersistResponses, so shard // counts and capacity scores will be fresh on the next try. } PersistFailureReason::NodeUnavailable + | PersistFailureReason::Internal | PersistFailureReason::WalFull | PersistFailureReason::Timeout => { unavailable_ingesters.insert(ingester_id.clone()); diff --git a/quickwit/quickwit-ingest/src/ingest_v2/state.rs b/quickwit/quickwit-ingest/src/ingest_v2/state.rs index a8745d03102..6fd10a165e4 100644 --- a/quickwit/quickwit-ingest/src/ingest_v2/state.rs +++ b/quickwit/quickwit-ingest/src/ingest_v2/state.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt; use std::ops::{Deref, DerefMut}; use std::path::Path; @@ -21,31 +21,41 @@ use std::time::{Duration, Instant}; use bytesize::ByteSize; use itertools::Itertools; +use mrecordlog::ResourceUsage; use mrecordlog::error::{DeleteQueueError, TruncateError}; use quickwit_cluster::Cluster; use quickwit_common::pretty::PrettyDisplay; +use quickwit_common::rate_limited_warn; use quickwit_common::rate_limiter::{RateLimiter, RateLimiterSettings}; use quickwit_common::shared_consts::INGESTER_STATUS_KEY; use quickwit_doc_mapper::DocMapper; -use quickwit_metrics::{gauge, histogram, labels}; +use quickwit_metrics::{counter, gauge, histogram, label_values, labels}; use quickwit_proto::control_plane::AdviseResetShardsResponse; -use quickwit_proto::ingest::ingester::IngesterStatus; -use quickwit_proto::ingest::{IngestV2Error, IngestV2Result, ShardIds, ShardState}; +use quickwit_proto::ingest::ingester::{ + IngesterStatus, PersistFailure, PersistFailureReason, PersistSubrequest, PersistSuccess, + RoutingUpdate, SourceShardUpdate, +}; +use quickwit_proto::ingest::{ + DocBatchV2, IngestV2Error, IngestV2Result, ParseFailure, ShardIds, ShardState, +}; use quickwit_proto::types::{ - DocMappingUid, IndexUid, Position, QueueId, SourceId, SourceUid, split_queue_id, + DocMappingUid, IndexUid, Position, QueueId, ShardId, SourceId, SourceUid, split_queue_id, }; use tokio::sync::{Mutex, MutexGuard, RwLock, RwLockMappedWriteGuard, RwLockWriteGuard, watch}; use tokio::task::JoinHandle; use tokio::time::MissedTickBehavior; -use tracing::{error, info, instrument}; +use tracing::{error, info, instrument, warn}; use super::local_shards::{ShardInfo, ShardInfos, ShardThroughputReadings}; use super::metrics::report_local_shards_metrics; use super::models::IngesterShard; -use super::mrecordlog_utils::read_queue_size; +use super::mrecordlog_utils::{ + AppendDocBatchError, append_non_empty_doc_batch, doc_batch_size, read_queue_size, +}; use super::rate_meter::RateMeter; use super::wal_capacity_tracker::WalCapacityTracker; use crate::OpenShardCounts; +use crate::metrics::{DOCS_BYTES_TOTAL, DOCS_TOTAL, VALIDITY}; use crate::mrecordlog_async::MultiRecordLogAsync; const LOCAL_SHARDS_SAMPLE_INTERVAL: Duration = if cfg!(any(test, feature = "testsuite")) { @@ -110,7 +120,7 @@ impl InnerIngesterState { } } - /// Returns the shard with the most available permits for this index and source. + /// Returns the shard with the smallesy queue size for this index and source. pub fn find_most_capacity_shard_mut( &mut self, index_uid: &IndexUid, @@ -121,9 +131,7 @@ impl InnerIngesterState { .filter(|shard| { shard.is_open() && shard.index_uid == *index_uid && shard.source_id == *source_id }) - .map(|shard| (shard.rate_limiter.available_permits(), shard)) - .max_by_key(|(available_permits, _)| *available_permits) - .map(|(_, shard)| shard) + .min_by_key(|shard| shard.queue_size) } /// Returns per-source open shard counts and closed shard IDs for all advertisable shards. @@ -161,6 +169,27 @@ impl InnerIngesterState { (open_counts, closed_shards) } + pub fn routing_update(&self, wal_usage: &ResourceUsage) -> RoutingUpdate { + let capacity_score = self.wal_capacity_tracker.score( + ByteSize::b(wal_usage.disk_used_bytes as u64), + ByteSize::b(wal_usage.memory_used_bytes as u64), + ) as u32; + let (open_shard_counts, closed_shards) = self.get_shard_snapshot(); + let source_shard_updates = open_shard_counts + .into_iter() + .map(|(index_uid, source_id, count)| SourceShardUpdate { + index_uid: Some(index_uid), + source_id, + open_shard_count: count as u32, + }) + .collect(); + RoutingUpdate { + capacity_score, + source_shard_updates, + closed_shards, + } + } + /// Every LOCAL_SHARDS_SAMPLE_INTERVAL, measure how much work each shard did, per source, to /// report to the control plane. The value kept is always the latest, as the only data that /// matters is the most recent snapshot. @@ -648,6 +677,167 @@ where } impl FullyLockedIngesterState<'_> { + pub fn begin_persist(&self, force_commit: bool) -> PersistContext { + PersistContext { + force_commit, + wal_usage: self.mrecordlog.resource_usage(), + reserved_capacity: ByteSize::default(), + shards_to_close: HashSet::new(), + shards_to_delete: HashSet::new(), + } + } + + pub async fn stage_persist_subrequest( + &mut self, + mut subrequest: PersistSubrequest, + context: &mut PersistContext, + ) -> Result { + let doc_batch = match subrequest.doc_batch.take() { + Some(doc_batch) if !doc_batch.is_empty() => doc_batch, + _ => { + warn!("received empty persist request"); + DocBatchV2::default() + } + }; + let failure = |reason: PersistFailureReason| PersistFailure { + subrequest_id: subrequest.subrequest_id, + index_uid: subrequest.index_uid.clone(), + source_id: subrequest.source_id.clone(), + reason: reason as i32, + }; + let requested_capacity = doc_batch_size(&doc_batch, context.force_commit); + let total_requested_capacity = context.reserved_capacity + requested_capacity; + let disk_used = ByteSize::b(context.wal_usage.disk_used_bytes as u64); + let memory_used = ByteSize::b(context.wal_usage.memory_used_bytes as u64); + if disk_used + total_requested_capacity > self.disk_capacity + || memory_used + total_requested_capacity > self.memory_capacity + { + rate_limited_warn!( + limit_per_min = 10, + "failed to stage persist request: WAL disk usage {}, disk capacity {}, memory \ + usage {}, memory capacity {}, requested capacity {}", + disk_used, + self.disk_capacity, + memory_used, + self.memory_capacity, + total_requested_capacity + ); + return Err(failure(PersistFailureReason::WalFull)); + } + + let Some(shard) = self + .inner + .find_most_capacity_shard_mut(subrequest.index_uid(), &subrequest.source_id) + else { + warn!( + index_uid=%subrequest.index_uid(), + source_id=%subrequest.source_id, + "no open shard found on ingester" + ); + return Err(failure(PersistFailureReason::NoShardsForSource)); + }; + shard.is_advertisable = true; + let (valid_doc_batch, parse_failures) = match validate_doc_batch(shard, doc_batch).await { + Ok(validated) => validated, + Err(error) => { + error!(queue_id=%shard.queue_id(), "failed to validate documents: {error}"); + return Err(failure(PersistFailureReason::Internal)); + } + }; + let batch_size = if parse_failures.is_empty() { + requested_capacity + } else { + doc_batch_size(&valid_doc_batch, context.force_commit) + }; + context.reserved_capacity += batch_size; + shard.queue_size += batch_size; + + Ok(StagedPersistRequest { + subrequest_id: subrequest.subrequest_id, + index_uid: subrequest.index_uid, + source_id: subrequest.source_id, + shard_id: shard.shard_id.clone(), + queue_id: shard.queue_id(), + num_docs: valid_doc_batch.num_docs() as u32, + doc_batch: valid_doc_batch, + batch_size, + parse_failures, + from_position_exclusive: shard.replication_position_inclusive.clone(), + }) + } + + pub async fn persist_subrequest( + &mut self, + staged_request: StagedPersistRequest, + context: &mut PersistContext, + ) -> Result { + let replication_position_inclusive = if staged_request.num_docs > 0 { + let append_result = append_non_empty_doc_batch( + &mut self.mrecordlog, + &staged_request.queue_id, + staged_request.doc_batch, + context.force_commit, + ) + .await; + let position = match append_result { + Ok(position) => position, + Err(append_error) => { + self.shards + .get_mut(&staged_request.queue_id) + .expect("shard should exist") + .queue_size -= staged_request.batch_size; + match append_error { + AppendDocBatchError::Io(io_error) => { + error!(queue_id=%staged_request.queue_id, "failed to persist records: {io_error}"); + context.shards_to_close.insert(staged_request.queue_id); + } + AppendDocBatchError::QueueNotFound(_) => { + error!(queue_id=%staged_request.queue_id, "failed to persist records: WAL queue not found"); + context.shards_to_delete.insert(staged_request.queue_id); + } + } + return Err(PersistFailure { + subrequest_id: staged_request.subrequest_id, + index_uid: staged_request.index_uid, + source_id: staged_request.source_id, + reason: PersistFailureReason::Internal as i32, + }); + } + }; + self.shards + .get_mut(&staged_request.queue_id) + .expect("shard should exist") + .set_replication_position_inclusive(position.clone(), Instant::now()); + position + } else { + staged_request.from_position_exclusive + }; + + Ok(PersistSuccess { + subrequest_id: staged_request.subrequest_id, + index_uid: staged_request.index_uid, + source_id: staged_request.source_id, + shard_id: Some(staged_request.shard_id), + replication_position_inclusive: Some(replication_position_inclusive), + num_persisted_docs: staged_request.num_docs, + parse_failures: staged_request.parse_failures, + }) + } + + pub fn finish_persist(&mut self, context: PersistContext) { + for queue_id in context.shards_to_close { + self.shards + .get_mut(&queue_id) + .expect("shard should exist") + .close(); + warn!("closed shard `{queue_id}` following IO error"); + } + for queue_id in context.shards_to_delete { + self.shards.remove(&queue_id); + warn!("deleted dangling shard `{queue_id}`"); + } + } + /// Reports the current WAL disk/memory usage and usage-ratio metrics against the configured /// capacity limits. Called after any operation that grows or shrinks WAL usage so that /// dashboards and alerts relying on these metrics stay accurate even when the ingester is @@ -788,6 +978,81 @@ impl WeakIngesterState { } } +pub(super) struct PersistContext { + force_commit: bool, + wal_usage: ResourceUsage, + reserved_capacity: ByteSize, + shards_to_close: HashSet, + shards_to_delete: HashSet, +} + +pub(super) struct StagedPersistRequest { + subrequest_id: u32, + index_uid: Option, + source_id: SourceId, + shard_id: ShardId, + queue_id: QueueId, + doc_batch: DocBatchV2, + batch_size: ByteSize, + num_docs: u32, + parse_failures: Vec, + from_position_exclusive: Position, +} + +async fn validate_doc_batch( + shard: &mut IngesterShard, + doc_batch: DocBatchV2, +) -> IngestV2Result<(DocBatchV2, Vec)> { + let original_batch_num_bytes = doc_batch.num_bytes() as u64; + let doc_mapper = shard.doc_mapper_opt.clone().expect("shard should be open"); + let (valid_doc_batch, parse_failures) = if shard.validate_docs { + super::doc_mapper::validate_doc_batch(doc_batch, doc_mapper).await? + } else { + (doc_batch, Vec::new()) + }; + doc_batch_metrics( + shard, + original_batch_num_bytes, + &valid_doc_batch, + &parse_failures, + ); + Ok((valid_doc_batch, parse_failures)) +} + +fn doc_batch_metrics( + shard: &mut IngesterShard, + original_batch_num_bytes: u64, + valid_doc_batch: &DocBatchV2, + parse_failures: &[ParseFailure], +) { + let valid_batch_num_bytes = valid_doc_batch.num_bytes() as u64; + if valid_doc_batch.is_empty() || !parse_failures.is_empty() { + counter!( + parent: DOCS_TOTAL, + labels: [label_values!(VALIDITY => "invalid")], + ) + .inc_by(parse_failures.len() as u64); + counter!( + parent: DOCS_BYTES_TOTAL, + labels: [label_values!(VALIDITY => "invalid")], + ) + .inc_by(original_batch_num_bytes - valid_batch_num_bytes); + } + if !valid_doc_batch.is_empty() { + counter!( + parent: DOCS_TOTAL, + labels: [label_values!(VALIDITY => "valid")], + ) + .inc_by(valid_doc_batch.num_docs() as u64); + counter!( + parent: DOCS_BYTES_TOTAL, + labels: [label_values!(VALIDITY => "valid")], + ) + .inc_by(valid_batch_num_bytes); + shard.rate_meter.update(valid_batch_num_bytes); + } +} + #[cfg(test)] mod tests { use bytesize::ByteSize; diff --git a/quickwit/quickwit-proto/protos/quickwit/ingester.proto b/quickwit/quickwit-proto/protos/quickwit/ingester.proto index cca5d243ddd..1b02624a22c 100644 --- a/quickwit/quickwit-proto/protos/quickwit/ingester.proto +++ b/quickwit/quickwit-proto/protos/quickwit/ingester.proto @@ -105,8 +105,9 @@ enum PersistFailureReason { PERSIST_FAILURE_REASON_UNSPECIFIED = 0; PERSIST_FAILURE_REASON_WAL_FULL = 4; PERSIST_FAILURE_REASON_TIMEOUT = 5; - PERSIST_FAILURE_REASON_NO_SHARDS_AVAILABLE = 6; + PERSIST_FAILURE_REASON_NO_SHARDS_FOR_SOURCE = 6; PERSIST_FAILURE_REASON_NODE_UNAVAILABLE = 7; + PERSIST_FAILURE_REASON_INTERNAL = 8; reserved 1; reserved 2; reserved 3; diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.ingest.ingester.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.ingest.ingester.rs index 8af01f90966..06b493c9c22 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.ingest.ingester.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.ingest.ingester.rs @@ -262,8 +262,9 @@ pub enum PersistFailureReason { Unspecified = 0, WalFull = 4, Timeout = 5, - NoShardsAvailable = 6, + NoShardsForSource = 6, NodeUnavailable = 7, + Internal = 8, } impl PersistFailureReason { /// String value of the enum field names used in the ProtoBuf definition. @@ -275,8 +276,9 @@ impl PersistFailureReason { Self::Unspecified => "PERSIST_FAILURE_REASON_UNSPECIFIED", Self::WalFull => "PERSIST_FAILURE_REASON_WAL_FULL", Self::Timeout => "PERSIST_FAILURE_REASON_TIMEOUT", - Self::NoShardsAvailable => "PERSIST_FAILURE_REASON_NO_SHARDS_AVAILABLE", + Self::NoShardsForSource => "PERSIST_FAILURE_REASON_NO_SHARDS_FOR_SOURCE", Self::NodeUnavailable => "PERSIST_FAILURE_REASON_NODE_UNAVAILABLE", + Self::Internal => "PERSIST_FAILURE_REASON_INTERNAL", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -285,8 +287,11 @@ impl PersistFailureReason { "PERSIST_FAILURE_REASON_UNSPECIFIED" => Some(Self::Unspecified), "PERSIST_FAILURE_REASON_WAL_FULL" => Some(Self::WalFull), "PERSIST_FAILURE_REASON_TIMEOUT" => Some(Self::Timeout), - "PERSIST_FAILURE_REASON_NO_SHARDS_AVAILABLE" => Some(Self::NoShardsAvailable), + "PERSIST_FAILURE_REASON_NO_SHARDS_FOR_SOURCE" => { + Some(Self::NoShardsForSource) + } "PERSIST_FAILURE_REASON_NODE_UNAVAILABLE" => Some(Self::NodeUnavailable), + "PERSIST_FAILURE_REASON_INTERNAL" => Some(Self::Internal), _ => None, } } diff --git a/quickwit/quickwit-proto/src/ingest/mod.rs b/quickwit/quickwit-proto/src/ingest/mod.rs index 9e6f81a26af..26fa0548935 100644 --- a/quickwit/quickwit-proto/src/ingest/mod.rs +++ b/quickwit/quickwit-proto/src/ingest/mod.rs @@ -311,10 +311,11 @@ impl From for IngestFailureReason { fn from(reason: PersistFailureReason) -> Self { match reason { PersistFailureReason::Unspecified => IngestFailureReason::Unspecified, - PersistFailureReason::NoShardsAvailable => IngestFailureReason::NoShardsAvailable, + PersistFailureReason::NoShardsForSource => IngestFailureReason::NoShardsAvailable, PersistFailureReason::WalFull => IngestFailureReason::WalFull, PersistFailureReason::Timeout => IngestFailureReason::Timeout, PersistFailureReason::NodeUnavailable => IngestFailureReason::NoShardsAvailable, + PersistFailureReason::Internal => IngestFailureReason::Internal, } } }