// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. use std::{ collections::VecDeque, fmt, sync::{ Arc, atomic::{AtomicU64, Ordering}, }, time::{Duration, Instant}, }; use collections::HashMap; use kvproto::{ kvrpcpb::{ApiVersion, Context}, metapb::RegionEpoch, }; use tikv_util::deadline::Deadline; use tokio::{sync::mpsc, time::Interval}; use txn_types::TimeStamp; use crate::{ storage, storage::{ metrics::{ ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC, ASYNC_RESOLVE_LOCK_DEQUEUE_COUNTER_VEC_STATIC, ASYNC_RESOLVE_LOCK_INQUEUE_COUNTER_VEC_STATIC, ASYNC_RESOLVE_LOCK_QUEUE_SIZE_GAUGE_VEC, AsyncResolveLockBgScheduleType, AsyncResolveLockDequeueType, AsyncResolveLockInqueueType, }, txn::{ Error, ErrorInner, commands::{Command, ResolveLockReadPhase, TypedCommand}, }, }, }; /// Stable key for batching async resolve lock signals. /// /// Region id alone is not enough because requests for different API versions, /// keyspaces, stores, or local peers must not share a resolve-lock context. #[derive(Clone, Debug, Eq, Hash, PartialEq)] pub struct AsyncResolveLockBatchKey { api_version: ApiVersion, keyspace_id: u32, store_id: u64, region_id: u64, peer_id: u64, } impl AsyncResolveLockBatchKey { /// Extracts the batch key from the request context. fn from_context(ctx: &Context) -> Self { let peer = ctx.get_peer(); Self { api_version: ctx.get_api_version(), keyspace_id: ctx.get_keyspace_id(), store_id: peer.get_store_id(), region_id: ctx.get_region_id(), peer_id: peer.get_id(), } } /// Returns the API version part of the batch key. #[cfg(test)] pub fn api_version(&self) -> ApiVersion { self.api_version } /// Returns the keyspace id part of the batch key. #[cfg(test)] pub fn keyspace_id(&self) -> u32 { self.keyspace_id } /// Returns the store id part of the batch key. #[cfg(test)] pub fn store_id(&self) -> u64 { self.store_id } /// Returns the region id part of the batch key. #[cfg(test)] pub fn region_id(&self) -> u64 { self.region_id } /// Returns the peer id part of the batch key. #[cfg(test)] pub fn peer_id(&self) -> u64 { self.peer_id } } /// Pending async resolve batch for one API version, keyspace, store, region, /// and peer. pub struct AsyncResolveLockBatch { /// Context reused when the batch is eventually converted to /// `ResolveLockReadPhase`. ctx: Context, /// Merged txn final statuses. /// The key is lock ts and the value is commit ts, where 0 means rollback. txn_status: HashMap, /// The time when the pending batch was first created. /// Merge operations do not refresh it. first_seen: Instant, } impl fmt::Debug for AsyncResolveLockBatch { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("AsyncResolveLockBatch") .field("region_id", &self.ctx.get_region_id()) .field("keyspace_id", &self.ctx.get_keyspace_id()) .field("keyspace_name", &self.ctx.get_keyspace_name()) .field("store_id", &self.ctx.get_peer().get_store_id()) .field("txn_status", &self.txn_status) .field("first_seen", &self.first_seen) .finish() } } impl AsyncResolveLockBatch { const DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE_SECS: u64 = 5 * 60; const DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE: Duration = Duration::from_secs(Self::DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE_SECS); const REQUEST_SOURCE: &'static str = "internal_async_batch_resolve_lock"; /// Creates a pending batch with a fixed first-seen timestamp. #[inline] fn new(ctx: &Context, txn_status: HashMap, first_seen: Instant) -> Self { Self { ctx: Self::new_batch_context(ctx), txn_status, first_seen, } } /// Builds the minimal context needed by background resolve-lock execution. /// /// The background task is detached from the original RPC, so the context is /// rebuilt by field category: /// - Routing and request identity are copied: `region_id`, `region_epoch`, /// `peer`, `api_version`, `keyspace_id`, `keyspace_name`, and /// `cluster_id`. They identify the target peer, key encoding, keyspace, /// and cluster for the background resolve-lock request. /// - Cache behavior follows the original request: `not_fill_cache` is /// copied so async background resolve-lock keeps the same cache-fill hint /// as ordinary region-level resolve-lock requests from the same caller. /// - Request source is replaced: `request_source` is set to the internal /// async resolve-lock source so metrics and logs can distinguish this /// background work from foreground RPCs. /// - RPC scheduling, read-mode, resource-control, and tracing fields stay /// at protobuf defaults because a merged batch can combine multiple /// foreground requests and must not synthesize these per-request choices /// from one arbitrary signal. This includes `term`, `priority`, /// `isolation_level`, `sync_log`, `record_time_stat`, `record_scan_stat`, /// `replica_read`, `resolved_locks`, `committed_locks`, /// `max_execution_duration_ms`, `applied_index`, `task_id`, `stale_read`, /// `resource_group_tag`, `resource_control_context`, `disk_full_opt`, /// `is_retry_request`, `txn_source`, `busy_threshold_ms`, /// `buckets_version`, `source_stmt`, `trace_context`, `trace_id`, and /// `trace_control_flags`. /// /// TODO: Define the resource-control attribution policy for merged /// background resolve-lock work. A single batch can combine signals from /// multiple foreground requests, so inheriting one request's resource group /// would misattribute the others. #[inline] fn new_batch_context(ctx: &Context) -> Context { let mut batch_ctx = Context::default(); batch_ctx.set_region_id(ctx.get_region_id()); batch_ctx.set_region_epoch(ctx.get_region_epoch().clone()); batch_ctx.set_peer(ctx.get_peer().clone()); batch_ctx.set_api_version(ctx.get_api_version()); batch_ctx.set_keyspace_id(ctx.get_keyspace_id()); batch_ctx.set_keyspace_name(ctx.get_keyspace_name().to_owned()); batch_ctx.set_cluster_id(ctx.get_cluster_id()); batch_ctx.set_not_fill_cache(ctx.get_not_fill_cache()); batch_ctx.set_request_source(Self::REQUEST_SOURCE.to_owned()); batch_ctx } /// Rebuilds the queue key from the batch context. #[inline] pub fn key(&self) -> AsyncResolveLockBatchKey { AsyncResolveLockBatchKey::from_context(&self.ctx) } /// Returns the number of unique txn statuses currently merged into this /// batch. #[inline] pub fn txn_count(&self) -> usize { self.txn_status.len() } /// Returns whether `incoming` can replace `current` as a strictly newer /// epoch without synthesizing a state from unrelated epoch fields. #[inline] fn is_epoch_strictly_newer(incoming: &RegionEpoch, current: &RegionEpoch) -> bool { incoming.get_conf_ver() >= current.get_conf_ver() && incoming.get_version() >= current.get_version() && (incoming.get_conf_ver() > current.get_conf_ver() || incoming.get_version() > current.get_version()) } /// Updates the stored context with a strictly newer epoch from the incoming /// request while keeping other context fields stable. #[inline] fn maybe_update_context(&mut self, ctx: &Context) { let current_epoch = self.ctx.get_region_epoch(); let incoming_epoch = ctx.get_region_epoch(); if Self::is_epoch_strictly_newer(incoming_epoch, current_epoch) { self.ctx.set_region_epoch(incoming_epoch.clone()); } } /// Merges txn final statuses into this pending batch. /// /// Duplicate `start_ts -> commit_ts` pairs are ignored. Conflicting /// statuses for the same `start_ts` are logged and the existing status /// wins. fn merge_txn_status( &mut self, ctx: &Context, txn_status: HashMap, ) -> usize { debug_assert_eq!(self.key(), AsyncResolveLockBatchKey::from_context(ctx)); self.maybe_update_context(ctx); let mut inserted = 0; for (start_ts, commit_ts) in txn_status { match self.txn_status.get(&start_ts) { Some(existing_commit_ts) if *existing_commit_ts == commit_ts => {} Some(existing_commit_ts) => { error!( "async resolve lock signal has conflicting txn status"; "key" => ?self.key(), "start_ts" => ?start_ts, "existing_status" => ?existing_commit_ts, "incoming_status" => ?commit_ts, ); } None => { self.txn_status.insert(start_ts, commit_ts); inserted += 1; } } } inserted } /// Converts this pending batch into the existing region-level resolve lock /// read phase. #[inline] pub fn into_command(self) -> TypedCommand<()> { Command::ResolveLockReadPhase(ResolveLockReadPhase { ctx: self.ctx, deadline: Deadline::from_now(Self::DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE), txn_status: self.txn_status, scan_key: None, is_background: true, }) .into() } } /// Completion token for one dispatched async resolve batch. /// /// Multiple inflight batches can exist for the same batch key, so completion /// must carry both the key and a unique id. #[derive(Clone, Debug, Eq, PartialEq)] pub struct AsyncResolveLockInflightToken { /// Batch key shared by the dispatched batch and future completion callback. key: AsyncResolveLockBatchKey, /// Per-worker unique id used to remove exactly one inflight batch. id: u64, } /// Txn statuses of one dispatched async resolve batch. #[derive(Debug)] struct AsyncResolveLockInflightBatch { /// Id matched by `AsyncResolveLockInflightToken`. id: u64, /// Time when the batch was dispatched to scheduler. started_at: Instant, /// Txn final statuses already being resolved by this inflight batch. txn_status: HashMap, } #[derive(Debug)] struct ExpiredAsyncResolveLockInflightBatch { token: AsyncResolveLockInflightToken, started_at: Instant, txn_status: HashMap, } #[derive(Clone, Default)] struct Clock { #[cfg(test)] now_fn: Option Instant + Send + Sync>>, } impl Clock { fn new() -> Clock { Self { #[cfg(test)] now_fn: None, } } #[cfg(test)] fn new_with(now_fn: Arc Instant + Send + Sync>) -> Clock { Self { now_fn: Some(now_fn), } } #[inline] fn now(&self) -> Instant { #[cfg(test)] { if let Some(now_fn) = self.now_fn.as_ref() { return now_fn(); } } Instant::now() } } /// Single-threaded in-memory queue owned by `run_async_resolve_lock_worker`. /// /// The queue has three layers: /// - `pending_by_key` stores at most one mergeable pending batch per key. /// - `region_queue` keeps pending keys in insertion order for periodic /// dispatch. /// - `blocking_set` keeps min-waited pending keys that are blocked by same-key /// inflight work. /// - `inflight_by_key` tracks dispatched batches so new duplicate txn statuses /// can be filtered while still allowing multiple inflight batches for one /// key. /// /// For one resolve item `(batch_key, start_ts)`, the state flow is: /// /// ```text /// ┌─────────────────┐ /// │ incoming signal │ /// └────────┬────────┘ /// │ /// ┌───────────────┴───────────────┐ /// │ covered by same-key inflight? │ /// └───────────┬───────────┬───────┘ /// │ yes │ no /// v v /// ┌──────────────────┐ ┌────────────────┐ /// │ drop duplicate │ │ pending_by_key │ /// └──────────────────┘ └────────┬───────┘ /// │ /// ┌────────────────────────┴────────────────────────┐ /// │ merge more statuses with the same batch key │ /// └────────────────────────┬────────────────────────┘ /// │ /// ┌────────────────────────────────────┴─────────────────────────┐ /// │ force dequeue? │ /// │ - txn count exceeds limit, or │ /// │ - cleanup tick: max wait reached │ /// └───────────────┬──────────────────────────────────┬────────────┘ /// │ yes │ no /// v v /// ┌────────────────────────┐ ┌───────────────────────────┐ /// │ schedule attempt │ │ min wait reached, no │ /// │ │ │ same-key inflight, and │ /// │ │ │ global inflight has room? │ /// └───────────┬────────────┘ └────────────┬──────────────┘ /// │ │ /// ┌───────────┴───────────┐ ┌───────┴────────┐ /// │ global inflight full? │ │ yes │ no /// └───────────┬───────┬───┘ v v /// │ yes │ no ┌────────────────┐ ┌──────────────────┐ /// v v │ schedule attempt│ │ stay pending │ /// ┌────────────┐ ┌────────────────┐ │ │ or blocking_set │ /// │ drop batch │ │ inflight_by_key│<──────────┘ └──────────────────┘ /// └────────────┘ └───────┬────────┘ /// │ /// ┌──────────┴───────────┐ /// │ background completes │ /// │ or timeout cleanup │ /// └──────────┬───────────┘ /// v /// ┌───────────┐ /// │ forgotten │ /// └───────────┘ /// ``` #[derive(Default)] pub struct AsyncResolveLockQueue { /// Clock override used by tests. clock: Clock, /// Pending merge state keyed by API version, keyspace, store, region, and /// peer. /// /// Each key has at most one pending batch. New signals for the same key /// merge into this batch until it becomes ready by time or size. pending_by_key: HashMap, /// Insertion-order index of pending batch keys and their first-seen times. /// /// A key appears at most once. When the pending batch is dispatched, the /// key is removed from this queue and from `pending_by_key` together. The /// first-seen time is duplicated here so the tick path can decide readiness /// without looking up the front batch in `pending_by_key`. region_queue: VecDeque<(AsyncResolveLockBatchKey, Instant)>, /// Min-waited pending keys that are waiting for same-key inflight work. /// /// A pending key appears in either `region_queue` or `blocking_set`, not /// both. The first-seen time is copied from the pending batch so blocked /// batches can be forced ready after `max_wait`. blocking_set: HashMap, /// Dispatched batches that have not completed yet, grouped by batch key. /// /// Multiple inflight batches are allowed for one key. This lets a large /// pending batch dispatch even if an older batch for the same region is /// still running, while still filtering duplicate or conflicting txn /// statuses. inflight_by_key: HashMap>, /// Number of unique txn statuses currently stored in all pending batches. /// /// This does not include txn statuses in inflight batches. pending_txn_count: usize, /// Number of inflight resolve lock tasks. inflight_task_count: usize, /// Monotonic id generator for inflight completion tokens. /// /// The id only needs to be unique within this worker process lifetime. next_inflight_id: u64, } impl AsyncResolveLockQueue { /// Creates an empty async resolve lock queue. pub fn new() -> Self { Self { ..Default::default() } } #[inline] pub fn pending_batch_count(&self) -> usize { self.pending_by_key.len() } /// Returns the number of unique txn statuses currently pending. #[cfg(test)] pub fn pending_txn_count(&self) -> usize { self.pending_txn_count } /// Returns the number of dispatched batches that have not completed yet. #[cfg(test)] pub fn inflight_task_count(&self) -> usize { self.inflight_task_count } /// Returns whether there is a pending batch for the given key. pub fn has_pending_batch(&self, ctx: &Context) -> bool { self.pending_by_key .contains_key(&AsyncResolveLockBatchKey::from_context(ctx)) } /// Inserts one async resolve signal into the pending queue. /// /// Returns the pending batch, if one exists, and the number of newly /// inserted txn statuses. A zero inserted count means all incoming txn /// statuses were already covered by pending or inflight state. pub fn insert_for_batch( &mut self, ctx: &Context, txn_status: HashMap, ) -> Result< ( AsyncResolveLockBatchKey, Option<&AsyncResolveLockBatch>, usize, ), Error, > { if txn_status.is_empty() { return Err(ErrorInner::Other(box_err!("txn_status is empty")).into()); } let key = AsyncResolveLockBatchKey::from_context(ctx); let txn_status = self.filter_covered_txn_status(&key, txn_status); if txn_status.is_empty() { let batch = self.pending_by_key.get(&key); return Ok((key, batch, 0)); } let inserted = if let Some(batch) = self.pending_by_key.get_mut(&key) { let inserted = batch.merge_txn_status(ctx, txn_status); self.pending_txn_count += inserted; inserted } else { let first_seen = self.clock.now(); let batch = AsyncResolveLockBatch::new(ctx, txn_status, first_seen); let inserted = batch.txn_count(); self.pending_txn_count += inserted; self.pending_by_key.insert(key.clone(), batch); self.region_queue.push_back((key.clone(), first_seen)); inserted }; let batch = self.pending_by_key.get(&key); Ok((key, batch, inserted)) } /// Removes a pending batch immediately and drops its queued key. pub fn take_pending_batch( &mut self, key: &AsyncResolveLockBatchKey, ) -> Option { let batch = self.remove_pending_batch_from_map(key)?; if self.blocking_set.remove(key).is_none() { self.region_queue .retain(|(queued_key, _first_seen)| queued_key != key); } Some(batch) } /// Marks a dispatched batch as inflight and returns its completion token. pub fn mark_inflight( &mut self, batch: &AsyncResolveLockBatch, ) -> AsyncResolveLockInflightToken { let key = batch.key(); let id = self.next_inflight_id; self.next_inflight_id = self.next_inflight_id.wrapping_add(1); self.inflight_by_key .entry(key.clone()) .or_default() .push(AsyncResolveLockInflightBatch { id, started_at: self.clock.now(), txn_status: batch.txn_status.clone(), }); self.inflight_task_count += 1; AsyncResolveLockInflightToken { key, id } } /// Takes pending batches that are ready to be dispatched. /// /// This method walks the due prefix of `region_queue` in queue order. Keys /// blocked by same-key inflight work are moved to `blocking_set`, so they /// do not make later ticks rescan the same blocked prefix. /// /// `region_queue` is also ordered by `first_seen`: a key is pushed to the /// back only when a new pending batch is created, and merging more txn /// statuses into an existing batch does not refresh `first_seen`. Because /// of that, once this scan reaches a batch newer than `now - min_wait`, all /// later batches are newer too and the scan can stop. /// /// Readiness is decided as follows: /// - A batch that has waited for less than `min_wait` stays pending. /// - After `min_wait`, a batch is ready if the global inflight limit still /// has room and there is no inflight batch with the same key. /// - If there is already inflight work with the same key, or the global /// inflight limit is full, the batch stays pending until it has waited /// for at least `max_wait`. /// /// ```text /// now.sub(max_wait) now.sub(min_wait) now /// <---------------|------------------------------------------|-----------------------------| /// ready ready not-ready /// (if no inflight batch with the same key) /// ``` /// /// Ready batches are removed from `pending_by_key` and appended to `ready`. /// Non-ready batches stay in the pending queue, so an earlier non-ready /// batch does not prevent later ready batches from being taken. pub fn take_ready_batches_from_pending( &mut self, min_wait: Duration, max_wait: Duration, max_inflight_batch_count: usize, ready: &mut Vec, ) { let max_wait = if max_wait < min_wait { min_wait } else { max_wait }; let now = self.clock.now(); let Some(min_wait_instant) = now.checked_sub(min_wait) else { debug_assert!(false, "invalid min_wait: {:?}", min_wait); return; }; let max_wait_instant = now.checked_sub(max_wait); while let Some((_, first_seen)) = self.region_queue.front() { // Too new to dispatch, even when no same-key batch is inflight. if *first_seen > min_wait_instant { // Later queued batches have greater or equal `first_seen`, so // they cannot be ready either. break; } let (key, first_seen) = self .region_queue .pop_front() .expect("front entry must exist"); let has_inflight = self.inflight_by_key.contains_key(&key); let max_wait_reached = max_wait_instant.is_some_and(|inst| first_seen <= inst); let inflight_full = self.inflight_task_count.saturating_add(ready.len()) >= max_inflight_batch_count; if inflight_full && !max_wait_reached { self.region_queue.push_front((key, first_seen)); break; } // Same-key work is still inflight. Keep this batch pending until // `max_wait` forces it to be dispatched. if has_inflight && !max_wait_reached { let previous = self.blocking_set.insert(key, first_seen); debug_assert!( previous.is_none(), "blocked key must not already exist in blocking_set" ); continue; } let Some(batch) = self.remove_pending_batch_from_map(&key) else { debug_assert!(false, "pending batch for queued key must exist"); continue; }; let dequeue_type = if max_wait_reached { AsyncResolveLockDequeueType::reason_max_wait } else { AsyncResolveLockDequeueType::reason_min_wait }; ASYNC_RESOLVE_LOCK_DEQUEUE_COUNTER_VEC_STATIC .get(dequeue_type) .inc(); ready.push(batch); } if !self.blocking_set.is_empty() { self.take_ready_batches_from_blocking_set( ready, max_wait_instant, max_inflight_batch_count, ); } } fn take_ready_batches_from_blocking_set( &mut self, ready: &mut Vec, max_wait_instant: Option, max_inflight_batch_count: usize, ) { let inflight_by_key = &self.inflight_by_key; let pending_by_key = &mut self.pending_by_key; let pending_txn_count = &mut self.pending_txn_count; let inflight_task_count = self.inflight_task_count; self.blocking_set.retain(|key, first_seen| { let has_inflight = inflight_by_key.contains_key(key); let max_wait_reached = max_wait_instant.is_some_and(|inst| *first_seen <= inst); let inflight_full = inflight_task_count.saturating_add(ready.len()) >= max_inflight_batch_count; if (has_inflight || inflight_full) && !max_wait_reached { return true; } let Some(batch) = pending_by_key.remove(key) else { debug_assert!(false, "pending batch for blocked key must exist"); return false; }; *pending_txn_count -= batch.txn_count(); let dequeue_type = if max_wait_reached { AsyncResolveLockDequeueType::reason_max_wait } else { AsyncResolveLockDequeueType::reason_inflight_finished }; ASYNC_RESOLVE_LOCK_DEQUEUE_COUNTER_VEC_STATIC .get(dequeue_type) .inc(); ready.push(batch); false }); } /// Removes one completed inflight batch identified by its completion token. fn complete_inflight( &mut self, token: &AsyncResolveLockInflightToken, ) -> Option { let (inflight_batch, remove_key) = { let inflight_batches = self.inflight_by_key.get_mut(&token.key)?; let index = inflight_batches .iter() .position(|batch| batch.id == token.id)?; let inflight_batch = inflight_batches.remove(index); (inflight_batch, inflight_batches.is_empty()) }; if remove_key { self.inflight_by_key.remove(&token.key); } self.inflight_task_count -= 1; Some(inflight_batch) } fn remove_expired_inflight( &mut self, now: Instant, timeout: Duration, ) -> Vec { let mut expired = Vec::new(); self.inflight_by_key.retain(|key, inflight_batches| { inflight_batches.retain(|batch| { if now.saturating_duration_since(batch.started_at) <= timeout { return true; } expired.push(ExpiredAsyncResolveLockInflightBatch { token: AsyncResolveLockInflightToken { key: key.clone(), id: batch.id, }, started_at: batch.started_at, txn_status: batch.txn_status.clone(), }); false }); !inflight_batches.is_empty() }); debug_assert!( self.inflight_task_count >= expired.len(), "inflight task count must cover expired inflight batches" ); self.inflight_task_count = self.inflight_task_count.saturating_sub(expired.len()); expired } /// Removes a pending batch by key and updates the aggregate pending txn /// count. fn remove_pending_batch_from_map( &mut self, key: &AsyncResolveLockBatchKey, ) -> Option { let batch = self.pending_by_key.remove(key)?; self.pending_txn_count -= batch.txn_count(); Some(batch) } #[cfg(test)] fn set_clock(&mut self, clock: Clock) { self.clock = clock } /// Drops txn statuses already covered by inflight batches of the same key. /// /// If the same `start_ts` is inflight with a different final status, the /// incoming status is treated as conflicting and is ignored. fn filter_covered_txn_status( &self, key: &AsyncResolveLockBatchKey, txn_status: HashMap, ) -> HashMap { let inflight_batches = self.inflight_by_key.get(key); txn_status .into_iter() .filter_map(|(start_ts, commit_ts)| { if let Some(inflight_batches) = inflight_batches { for inflight_batch in inflight_batches { if let Some(existing_commit_ts) = inflight_batch.txn_status.get(&start_ts) { if *existing_commit_ts != commit_ts { error!( "async resolve lock signal conflicts with inflight txn status"; "key" => ?key, "start_ts" => ?start_ts, "inflight_status" => ?existing_commit_ts, "incoming_status" => ?commit_ts, ); } return None; } } } Some((start_ts, commit_ts)) }) .collect() } } pub struct AsyncResolveLockTask { ctx: Context, txn_status: HashMap, } impl fmt::Debug for AsyncResolveLockTask { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("AsyncResolveLockTask") .field("region_id", &self.ctx.get_region_id()) .field("keyspace_id", &self.ctx.get_keyspace_id()) .field("keyspace_name", &self.ctx.get_keyspace_name()) .field("store_id", &self.ctx.get_peer().get_store_id()) .field("txn_status", &self.txn_status) .finish() } } impl AsyncResolveLockTask { /// Builds a task from an async region-level resolve signal. pub fn new(ctx: Context, txn_status: HashMap) -> Self { Self { ctx, txn_status } } /// Splits the task so the queue can merge txn statuses while preserving the /// request context. fn into_parts(self) -> (Context, HashMap) { (self.ctx, self.txn_status) } } type ResolveLockBatchCallback = Box) + Send>; type AsyncResolveLockInflightCompleteResult = (AsyncResolveLockInflightToken, storage::Result<()>); #[derive(Debug)] struct RateLimitLogState { /// Last time the event emitted a log. last_log_at: Instant, /// Number of logs skipped since `last_log_at` due to rate limiting. skipped_log_count: u64, } pub struct BatchAsyncResolveWorker { task_rx: mpsc::Receiver, inflight_complete_tx: mpsc::UnboundedSender, inflight_complete_rx: mpsc::UnboundedReceiver, queue: AsyncResolveLockQueue, rate_limited_log_state_by_event_name: HashMap<&'static str, RateLimitLogState>, last_inflight_cleanup_at: Option, pending_batch_wait_us: Arc, pending_batch_max_wait: Duration, pending_batch_max_txn_count: usize, max_pending_batch_count: usize, max_inflight_batch_count: usize, clock: Clock, } impl BatchAsyncResolveWorker { const TASK_DISPATCH_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_TICKER_INTERVAL: Duration = Duration::from_millis(20); const DEFAULT_PENDING_BATCH_MAX_WAIT: Duration = Duration::from_millis(5000); const DEFAULT_INFLIGHT_TIMEOUT: Duration = Duration::from_secs( // Keep the inflight record slightly longer than the command deadline // so the scheduler's DeadlineExceeded callback can complete it first. AsyncResolveLockBatch::DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE_SECS + 1, ); const DEFAULT_INFLIGHT_CLEANUP_INTERVAL: Duration = Duration::from_secs(5); const DEFAULT_PENDING_BATCH_MAX_TXN_COUNT: usize = 512; const DEFAULT_MAX_PENDING_BATCH_COUNT: usize = 20 * 1024; const DEFAULT_MAX_INFLIGHT_BATCH_COUNT: usize = 2048; const DEFAULT_RATE_LIMITED_LOG_INTERVAL: Duration = Duration::from_secs(1); pub fn new( pending_batch_wait_us: Arc, ) -> (Self, mpsc::Sender) { Self::new_with_params( Self::TASK_DISPATCH_CHANNEL_CAPACITY, pending_batch_wait_us, Self::DEFAULT_PENDING_BATCH_MAX_WAIT, Self::DEFAULT_PENDING_BATCH_MAX_TXN_COUNT, Self::DEFAULT_MAX_PENDING_BATCH_COUNT, Self::DEFAULT_MAX_INFLIGHT_BATCH_COUNT, ) } fn new_with_params( task_channel_cap: usize, pending_batch_wait_us: Arc, pending_batch_max_wait: Duration, pending_batch_max_txn_count: usize, max_pending_batch_count: usize, max_inflight_batch_count: usize, ) -> (Self, mpsc::Sender) { let (task_tx, task_rx) = mpsc::channel(task_channel_cap); let (inflight_complete_tx, inflight_complete_rx) = mpsc::unbounded_channel(); ( Self { clock: Clock::new(), task_rx, queue: AsyncResolveLockQueue::new(), inflight_complete_tx, inflight_complete_rx, pending_batch_wait_us, pending_batch_max_wait, pending_batch_max_txn_count, max_pending_batch_count, max_inflight_batch_count, rate_limited_log_state_by_event_name: HashMap::default(), last_inflight_cleanup_at: None, }, task_tx, ) } fn pending_batch_wait(&self) -> Duration { Duration::from_micros(self.pending_batch_wait_us.load(Ordering::Relaxed)) } pub async fn loop_with_scheduler(mut self, scheduler: S) where S: Fn(AsyncResolveLockBatch, ResolveLockBatchCallback) + Send + Sync + 'static, { let mut ticker = tokio::time::interval(Self::DEFAULT_TICKER_INTERVAL); let mut ready_tmp_buff = Vec::with_capacity(32); loop { if !self .loop_once(&scheduler, &mut ticker, &mut ready_tmp_buff) .await { break; } } } pub async fn loop_once( &mut self, scheduler: &S, ticker: &mut Interval, ready_tmp_buf: &mut Vec, ) -> bool where S: Fn(AsyncResolveLockBatch, ResolveLockBatchCallback) + Send + Sync + ?Sized, { tokio::select! { task = self.task_rx.recv() => { let Some(task) = task else { return false }; let (ctx, txn_status) = task.into_parts(); let pending_batch_count = self.queue.pending_batch_count(); let max_pending_batch_count = self.max_pending_batch_count; let pending_batch_max_txn_count = self.pending_batch_max_txn_count; let has_pending_batch = self.queue.has_pending_batch(&ctx); if pending_batch_count >= max_pending_batch_count && !has_pending_batch { // This is a rare memory-protection guard for the // best-effort background path. A foreground success only // means the signal reached the worker channel; the worker // may still drop it later if accepting a new batch would // grow the pending queue beyond the safety limit. ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::task_dropped) .inc(); self.log_rate_limited("batch_exceeds_limit", |skipped_log_count| { warn!( "async resolve lock worker pending batch count exceeds the limit, dropping new task"; "pending_batch_count" => pending_batch_count, "max_pending_batch_count" => max_pending_batch_count, "region_id" => ctx.get_region_id(), "txn_status" => ?txn_status, "skipped_log_count" => skipped_log_count, ); }); return true; } match self.queue.insert_for_batch(&ctx, txn_status) { Ok((key, batch, inserted)) => { let inqueue_type = if inserted == 0 { AsyncResolveLockInqueueType::dedup } else if batch .map(|batch| inserted == batch.txn_count()) .unwrap_or(false) { AsyncResolveLockInqueueType::inserted } else { AsyncResolveLockInqueueType::merged }; ASYNC_RESOLVE_LOCK_INQUEUE_COUNTER_VEC_STATIC .get(inqueue_type) .inc(); if inserted > 0 && batch .map(|batch| batch.txn_count() > pending_batch_max_txn_count) .unwrap_or(false) { if let Some(batch) = self.queue.take_pending_batch(&key) { ASYNC_RESOLVE_LOCK_DEQUEUE_COUNTER_VEC_STATIC .get(AsyncResolveLockDequeueType::reason_txn_count) .inc(); self.schedule_resolve_lock_command(scheduler, batch); } self.update_all_queue_size_metrics(); } else { self.update_pending_queue_size_metrics(); } }, Err(err) => { self.log_rate_limited("queue_insert_failed", |skipped_log_count| { warn!( "failed to insert async resolve lock task into queue"; "region" => ctx.get_region_id(), "error" => ?err, "skipped_log_count" => skipped_log_count, ); }); }, } }, completed = self.inflight_complete_rx.recv() => { let Some((token, result)) = completed else { return false; }; let completed_batch = self.queue.complete_inflight(&token); if completed_batch.is_none() { ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::inflight_missing) .inc(); self.log_rate_limited( "inflight_missing", |skipped_log_count| { warn!( "async resolve lock inflight completion is not found"; "inflight_token" => ?token, "skipped_log_count" => skipped_log_count, ); }, ); } if result.is_err() { ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::task_failed) .inc(); fail::fail_point!("async_batch_resolve_lock_background_task_failed"); } self.update_inflight_queue_size_metric(); }, _ = ticker.tick() => { self.maybe_cleanup_resolve_queue(); ready_tmp_buf.clear(); self.queue.take_ready_batches_from_pending( self.pending_batch_wait(), self.pending_batch_max_wait, self.max_inflight_batch_count, ready_tmp_buf, ); for batch in ready_tmp_buf.drain(..) { self.schedule_resolve_lock_command(scheduler, batch); } self.update_all_queue_size_metrics(); } } true } fn maybe_cleanup_resolve_queue(&mut self) { let now = self.clock.now(); if let Some(last_cleanup_at) = self.last_inflight_cleanup_at { if now.saturating_duration_since(last_cleanup_at) < Self::DEFAULT_INFLIGHT_CLEANUP_INTERVAL { return; } } self.last_inflight_cleanup_at = Some(now); for expired in self .queue .remove_expired_inflight(now, Self::DEFAULT_INFLIGHT_TIMEOUT) { ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::inflight_expired) .inc(); self.log_rate_limited("inflight_expired", |skipped_log_count| { warn!( "async resolve lock inflight batch expired and is force removed"; "inflight_token" => ?expired.token, "elapsed" => ?now.saturating_duration_since(expired.started_at), "timeout" => ?Self::DEFAULT_INFLIGHT_TIMEOUT, "txn_status" => ?expired.txn_status, "skipped_log_count" => skipped_log_count, ); }); } } fn log_rate_limited(&mut self, event_name: &'static str, log: F) where F: FnOnce(u64), { let now = self.clock.now(); let Some(state) = self .rate_limited_log_state_by_event_name .get_mut(event_name) else { self.rate_limited_log_state_by_event_name.insert( event_name, RateLimitLogState { last_log_at: now, skipped_log_count: 0, }, ); log(0); return; }; if now.saturating_duration_since(state.last_log_at) < Self::DEFAULT_RATE_LIMITED_LOG_INTERVAL { state.skipped_log_count = state.skipped_log_count.saturating_add(1); return; } let skipped_log_count_to_report = state.skipped_log_count; state.last_log_at = now; state.skipped_log_count = 0; log(skipped_log_count_to_report); } fn schedule_resolve_lock_command(&mut self, scheduler: &S, batch: AsyncResolveLockBatch) where S: Fn(AsyncResolveLockBatch, ResolveLockBatchCallback) + Send + Sync + ?Sized, { if self.queue.inflight_task_count >= self.max_inflight_batch_count { let inflight_task_count = self.queue.inflight_task_count; let max_inflight_batch_count = self.max_inflight_batch_count; ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::task_dropped) .inc(); self.log_rate_limited("inflight_full", |skipped_log_count| { warn!( "failed to schedule async resolve lock because inflight batch count reaches the limit"; "inflight_batch_count" => inflight_task_count, "max_inflight_batch_count" => max_inflight_batch_count, "txn_status" => ?batch.txn_status, "skipped_log_count" => skipped_log_count, ); }); return; } ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::task) .inc(); ASYNC_RESOLVE_LOCK_BG_SCHEDULE_COUNTER_VEC_STATIC .get(AsyncResolveLockBgScheduleType::txn) .inc_by(batch.txn_count() as u64); let inflight_token = self.queue.mark_inflight(&batch); let complete_tx = self.inflight_complete_tx.clone(); fail::fail_point!("before_async_batch_resolve_lock_schedule"); scheduler( batch, Box::new(move |result| { let _ = complete_tx.send((inflight_token, result)); // if complete_tx is closed, the worker must have been stopped, // so just drop the result. }), ); } fn update_pending_queue_size_metrics(&self) { ASYNC_RESOLVE_LOCK_QUEUE_SIZE_GAUGE_VEC .pending_batch .set(self.queue.pending_by_key.len() as i64); ASYNC_RESOLVE_LOCK_QUEUE_SIZE_GAUGE_VEC .pending_txn .set(self.queue.pending_txn_count as i64); } fn update_inflight_queue_size_metric(&self) { ASYNC_RESOLVE_LOCK_QUEUE_SIZE_GAUGE_VEC .inflight_batch .set(self.queue.inflight_task_count as i64); } fn update_blocked_queue_size_metric(&self) { ASYNC_RESOLVE_LOCK_QUEUE_SIZE_GAUGE_VEC .blocked_batch .set(self.queue.blocking_set.len() as i64); } fn update_all_queue_size_metrics(&self) { self.update_pending_queue_size_metrics(); self.update_blocked_queue_size_metric(); self.update_inflight_queue_size_metric(); } #[cfg(test)] fn set_clock(&mut self, clock: Clock) { self.clock = clock.clone(); self.queue.set_clock(clock); } } pub fn duration_to_micros(duration: Duration) -> u64 { duration.as_micros().min(u64::MAX as u128) as u64 } #[cfg(test)] mod tests { use std::sync::Mutex; use kvproto::{kvrpcpb::Context_oneof_keyspace, metapb::Peer}; use protobuf::SingularPtrField; use super::*; use crate::storage::{metrics::CommandKind, txn::commands::Command}; fn region_epoch(conf_ver: u64, version: u64) -> RegionEpoch { RegionEpoch { conf_ver, version, ..Default::default() } } fn peer(store_id: u64, peer_id: u64) -> Peer { Peer { id: peer_id, store_id, ..Default::default() } } fn new_test_req_ctx(region_id: u64) -> Context { Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "ks42".to_owned(), region_id, region_epoch: SingularPtrField::some(region_epoch(2, 3)), peer: SingularPtrField::some(peer(1001, 3003)), ..Default::default() } } fn ts(ts: u64) -> TimeStamp { TimeStamp::new(ts) } fn txn_status(pairs: &[(u64, u64)]) -> HashMap { let mut txn_status = HashMap::default(); for (start_ts, commit_ts) in pairs { txn_status.insert(ts(*start_ts), ts(*commit_ts)); } txn_status } fn set_queue_now(queue: &mut AsyncResolveLockQueue, now: Instant) { queue.set_clock(Clock::new_with(Arc::new(move || now))); } fn queued_front_key(queue: &AsyncResolveLockQueue) -> Option<&AsyncResolveLockBatchKey> { queue.region_queue.front().map(|(key, _first_seen)| key) } fn queued_back_key(queue: &AsyncResolveLockQueue) -> Option<&AsyncResolveLockBatchKey> { queue.region_queue.back().map(|(key, _first_seen)| key) } fn queued_front_first_seen(queue: &AsyncResolveLockQueue) -> Option { queue .region_queue .front() .map(|(_key, first_seen)| *first_seen) } fn blocked_first_seen( queue: &AsyncResolveLockQueue, key: &AsyncResolveLockBatchKey, ) -> Option { queue.blocking_set.get(key).copied() } fn set_worker_now(worker: &mut BatchAsyncResolveWorker, now: Instant) { worker.set_clock(Clock::new_with(Arc::new(move || now))); } async fn new_blocked_ticker() -> Interval { let mut ticker = tokio::time::interval(Duration::from_secs(3600)); // `tokio::time::interval` returns one immediately-ready tick first, so // consume it here to keep `loop_once` tests on the `task_rx.recv()` path. ticker.tick().await; ticker } fn assert_epoch(epoch: &RegionEpoch, conf_ver: u64, version: u64) { assert_eq!(epoch.get_conf_ver(), conf_ver); assert_eq!(epoch.get_version(), version); } #[test] fn test_batch_key_uses_api_keyspace_store_region_peer() { let ctx = new_test_req_ctx(2002); let key = AsyncResolveLockBatchKey::from_context(&ctx); assert_eq!(key.api_version(), ApiVersion::V2); assert_eq!(key.keyspace_id(), 42); assert_eq!(key.store_id(), 1001); assert_eq!(key.region_id(), 2002); assert_eq!(key.peer_id(), 3003); } #[test] fn test_batch_context_keeps_required_fields_and_normalizes_background_fields() { let ctx = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "ks42".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(4, 5)), peer: SingularPtrField::some(peer(1001, 3003)), term: 99, buckets_version: 88, cluster_id: 666, not_fill_cache: true, request_source: "foreground_resolve_lock".to_owned(), ..Default::default() }; let first_seen = Instant::now(); let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), first_seen); let batch_ctx = &batch.ctx; assert_eq!(batch.key(), AsyncResolveLockBatchKey::from_context(&ctx)); assert_eq!(batch.txn_count(), 1); assert_eq!(batch.first_seen, first_seen); assert_eq!(batch_ctx.get_region_id(), 2002); assert_epoch(batch_ctx.get_region_epoch(), 4, 5); assert_eq!(batch_ctx.get_peer().get_store_id(), 1001); assert_eq!(batch_ctx.get_peer().get_id(), 3003); assert_eq!(batch_ctx.get_api_version(), ApiVersion::V2); assert_eq!(batch_ctx.get_keyspace_id(), 42); assert_eq!(batch_ctx.get_keyspace_name(), "ks42"); assert_eq!(batch_ctx.get_cluster_id(), 666); assert!(batch_ctx.get_not_fill_cache()); assert_eq!( batch_ctx.get_request_source(), "internal_async_batch_resolve_lock" ); assert_eq!(batch_ctx.get_term(), 0); assert_eq!(batch_ctx.buckets_version, 0); } #[test] fn test_epoch_comparison_requires_componentwise_strictly_newer() { let current = region_epoch(2, 3); assert!(AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(2, 4), ¤t )); assert!(AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(3, 3), ¤t )); assert!(AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(3, 4), ¤t )); assert!(!AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(2, 3), ¤t )); assert!(!AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(2, 2), ¤t )); assert!(!AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(1, 4), ¤t )); assert!(!AsyncResolveLockBatch::is_epoch_strictly_newer( ®ion_epoch(3, 2), ¤t )); } #[test] fn test_newer_epoch_updates_only_region_epoch() { let ctx = new_test_req_ctx(2002); let mut batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), Instant::now()); let incoming = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "changed".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(3, 4)), peer: SingularPtrField::some(peer(1001, 3003)), term: 99, cluster_id: 666, not_fill_cache: false, request_source: "foreground_resolve_lock".to_owned(), ..Default::default() }; batch.maybe_update_context(&incoming); assert_epoch(batch.ctx.get_region_epoch(), 3, 4); assert_eq!(batch.ctx.get_keyspace_name(), "ks42"); assert!(!batch.ctx.get_not_fill_cache()); assert_eq!( batch.ctx.get_request_source(), "internal_async_batch_resolve_lock" ); assert_eq!(batch.ctx.get_term(), 0); let incomparable = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "changed".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(4, 3)), peer: SingularPtrField::some(peer(1001, 3003)), ..Default::default() }; batch.maybe_update_context(&incomparable); assert_epoch(batch.ctx.get_region_epoch(), 3, 4); let older = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "changed".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(3, 3)), peer: SingularPtrField::some(peer(1001, 3003)), ..Default::default() }; batch.maybe_update_context(&older); assert_epoch(batch.ctx.get_region_epoch(), 3, 4); } #[test] fn test_merge_txn_status_inserts_duplicates_and_conflicts() { let ctx = new_test_req_ctx(2002); let mut batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(9, 19)]), Instant::now()); assert_eq!( batch.merge_txn_status(&ctx, txn_status(&[(10, 20), (11, 0)])), 2 ); assert_eq!(batch.txn_count(), 3); assert_eq!( batch.merge_txn_status(&ctx, txn_status(&[(10, 20), (12, 30)])), 1 ); assert_eq!(batch.txn_count(), 4); assert_eq!(batch.merge_txn_status(&ctx, txn_status(&[(10, 21)])), 0); assert_eq!(batch.txn_count(), 4); assert_eq!(batch.txn_status.get(&ts(9)), Some(&ts(19))); assert_eq!(batch.txn_status.get(&ts(10)), Some(&ts(20))); } #[test] fn test_into_command_builds_resolve_lock_read_phase() { let ctx = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "ks42".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(3, 4)), peer: SingularPtrField::some(peer(1001, 3003)), term: 99, cluster_id: 666, not_fill_cache: false, request_source: "foreground_resolve_lock".to_owned(), ..Default::default() }; let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20), (11, 0)]), Instant::now()); let typed_cmd = batch.into_command(); assert_eq!(typed_cmd.cmd.tag(), CommandKind::resolve_lock_background); let Command::ResolveLockReadPhase(cmd) = typed_cmd.cmd else { panic!("expect resolve lock read phase command"); }; assert_eq!(cmd.txn_status, txn_status(&[(10, 20), (11, 0)])); assert!(cmd.scan_key.is_none()); let deadline_left = cmd.deadline.check().unwrap(); assert!(deadline_left <= AsyncResolveLockBatch::DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE); assert!( deadline_left > AsyncResolveLockBatch::DEFAULT_BACKGROUND_RESOLVE_LOCK_DEADLINE - Duration::from_secs(1) ); assert_eq!(cmd.ctx.get_region_id(), 2002); assert_epoch(cmd.ctx.get_region_epoch(), 3, 4); assert_eq!(cmd.ctx.get_peer().get_store_id(), 1001); assert_eq!(cmd.ctx.get_peer().get_id(), 3003); assert_eq!(cmd.ctx.get_api_version(), ApiVersion::V2); assert_eq!(cmd.ctx.get_keyspace_id(), 42); assert_eq!(cmd.ctx.get_keyspace_name(), "ks42"); assert_eq!(cmd.ctx.get_cluster_id(), 666); assert!(!cmd.ctx.get_not_fill_cache()); assert_eq!( cmd.ctx.get_request_source(), "internal_async_batch_resolve_lock" ); assert_eq!(cmd.ctx.get_term(), 0); } #[test] fn test_queue_insert_for_batch_creates_and_merges_pending_batch() { let ctx = new_test_req_ctx(2002); let first_seen1 = Instant::now(); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, first_seen1); assert!(!queue.has_pending_batch(&ctx)); assert_eq!(queue.pending_txn_count(), 0); // case: inserting the first task creates one pending batch. let (key, txn_count, batch_first_seen, inserted) = { let (key, batch, inserted) = queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); let batch = batch.unwrap(); (key, batch.txn_count(), batch.first_seen, inserted) }; assert_eq!(key, AsyncResolveLockBatchKey::from_context(&ctx)); assert_eq!(txn_count, 1); assert_eq!(inserted, 1); assert_eq!(batch_first_seen, first_seen1); assert_eq!(queue.pending_batch_count(), 1); assert!(queue.has_pending_batch(&ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queue.region_queue.len(), 1); assert_eq!(queued_front_key(&queue), Some(&key)); assert_eq!(queued_front_first_seen(&queue), Some(first_seen1)); // case: inserting the same region with a new txn status merges into the // existing batch. let first_seen2 = first_seen1 + Duration::from_secs(10); set_queue_now(&mut queue, first_seen2); let (merged_key, merged_txn_count, merged_first_seen, inserted) = { let (merged_key, batch, inserted) = queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap(); let batch = batch.unwrap(); (merged_key, batch.txn_count(), batch.first_seen, inserted) }; assert_eq!(merged_key, key); assert_eq!(merged_txn_count, 2); assert_eq!(inserted, 1); assert_eq!(merged_first_seen, first_seen1); assert_eq!(queue.pending_batch_count(), 1); assert!(queue.has_pending_batch(&ctx)); assert_eq!(queue.pending_txn_count(), 2); assert_eq!(queue.region_queue.len(), 1); // case: inserting only duplicate txn status does not create new pending work. let (dedup_key, dedup_txn_count, inserted) = { let (dedup_key, batch, inserted) = queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); let batch = batch.unwrap(); (dedup_key, batch.txn_count(), inserted) }; assert_eq!(dedup_key, key); assert_eq!(dedup_txn_count, 2); assert_eq!(inserted, 0); assert_eq!(queue.pending_batch_count(), 1); assert!(queue.has_pending_batch(&ctx)); assert_eq!(queue.pending_txn_count(), 2); assert_eq!(queue.region_queue.len(), 1); // case: inserting a different region creates a separate pending batch. let other_region_ctx = new_test_req_ctx(2003); assert!(!queue.has_pending_batch(&other_region_ctx)); let (other_region_key, other_region_txn_count, other_region_first_seen, inserted) = { let (other_region_key, batch, inserted) = queue .insert_for_batch(&other_region_ctx, txn_status(&[(20, 40)])) .unwrap(); let batch = batch.unwrap(); ( other_region_key, batch.txn_count(), batch.first_seen, inserted, ) }; assert_eq!( other_region_key, AsyncResolveLockBatchKey::from_context(&other_region_ctx) ); assert_ne!(other_region_key, key); assert_eq!(other_region_txn_count, 1); assert_eq!(inserted, 1); assert_eq!(other_region_first_seen, first_seen2); assert_eq!(queue.pending_batch_count(), 2); assert!(queue.has_pending_batch(&ctx)); assert!(queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 3); assert_eq!(queue.region_queue.len(), 2); assert_eq!(queued_front_key(&queue), Some(&key)); assert_eq!(queued_back_key(&queue), Some(&other_region_key)); // case: inserting the same region with a newer epoch updates the merged batch // context. let newer_epoch_ctx = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "ks42".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(3, 4)), peer: SingularPtrField::some(peer(1001, 3003)), ..Default::default() }; let (newer_epoch_key, newer_epoch_txn_count, inserted) = { let (newer_epoch_key, batch, inserted) = queue .insert_for_batch(&newer_epoch_ctx, txn_status(&[(12, 30)])) .unwrap(); let batch = batch.unwrap(); assert_epoch(batch.ctx.get_region_epoch(), 3, 4); (newer_epoch_key, batch.txn_count(), inserted) }; assert_eq!(newer_epoch_key, key); assert_eq!(newer_epoch_txn_count, 3); assert_eq!(inserted, 1); assert_eq!(queue.pending_batch_count(), 2); assert!(queue.has_pending_batch(&newer_epoch_ctx)); assert_eq!(queue.pending_txn_count(), 4); assert_eq!(queue.region_queue.len(), 2); // case: inserting the same region with an older epoch keeps the newer batch // context. let older_epoch_ctx = Context { api_version: ApiVersion::V2, keyspace: Some(Context_oneof_keyspace::KeyspaceId(42)), keyspace_name: "ks42".to_owned(), region_id: 2002, region_epoch: SingularPtrField::some(region_epoch(2, 2)), peer: SingularPtrField::some(peer(1001, 3003)), ..Default::default() }; let (older_epoch_key, older_epoch_txn_count, inserted) = { let (older_epoch_key, batch, inserted) = queue .insert_for_batch(&older_epoch_ctx, txn_status(&[(13, 0)])) .unwrap(); let batch = batch.unwrap(); assert_epoch(batch.ctx.get_region_epoch(), 3, 4); (older_epoch_key, batch.txn_count(), inserted) }; assert_eq!(older_epoch_key, key); assert_eq!(older_epoch_txn_count, 4); assert_eq!(inserted, 1); assert_eq!(queue.pending_batch_count(), 2); assert!(queue.has_pending_batch(&older_epoch_ctx)); assert_eq!(queue.pending_txn_count(), 5); assert_eq!(queue.region_queue.len(), 2); } #[test] fn test_queue_insert_for_batch_dedups_inflight_without_pending_batch() { let ctx = new_test_req_ctx(2002); let first_seen = Instant::now(); let mut queue = AsyncResolveLockQueue::new(); let inflight_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), first_seen); let inflight_token = queue.mark_inflight(&inflight_batch); // case: an incoming txn status already covered by inflight work does not // create a pending batch, so there is no batch reference to return. let (key, batch, inserted) = queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); assert_eq!(key, inflight_token.key); assert!(batch.is_none()); assert_eq!(inserted, 0); assert_eq!(queue.pending_batch_count(), 0); assert!(!queue.has_pending_batch(&ctx)); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); assert_eq!(queue.inflight_task_count(), 1); } #[test] fn test_queue_insert_for_batch_accepts_completed_resolve_again() { let ctx = new_test_req_ctx(2002); let first_seen = Instant::now(); let mut queue = AsyncResolveLockQueue::new(); let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), first_seen); let token = queue.mark_inflight(&batch); let _completed = queue.complete_inflight(&token).unwrap(); // case: once the previous resolve is no longer inflight, the same // signal is accepted again. A later prewrite for the same transaction // may have written a new lock after the previous background scan. let (_key, batch, inserted) = queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); assert_eq!(inserted, 1); assert_eq!(batch.unwrap().txn_count(), 1); assert_eq!(queue.pending_batch_count(), 1); } #[test] fn test_queue_take_pending_batch_removes_batch_and_updates_pending_state() { let ctx = new_test_req_ctx(2002); let other_region_ctx = new_test_req_ctx(2003); let first_seen = Instant::now(); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, first_seen); // Prepare two pending batches so taking one must preserve the other. let key = queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap() .0; let other_region_key = queue .insert_for_batch(&other_region_ctx, txn_status(&[(20, 40)])) .unwrap() .0; assert_eq!(queue.pending_batch_count(), 2); assert!(queue.has_pending_batch(&ctx)); assert!(queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 3); assert_eq!(queue.region_queue.len(), 2); // case: taking an existing key returns its batch and removes its queued key. let batch = queue.take_pending_batch(&key).unwrap(); assert_eq!(batch.key(), key); assert_eq!(batch.txn_count(), 2); assert_eq!(batch.txn_status.get(&ts(10)), Some(&ts(20))); assert_eq!(batch.txn_status.get(&ts(11)), Some(&ts(0))); assert_eq!(queue.pending_batch_count(), 1); assert!(!queue.has_pending_batch(&ctx)); assert!(queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queue.region_queue.len(), 1); assert_eq!(queued_front_key(&queue), Some(&other_region_key)); // case: taking the same key again returns None and leaves pending state // unchanged. assert!(queue.take_pending_batch(&key).is_none()); assert_eq!(queue.pending_batch_count(), 1); assert!(!queue.has_pending_batch(&ctx)); assert!(queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queue.region_queue.len(), 1); // case: taking a missing key returns None and leaves pending state unchanged. let missing_ctx = new_test_req_ctx(2004); let missing_key = AsyncResolveLockBatchKey::from_context(&missing_ctx); assert!(queue.take_pending_batch(&missing_key).is_none()); assert_eq!(queue.pending_batch_count(), 1); assert!(!queue.has_pending_batch(&missing_ctx)); assert!(queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queue.region_queue.len(), 1); // case: taking the last remaining key leaves the pending queue empty. let batch = queue.take_pending_batch(&other_region_key).unwrap(); assert_eq!(batch.key(), other_region_key); assert_eq!(batch.txn_count(), 1); assert_eq!(queue.pending_batch_count(), 0); assert!(!queue.has_pending_batch(&ctx)); assert!(!queue.has_pending_batch(&other_region_ctx)); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); } #[test] fn test_queue_take_pending_batch_removes_blocking_batch() { let ctx = new_test_req_ctx(2002); let base = Instant::now() + Duration::from_secs(60); let min_wait = Duration::from_secs(5); let max_wait = Duration::from_secs(20); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let key = queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap() .0; let inflight_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(12, 22)]), base); queue.mark_inflight(&inflight_batch); let mut ready = Vec::new(); set_queue_now(&mut queue, base + min_wait + Duration::from_secs(1)); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert!(ready.is_empty()); assert!(queue.region_queue.is_empty()); assert_eq!(queue.blocking_set.len(), 1); assert_eq!(blocked_first_seen(&queue, &key), Some(base)); let batch = queue.take_pending_batch(&key).unwrap(); assert_eq!(batch.key(), key); assert_eq!(batch.txn_count(), 2); assert_eq!(queue.pending_batch_count(), 0); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); assert!(queue.blocking_set.is_empty()); } #[test] fn test_queue_mark_inflight_tracks_dispatched_batches() { let ctx = new_test_req_ctx(2002); let other_region_ctx = new_test_req_ctx(2003); let first_seen = Instant::now(); let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), first_seen); let same_region_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(11, 0)]), first_seen); let other_region_batch = AsyncResolveLockBatch::new(&other_region_ctx, txn_status(&[(20, 40)]), first_seen); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, first_seen); // case: marking the first batch creates one inflight record and token. let token = queue.mark_inflight(&batch); assert_eq!(token.key, batch.key()); assert_eq!(token.id, 0); assert_eq!(queue.inflight_task_count(), 1); assert_eq!(queue.inflight_by_key.len(), 1); let inflight_batches = queue.inflight_by_key.get(&token.key).unwrap(); assert_eq!(inflight_batches.len(), 1); assert_eq!(inflight_batches[0].id, token.id); assert_eq!(inflight_batches[0].started_at, first_seen); assert_eq!(inflight_batches[0].txn_status.get(&ts(10)), Some(&ts(20))); assert_eq!(queue.pending_batch_count(), 0); assert_eq!(queue.pending_txn_count(), 0); // case: marking another batch for the same region keeps both inflight records. let same_region_token = queue.mark_inflight(&same_region_batch); assert_eq!(same_region_token.key, token.key); assert_eq!(same_region_token.id, 1); assert_eq!(queue.inflight_task_count(), 2); let inflight_batches = queue.inflight_by_key.get(&token.key).unwrap(); assert_eq!(inflight_batches.len(), 2); assert_eq!(inflight_batches[1].id, same_region_token.id); assert_eq!(inflight_batches[1].started_at, first_seen); assert_eq!(inflight_batches[1].txn_status.get(&ts(11)), Some(&ts(0))); // case: marking a batch for another region creates a separate inflight group. let other_region_token = queue.mark_inflight(&other_region_batch); assert_eq!(other_region_token.key, other_region_batch.key()); assert_eq!(other_region_token.id, 2); assert_eq!(queue.inflight_task_count(), 3); assert_eq!(queue.inflight_by_key.len(), 2); assert_eq!( queue .inflight_by_key .get(&other_region_token.key) .unwrap() .len(), 1 ); } #[test] fn test_queue_complete_inflight_removes_matching_token() { let ctx = new_test_req_ctx(2002); let other_region_ctx = new_test_req_ctx(2003); let first_seen = Instant::now(); let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), first_seen); let same_region_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(11, 0)]), first_seen); let other_region_batch = AsyncResolveLockBatch::new(&other_region_ctx, txn_status(&[(20, 40)]), first_seen); let mut queue = AsyncResolveLockQueue::new(); // Prepare two inflight batches for one region and one inflight batch for // another. set_queue_now(&mut queue, first_seen); let token = queue.mark_inflight(&batch); let same_region_token = queue.mark_inflight(&same_region_batch); let other_region_token = queue.mark_inflight(&other_region_batch); assert_eq!(queue.inflight_task_count(), 3); assert_eq!(queue.inflight_by_key.len(), 2); // case: completing a missing key leaves inflight state unchanged. let missing_token = AsyncResolveLockInflightToken { key: AsyncResolveLockBatchKey::from_context(&new_test_req_ctx(2004)), id: token.id, }; assert!(queue.complete_inflight(&missing_token).is_none()); assert_eq!(queue.inflight_task_count(), 3); assert_eq!(queue.inflight_by_key.len(), 2); // case: completing a missing id under an existing key leaves inflight state // unchanged. let wrong_id_token = AsyncResolveLockInflightToken { key: token.key.clone(), id: other_region_token.id.wrapping_add(1), }; assert!(queue.complete_inflight(&wrong_id_token).is_none()); assert_eq!(queue.inflight_task_count(), 3); assert_eq!(queue.inflight_by_key.get(&token.key).unwrap().len(), 2); // case: completing one of several batches for the same key keeps the group. let completed = queue.complete_inflight(&token).unwrap(); assert_eq!(completed.id, token.id); assert_eq!(completed.txn_status.get(&ts(10)), Some(&ts(20))); assert_eq!(queue.inflight_task_count(), 2); let inflight_batches = queue.inflight_by_key.get(&same_region_token.key).unwrap(); assert_eq!(inflight_batches.len(), 1); assert_eq!(inflight_batches[0].id, same_region_token.id); // case: completing the last batch for a key removes that key's group. let completed = queue.complete_inflight(&same_region_token).unwrap(); assert_eq!(completed.id, same_region_token.id); assert_eq!(completed.txn_status.get(&ts(11)), Some(&ts(0))); assert_eq!(queue.inflight_task_count(), 1); assert!(!queue.inflight_by_key.contains_key(&same_region_token.key)); assert!(queue.inflight_by_key.contains_key(&other_region_token.key)); // case: completing the last inflight batch leaves the inflight queue empty. let completed = queue.complete_inflight(&other_region_token).unwrap(); assert_eq!(completed.id, other_region_token.id); assert_eq!(completed.txn_status.get(&ts(20)), Some(&ts(40))); assert_eq!(queue.inflight_task_count(), 0); assert!(queue.inflight_by_key.is_empty()); } #[test] fn test_queue_remove_expired_inflight_removes_only_timed_out_batches() { let expired_ctx = new_test_req_ctx(2002); let fresh_ctx = new_test_req_ctx(2003); let base = Instant::now() + Duration::from_secs(60); let fresh_started_at = base + Duration::from_secs(10); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let expired_batch = AsyncResolveLockBatch::new(&expired_ctx, txn_status(&[(10, 20)]), base); let expired_token = queue.mark_inflight(&expired_batch); set_queue_now(&mut queue, fresh_started_at); let fresh_batch = AsyncResolveLockBatch::new(&fresh_ctx, txn_status(&[(30, 40)]), fresh_started_at); let fresh_token = queue.mark_inflight(&fresh_batch); assert_eq!(queue.inflight_task_count(), 2); let expired = queue.remove_expired_inflight( base + BatchAsyncResolveWorker::DEFAULT_INFLIGHT_TIMEOUT + Duration::from_secs(1), BatchAsyncResolveWorker::DEFAULT_INFLIGHT_TIMEOUT, ); assert_eq!(expired.len(), 1); assert_eq!(expired[0].token, expired_token); assert_eq!(expired[0].started_at, base); assert_eq!(expired[0].txn_status.get(&ts(10)), Some(&ts(20))); assert_eq!(queue.inflight_task_count(), 1); assert!(!queue.inflight_by_key.contains_key(&expired_token.key)); assert_eq!( queue.inflight_by_key.get(&fresh_token.key).unwrap().len(), 1 ); assert!(queue.complete_inflight(&expired_token).is_none()); let completed = queue.complete_inflight(&fresh_token).unwrap(); assert_eq!(completed.id, fresh_token.id); assert_eq!(completed.txn_status.get(&ts(30)), Some(&ts(40))); assert_eq!(queue.inflight_task_count(), 0); } #[test] fn test_queue_take_ready_batches_from_pending() { let ready_ctx = new_test_req_ctx(2002); let blocked_ctx = new_test_req_ctx(2003); let max_wait_ctx = new_test_req_ctx(2004); let not_ready_ctx = new_test_req_ctx(2005); let base = Instant::now() + Duration::from_secs(60); let max_ready_seen = base; let min_ready_seen = base + Duration::from_secs(3); let not_ready_seen = base + Duration::from_secs(6); let min_wait = Duration::from_secs(5); let max_wait = Duration::from_secs(10); let now = base + Duration::from_secs(10); let mut queue = AsyncResolveLockQueue::new(); // Prepare pending batches with deterministic first-seen timestamps. set_queue_now(&mut queue, max_ready_seen); let max_wait_key = queue .insert_for_batch(&max_wait_ctx, txn_status(&[(30, 60)])) .unwrap() .0; set_queue_now(&mut queue, min_ready_seen); let ready_key = queue .insert_for_batch(&ready_ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap() .0; let blocked_key = queue .insert_for_batch(&blocked_ctx, txn_status(&[(20, 40)])) .unwrap() .0; set_queue_now(&mut queue, not_ready_seen); let not_ready_key = queue .insert_for_batch(¬_ready_ctx, txn_status(&[(40, 80)])) .unwrap() .0; assert_eq!(queue.pending_batch_count(), 4); assert!(queue.has_pending_batch(&max_wait_ctx)); assert!(queue.has_pending_batch(&ready_ctx)); assert!(queue.has_pending_batch(&blocked_ctx)); assert!(queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 5); assert_eq!(queued_front_key(&queue), Some(&max_wait_key)); assert!(queue.blocking_set.is_empty()); // Prepare inflight batches for one min-waited batch and one max-waited // batch. let blocked_batch = AsyncResolveLockBatch::new(&blocked_ctx, txn_status(&[(21, 42)]), min_ready_seen); let max_wait_batch = AsyncResolveLockBatch::new(&max_wait_ctx, txn_status(&[(31, 62)]), max_ready_seen); set_queue_now(&mut queue, min_ready_seen + Duration::from_secs(1)); let blocked_token = queue.mark_inflight(&blocked_batch); let _max_wait_token = queue.mark_inflight(&max_wait_batch); assert_eq!(queue.inflight_task_count(), 2); let mut ready = Vec::new(); // case: before min_wait, no pending batch is ready. set_queue_now(&mut queue, base + Duration::from_secs(4)); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert!(ready.is_empty()); assert_eq!(queue.pending_batch_count(), 4); assert!(queue.has_pending_batch(&max_wait_ctx)); assert!(queue.has_pending_batch(&ready_ctx)); assert!(queue.has_pending_batch(&blocked_ctx)); assert!(queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 5); assert_eq!(queued_front_key(&queue), Some(&max_wait_key)); assert!(queue.blocking_set.is_empty()); // case: max-waited batches are taken even with inflight work, and // min-waited batches are taken only when their key has no inflight work. set_queue_now(&mut queue, now); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert_eq!(ready.len(), 2); assert_eq!(ready[0].key(), max_wait_key); assert_eq!(ready[0].txn_count(), 1); assert_eq!(ready[1].key(), ready_key); assert_eq!(ready[1].txn_count(), 2); assert_eq!(queue.pending_batch_count(), 2); assert!(!queue.has_pending_batch(&max_wait_ctx)); assert!(!queue.has_pending_batch(&ready_ctx)); assert!(queue.has_pending_batch(&blocked_ctx)); assert!(queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 2); assert_eq!(queue.region_queue.len(), 1); assert_eq!(queued_front_key(&queue), Some(¬_ready_key)); assert_eq!(queue.blocking_set.len(), 1); assert_eq!( blocked_first_seen(&queue, &blocked_key), Some(min_ready_seen) ); // case: after inflight completion, a min-waited retained batch becomes ready. ready.clear(); let _ = queue.complete_inflight(&blocked_token); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), blocked_key); assert_eq!(ready[0].txn_count(), 1); assert_eq!(queue.pending_batch_count(), 1); assert!(!queue.has_pending_batch(&blocked_ctx)); assert!(queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queued_front_key(&queue), Some(¬_ready_key)); assert!(queue.blocking_set.is_empty()); // case: a retained batch below min_wait remains pending. ready.clear(); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert!(ready.is_empty()); assert_eq!(queue.pending_batch_count(), 1); assert!(queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queued_front_key(&queue), Some(¬_ready_key)); assert!(queue.blocking_set.is_empty()); // case: the last retained batch is taken after it reaches min_wait. set_queue_now(&mut queue, not_ready_seen + min_wait); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), not_ready_key); assert_eq!(ready[0].txn_count(), 1); assert_eq!(queue.pending_batch_count(), 0); assert!(!queue.has_pending_batch(¬_ready_ctx)); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); assert!(queue.blocking_set.is_empty()); } #[test] fn test_queue_take_ready_batches_skips_blocked_batch() { let blocked_ctx = new_test_req_ctx(2002); let ready_ctx = new_test_req_ctx(2003); let not_ready_ctx = new_test_req_ctx(2004); let base = Instant::now() + Duration::from_secs(60); let min_wait = Duration::from_secs(5); let max_wait = Duration::from_secs(20); let min_ready_seen = base; let not_ready_seen = base + Duration::from_secs(4); let now = base + Duration::from_secs(6); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, min_ready_seen); let blocked_key = queue .insert_for_batch(&blocked_ctx, txn_status(&[(10, 20)])) .unwrap() .0; let ready_key = queue .insert_for_batch(&ready_ctx, txn_status(&[(11, 21)])) .unwrap() .0; set_queue_now(&mut queue, not_ready_seen); let not_ready_key = queue .insert_for_batch(¬_ready_ctx, txn_status(&[(12, 22)])) .unwrap() .0; let inflight_batch = AsyncResolveLockBatch::new(&blocked_ctx, txn_status(&[(13, 23)]), min_ready_seen); queue.mark_inflight(&inflight_batch); let mut ready = Vec::new(); set_queue_now(&mut queue, now); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), ready_key); assert_eq!(queue.pending_batch_count(), 2); assert_eq!(queue.pending_txn_count(), 2); assert_eq!(queue.region_queue.len(), 1); assert_eq!(queued_front_key(&queue), Some(¬_ready_key)); assert_eq!(queue.blocking_set.len(), 1); assert_eq!( blocked_first_seen(&queue, &blocked_key), Some(min_ready_seen) ); assert!(queue.has_pending_batch(&blocked_ctx)); assert!(!queue.has_pending_batch(&ready_ctx)); assert!(queue.has_pending_batch(¬_ready_ctx)); } #[test] fn test_queue_take_ready_batches_dispatches_blocking_batch_after_max_wait() { let blocked_ctx = new_test_req_ctx(2002); let base = Instant::now() + Duration::from_secs(60); let min_wait = Duration::from_secs(5); let max_wait = Duration::from_secs(10); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let blocked_key = queue .insert_for_batch(&blocked_ctx, txn_status(&[(10, 20)])) .unwrap() .0; let inflight_batch = AsyncResolveLockBatch::new(&blocked_ctx, txn_status(&[(11, 21)]), base); queue.mark_inflight(&inflight_batch); let mut ready = Vec::new(); set_queue_now(&mut queue, base + min_wait + Duration::from_secs(1)); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert!(ready.is_empty()); assert!(queue.region_queue.is_empty()); assert_eq!(queue.blocking_set.len(), 1); assert_eq!(blocked_first_seen(&queue, &blocked_key), Some(base)); set_queue_now(&mut queue, base + max_wait); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), blocked_key); assert_eq!(ready[0].txn_count(), 1); assert_eq!(queue.pending_batch_count(), 0); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); assert!(queue.blocking_set.is_empty()); // case: blocking batches also count ready batches already taken in the // current tick against the inflight limit. let full_ctx = new_test_req_ctx(2101); let ctx_1 = new_test_req_ctx(2102); let ctx_2 = new_test_req_ctx(2103); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let key_1 = queue .insert_for_batch(&ctx_1, txn_status(&[(10, 20)])) .unwrap() .0; let key_2 = queue .insert_for_batch(&ctx_2, txn_status(&[(20, 30)])) .unwrap() .0; let inflight_1 = AsyncResolveLockBatch::new(&ctx_1, txn_status(&[(11, 21)]), base); let inflight_2 = AsyncResolveLockBatch::new(&ctx_2, txn_status(&[(21, 31)]), base); let token_1 = queue.mark_inflight(&inflight_1); let token_2 = queue.mark_inflight(&inflight_2); let mut ready = Vec::new(); set_queue_now(&mut queue, base + min_wait); queue.take_ready_batches_from_pending(min_wait, max_wait, usize::MAX, &mut ready); assert!(ready.is_empty()); assert_eq!(queue.blocking_set.len(), 2); queue.complete_inflight(&token_1); queue.complete_inflight(&token_2); let full_batch = AsyncResolveLockBatch::new(&full_ctx, txn_status(&[(30, 40)]), base); queue.mark_inflight(&full_batch); queue.take_ready_batches_from_pending(min_wait, max_wait, 2, &mut ready); assert_eq!(ready.len(), 1); let ready_key = ready[0].key(); assert!(ready_key == key_1 || ready_key == key_2); assert_eq!(queue.pending_batch_count(), 1); assert_eq!(queue.pending_txn_count(), 1); let blocked_key = if ready_key == key_1 { key_2 } else { key_1 }; assert_eq!(blocked_first_seen(&queue, &blocked_key), Some(base)); assert!(queue.region_queue.is_empty()); } #[test] fn test_queue_take_ready_batches_waits_when_inflight_full_until_max_wait() { let inflight_ctx = new_test_req_ctx(2001); let pending_ctx = new_test_req_ctx(2002); let base = Instant::now() + Duration::from_secs(60); let min_wait = Duration::from_secs(5); let max_wait = Duration::from_secs(10); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let pending_key = queue .insert_for_batch(&pending_ctx, txn_status(&[(10, 20)])) .unwrap() .0; let inflight_batch = AsyncResolveLockBatch::new(&inflight_ctx, txn_status(&[(11, 21)]), base); queue.mark_inflight(&inflight_batch); let mut ready = Vec::new(); set_queue_now(&mut queue, base + min_wait + Duration::from_secs(1)); queue.take_ready_batches_from_pending(min_wait, max_wait, 1, &mut ready); assert!(ready.is_empty()); assert_eq!(queue.pending_batch_count(), 1); assert_eq!(queue.pending_txn_count(), 1); assert_eq!(queued_front_key(&queue), Some(&pending_key)); assert!(queue.blocking_set.is_empty()); set_queue_now(&mut queue, base + max_wait); queue.take_ready_batches_from_pending(min_wait, max_wait, 1, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), pending_key); assert_eq!(ready[0].txn_count(), 1); assert_eq!(queue.pending_batch_count(), 0); assert_eq!(queue.pending_txn_count(), 0); assert!(queue.region_queue.is_empty()); assert!(queue.blocking_set.is_empty()); // case: ready batches already taken in the current tick consume // inflight slots, so later min-waited batches stay pending. let inflight_ctx = new_test_req_ctx(2101); let ctx_1 = new_test_req_ctx(2102); let ctx_2 = new_test_req_ctx(2103); let ctx_3 = new_test_req_ctx(2104); let mut queue = AsyncResolveLockQueue::new(); set_queue_now(&mut queue, base); let key_1 = queue .insert_for_batch(&ctx_1, txn_status(&[(10, 20)])) .unwrap() .0; let key_2 = queue .insert_for_batch(&ctx_2, txn_status(&[(20, 30)])) .unwrap() .0; queue .insert_for_batch(&ctx_3, txn_status(&[(30, 40)])) .unwrap(); let inflight_batch = AsyncResolveLockBatch::new(&inflight_ctx, txn_status(&[(40, 50)]), base); queue.mark_inflight(&inflight_batch); let mut ready = Vec::new(); set_queue_now(&mut queue, base + min_wait); queue.take_ready_batches_from_pending(min_wait, max_wait, 2, &mut ready); assert_eq!(ready.len(), 1); assert_eq!(ready[0].key(), key_1); assert_eq!(queue.pending_batch_count(), 2); assert_eq!(queue.pending_txn_count(), 2); assert_eq!(queued_front_key(&queue), Some(&key_2)); assert!(queue.blocking_set.is_empty()); } fn new_test_worker( pending_batch_max_txn_count: usize, ) -> (BatchAsyncResolveWorker, mpsc::Sender) { new_test_worker_with_max_inflight(pending_batch_max_txn_count, 16) } fn new_test_worker_with_max_inflight( pending_batch_max_txn_count: usize, max_inflight_batch_count: usize, ) -> (BatchAsyncResolveWorker, mpsc::Sender) { BatchAsyncResolveWorker::new_with_params( 4, Arc::new(AtomicU64::new(duration_to_micros(Duration::from_secs(1)))), Duration::from_secs(20), pending_batch_max_txn_count, 16, max_inflight_batch_count, ) } fn take_scheduled_callback( callbacks: &mut Vec, index: usize, ) -> ResolveLockBatchCallback { callbacks.remove(index) } async fn send_worker_task( task_tx: &mpsc::Sender, ctx: Context, txn_status: HashMap, ) { task_tx .send(AsyncResolveLockTask::new(ctx, txn_status)) .await .unwrap(); } async fn run_worker_loop_once_with_ticker( worker: &mut BatchAsyncResolveWorker, ticker: &mut Interval, ) -> ( bool, Vec, Vec, ) { let scheduled = Arc::new(Mutex::new(Vec::new())); let callbacks = Arc::new(Mutex::new(Vec::new())); let scheduled_recorder = scheduled.clone(); let callbacks_recorder = callbacks.clone(); let scheduler = move |batch: AsyncResolveLockBatch, cb: ResolveLockBatchCallback| { scheduled_recorder.lock().unwrap().push(batch); callbacks_recorder.lock().unwrap().push(cb); }; let keep_running = worker.loop_once(&scheduler, ticker, &mut vec![]).await; let scheduled = std::mem::take(&mut *scheduled.lock().unwrap()); let callbacks = std::mem::take(&mut *callbacks.lock().unwrap()); (keep_running, scheduled, callbacks) } async fn run_worker_loop_once( worker: &mut BatchAsyncResolveWorker, ) -> ( bool, Vec, Vec, ) { let mut ticker = new_blocked_ticker().await; run_worker_loop_once_with_ticker(worker, &mut ticker).await } async fn run_worker_loop_once_on_tick( worker: &mut BatchAsyncResolveWorker, ) -> ( bool, Vec, Vec, ) { let mut ticker = tokio::time::interval(Duration::from_secs(3600)); // Keep the first immediately-ready tick so `loop_once` enters the tick branch. run_worker_loop_once_with_ticker(worker, &mut ticker).await } fn assert_one_scheduled_batch( queue: &AsyncResolveLockQueue, scheduled: &[AsyncResolveLockBatch], index: usize, expected_key: &AsyncResolveLockBatchKey, expected_txn_status: &[(u64, u64)], ) { let expected_txn_status = txn_status(expected_txn_status); let batch = &scheduled[index]; assert_eq!(batch.key(), *expected_key); assert_eq!(batch.txn_status, expected_txn_status); let inflight_batches = queue.inflight_by_key.get(expected_key).unwrap(); assert!( inflight_batches .iter() .any(|batch| batch.txn_status == expected_txn_status) ); } #[test] fn test_worker_new_uses_default_params() { let (worker, task_tx) = BatchAsyncResolveWorker::new(Arc::new(AtomicU64::new(500000))); // case: the task channel uses the default capacity. assert_eq!( task_tx.max_capacity(), BatchAsyncResolveWorker::TASK_DISPATCH_CHANNEL_CAPACITY ); assert_eq!( task_tx.capacity(), BatchAsyncResolveWorker::TASK_DISPATCH_CHANNEL_CAPACITY ); // case: worker timing and size thresholds use the default values. assert_eq!(worker.pending_batch_wait(), Duration::from_millis(500),); assert_eq!( worker.pending_batch_max_wait, BatchAsyncResolveWorker::DEFAULT_PENDING_BATCH_MAX_WAIT ); assert_eq!( worker.pending_batch_max_txn_count, BatchAsyncResolveWorker::DEFAULT_PENDING_BATCH_MAX_TXN_COUNT ); assert_eq!( worker.max_pending_batch_count, BatchAsyncResolveWorker::DEFAULT_MAX_PENDING_BATCH_COUNT ); assert_eq!( worker.max_inflight_batch_count, BatchAsyncResolveWorker::DEFAULT_MAX_INFLIGHT_BATCH_COUNT ); // case: constructor starts with empty queue and empty rate-limited log state. assert_eq!(worker.queue.pending_batch_count(), 0); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 0); assert!(worker.rate_limited_log_state_by_event_name.is_empty()); assert_eq!(worker.last_inflight_cleanup_at, None); } #[test] fn test_worker_new_with_params_uses_custom_params() { let task_channel_cap = 7; let pending_batch_wait = Duration::from_millis(30); let pending_batch_max_wait = Duration::from_millis(70); let pending_batch_max_txn_count = 11; let max_pending_batch_count = 13; let max_inflight_batch_count = 17; let (worker, task_tx) = BatchAsyncResolveWorker::new_with_params( task_channel_cap, Arc::new(AtomicU64::new(duration_to_micros(pending_batch_wait))), pending_batch_max_wait, pending_batch_max_txn_count, max_pending_batch_count, max_inflight_batch_count, ); // case: the task channel uses the custom capacity. assert_eq!(task_tx.max_capacity(), task_channel_cap); assert_eq!(task_tx.capacity(), task_channel_cap); // case: worker timing and size thresholds use the custom values. assert_eq!(worker.pending_batch_wait(), pending_batch_wait); assert_eq!(worker.pending_batch_max_wait, pending_batch_max_wait); assert_eq!( worker.pending_batch_max_txn_count, pending_batch_max_txn_count ); assert_eq!(worker.max_pending_batch_count, max_pending_batch_count); assert_eq!(worker.max_inflight_batch_count, max_inflight_batch_count); // case: constructor starts with empty queue and empty rate-limited log state. assert_eq!(worker.queue.pending_batch_count(), 0); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 0); assert!(worker.rate_limited_log_state_by_event_name.is_empty()); assert_eq!(worker.last_inflight_cleanup_at, None); } #[test] fn test_worker_log_rate_limited_tracks_skipped_count_by_event_name() { let (mut worker, _task_tx) = new_test_worker(8); let base = Instant::now() + Duration::from_secs(60); let mut event_a_skipped_counts = Vec::new(); let mut event_b_skipped_counts = Vec::new(); set_worker_now(&mut worker, base); worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); worker.log_rate_limited("event_b", |skipped_log_count| { event_b_skipped_counts.push(skipped_log_count); }); assert_eq!(event_a_skipped_counts, vec![0]); assert_eq!(event_b_skipped_counts, vec![0]); // case: skipped logs are counted per event and passed to the next // emitted log once the rate limit interval elapses. set_worker_now( &mut worker, base + BatchAsyncResolveWorker::DEFAULT_RATE_LIMITED_LOG_INTERVAL - Duration::from_millis(1), ); worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); assert_eq!(event_a_skipped_counts, vec![0]); set_worker_now( &mut worker, base + BatchAsyncResolveWorker::DEFAULT_RATE_LIMITED_LOG_INTERVAL, ); worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); assert_eq!(event_a_skipped_counts, vec![0, 2]); // case: after an emitted log reports the skipped count, the counter is // reset for the following interval. worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); set_worker_now( &mut worker, base + BatchAsyncResolveWorker::DEFAULT_RATE_LIMITED_LOG_INTERVAL * 2, ); worker.log_rate_limited("event_a", |skipped_log_count| { event_a_skipped_counts.push(skipped_log_count); }); assert_eq!(event_a_skipped_counts, vec![0, 2, 1]); } #[test] fn test_worker_maybe_cleanup_resolve_queue_throttles_cleanup_interval() { let (mut worker, _task_tx) = new_test_worker(8); let base = Instant::now() + Duration::from_secs(60); let ctx = new_test_req_ctx(2002); set_worker_now(&mut worker, base); let batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(10, 20)]), base); let token = worker.queue.mark_inflight(&batch); assert_eq!(worker.queue.inflight_task_count(), 1); let last_cleanup_at = base + BatchAsyncResolveWorker::DEFAULT_INFLIGHT_TIMEOUT + Duration::from_secs(1); worker.last_inflight_cleanup_at = Some(last_cleanup_at); // case: an expired inflight batch is not scanned again before the // cleanup interval elapses. let before_interval = last_cleanup_at + BatchAsyncResolveWorker::DEFAULT_INFLIGHT_CLEANUP_INTERVAL - Duration::from_millis(1); set_worker_now(&mut worker, before_interval); worker.maybe_cleanup_resolve_queue(); assert_eq!(worker.queue.inflight_task_count(), 1); assert!(worker.queue.inflight_by_key.contains_key(&token.key)); assert_eq!(worker.last_inflight_cleanup_at, Some(last_cleanup_at)); // case: once the interval elapses, expired inflight state is removed. let at_interval = last_cleanup_at + BatchAsyncResolveWorker::DEFAULT_INFLIGHT_CLEANUP_INTERVAL; set_worker_now(&mut worker, at_interval); worker.maybe_cleanup_resolve_queue(); assert_eq!(worker.queue.inflight_task_count(), 0); assert!(!worker.queue.inflight_by_key.contains_key(&token.key)); assert_eq!(worker.last_inflight_cleanup_at, Some(at_interval)); assert!(worker.queue.complete_inflight(&token).is_none()); } #[tokio::test] async fn test_worker_loop_once_returns_false_when_task_rx_is_closed() { let (mut worker, task_tx) = new_test_worker(8); drop(task_tx); // case: task receiver returns None after all task senders are dropped. let (keep_running, _scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(!keep_running); } #[tokio::test] async fn test_worker_loop_once_inserts_task_without_scheduling_small_batch() { let (mut worker, task_tx) = new_test_worker(2); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); let pending_key = AsyncResolveLockBatchKey::from_context(&ctx); send_worker_task(&task_tx, ctx.clone(), txn_status(&[(10, 20), (11, 0)])).await; // case: accepted task stays pending when merged txn count does not exceed the // threshold. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 1); assert!(worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 2); assert_eq!(worker.queue.inflight_task_count(), 0); assert_eq!( worker .queue .pending_by_key .get(&pending_key) .unwrap() .txn_status, txn_status(&[(10, 20), (11, 0)]) ); } #[tokio::test] async fn test_worker_loop_once_ignores_task_with_empty_txn_status() { let (mut worker, task_tx) = new_test_worker(2); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); send_worker_task(&task_tx, ctx.clone(), HashMap::default()).await; // case: empty txn_status is rejected by queue insertion and does not schedule // or leave pending state behind. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 0); } #[tokio::test] async fn test_worker_loop_once_schedules_large_batch_without_inflight() { let (mut worker, task_tx) = new_test_worker(2); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); let expected_key = AsyncResolveLockBatchKey::from_context(&ctx); send_worker_task( &task_tx, ctx.clone(), txn_status(&[(10, 20), (11, 0), (12, 30)]), ) .await; // case: accepted task is scheduled immediately when merged txn count exceeds // the threshold. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 1); assert_eq!(worker.queue.next_inflight_id, 1); assert_eq!(scheduled.len(), 1); assert_one_scheduled_batch( &worker.queue, &scheduled, 0, &expected_key, &[(10, 20), (11, 0), (12, 30)], ); } #[tokio::test] async fn test_worker_loop_once_schedules_large_batch_even_with_same_key_inflight() { let (mut worker, task_tx) = new_test_worker(2); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); let expected_key = AsyncResolveLockBatchKey::from_context(&ctx); let inflight_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(9, 19)]), now); let inflight_token = worker.queue.mark_inflight(&inflight_batch); send_worker_task( &task_tx, ctx.clone(), txn_status(&[(10, 20), (11, 0), (12, 30)]), ) .await; // case: exceeding the threshold still schedules immediately even when same-key // work is inflight. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 2); assert_eq!(worker.queue.next_inflight_id, 2); assert_eq!(scheduled.len(), 1); assert_one_scheduled_batch( &worker.queue, &scheduled, 0, &expected_key, &[(10, 20), (11, 0), (12, 30)], ); assert_eq!( worker .queue .inflight_by_key .get(&inflight_token.key) .unwrap() .len(), 2 ); } #[tokio::test] async fn test_worker_loop_once_drops_large_batch_when_inflight_full() { let (mut worker, task_tx) = new_test_worker_with_max_inflight(2, 1); let now = Instant::now(); set_worker_now(&mut worker, now); let inflight_ctx = new_test_req_ctx(2001); let ctx = new_test_req_ctx(2002); let inflight_batch = AsyncResolveLockBatch::new(&inflight_ctx, txn_status(&[(9, 19)]), now); worker.queue.mark_inflight(&inflight_batch); send_worker_task( &task_tx, ctx.clone(), txn_status(&[(10, 20), (11, 0), (12, 30)]), ) .await; // case: a batch exceeding the txn-count threshold is forced out of the // pending queue, and schedule rejects it while the inflight cap is full. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 1); } #[tokio::test] async fn test_worker_loop_once_schedules_when_pending_and_new_task_merge_over_threshold() { let (mut worker, task_tx) = new_test_worker(3); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); // Prepare one pending batch whose size alone does not exceed the threshold. let pending_key = worker .queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap() .0; assert_eq!(worker.queue.pending_batch_count(), 1); assert!(worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 2); send_worker_task(&task_tx, ctx.clone(), txn_status(&[(12, 30), (13, 0)])).await; // case: existing pending work and incoming task are both below the threshold, // but the merged batch is scheduled immediately once it exceeds the threshold. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 1); assert_eq!(worker.queue.next_inflight_id, 1); assert_eq!(scheduled.len(), 1); assert_one_scheduled_batch( &worker.queue, &scheduled, 0, &pending_key, &[(10, 20), (11, 0), (12, 30), (13, 0)], ); } #[tokio::test] async fn test_worker_loop_once_keeps_pending_when_pending_and_new_task_merge_not_over_threshold() { let (mut worker, task_tx) = new_test_worker(3); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); // Prepare one pending batch whose size alone does not exceed the threshold. let pending_key = worker .queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap() .0; assert_eq!(worker.queue.pending_batch_count(), 1); assert!(worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 2); send_worker_task(&task_tx, ctx.clone(), txn_status(&[(12, 30)])).await; // case: existing pending work and incoming task are both below the threshold, // and the merged batch stays pending when it still does not exceed the // threshold. let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 1); assert!(worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 3); assert_eq!( worker .queue .pending_by_key .get(&pending_key) .unwrap() .txn_status, txn_status(&[(10, 20), (11, 0), (12, 30)]) ); assert_eq!(worker.queue.inflight_task_count(), 0); } #[tokio::test] async fn test_worker_loop_once_tick_keeps_pending_when_no_batch_is_ready() { let (mut worker, _task_tx) = new_test_worker(8); let base = Instant::now(); let ctx = new_test_req_ctx(2002); let pending_key = AsyncResolveLockBatchKey::from_context(&ctx); set_queue_now(&mut worker.queue, base); worker .queue .insert_for_batch(&ctx, txn_status(&[(10, 20), (11, 0)])) .unwrap(); let now = base + Duration::from_millis(500); set_queue_now(&mut worker.queue, now); set_worker_now(&mut worker, now); // case: tick does not schedule anything before the pending batch reaches // min_wait. let (keep_running, scheduled, _callbacks) = run_worker_loop_once_on_tick(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 1); assert!(worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 2); assert_eq!(worker.queue.inflight_task_count(), 0); assert_eq!(worker.queue.next_inflight_id, 0); assert_eq!( worker .queue .pending_by_key .get(&pending_key) .unwrap() .txn_status, txn_status(&[(10, 20), (11, 0)]) ); } #[tokio::test] async fn test_worker_loop_once_tick_schedules_multiple_ready_batches() { let (mut worker, _task_tx) = new_test_worker(8); let base = Instant::now(); let ctx_1 = new_test_req_ctx(2002); let ctx_2 = new_test_req_ctx(2003); let key_1 = AsyncResolveLockBatchKey::from_context(&ctx_1); let key_2 = AsyncResolveLockBatchKey::from_context(&ctx_2); set_queue_now(&mut worker.queue, base); worker .queue .insert_for_batch(&ctx_1, txn_status(&[(10, 20)])) .unwrap(); worker .queue .insert_for_batch(&ctx_2, txn_status(&[(30, 40), (31, 0)])) .unwrap(); let now = base + Duration::from_secs(1); set_queue_now(&mut worker.queue, now); set_worker_now(&mut worker, now); // case: one tick schedules every ready pending batch in queue order. let (keep_running, scheduled, _callbacks) = run_worker_loop_once_on_tick(&mut worker).await; assert!(keep_running); assert_eq!(scheduled.len(), 2); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx_1)); assert!(!worker.queue.has_pending_batch(&ctx_2)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 2); assert_eq!(worker.queue.next_inflight_id, 2); assert_one_scheduled_batch(&worker.queue, &scheduled, 0, &key_1, &[(10, 20)]); assert_one_scheduled_batch(&worker.queue, &scheduled, 1, &key_2, &[(30, 40), (31, 0)]); } #[tokio::test] async fn test_worker_loop_once_tick_drops_max_waited_batch_when_inflight_full() { let (mut worker, _task_tx) = new_test_worker_with_max_inflight(8, 1); let base = Instant::now() + Duration::from_secs(60); let inflight_ctx = new_test_req_ctx(2001); let ctx = new_test_req_ctx(2002); set_queue_now(&mut worker.queue, base); worker .queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); let inflight_batch = AsyncResolveLockBatch::new(&inflight_ctx, txn_status(&[(11, 21)]), base); worker.queue.mark_inflight(&inflight_batch); let now = base + worker.pending_batch_max_wait; set_queue_now(&mut worker.queue, now); set_worker_now(&mut worker, now); // case: max-wait still forces dequeue, but scheduling refuses the batch // while the global inflight cap is full. let (keep_running, scheduled, _callbacks) = run_worker_loop_once_on_tick(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.pending_batch_count(), 0); assert!(!worker.queue.has_pending_batch(&ctx)); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 1); } #[tokio::test] async fn test_worker_loop_once_tick_expires_inflight_before_taking_ready_pending() { let (mut worker, _task_tx) = new_test_worker(8); let base = Instant::now() + Duration::from_secs(60); let ctx = new_test_req_ctx(2002); let key = AsyncResolveLockBatchKey::from_context(&ctx); set_worker_now(&mut worker, base); let expired_batch = AsyncResolveLockBatch::new(&ctx, txn_status(&[(9, 19)]), base); let expired_token = worker.queue.mark_inflight(&expired_batch); worker .queue .insert_for_batch(&ctx, txn_status(&[(10, 20)])) .unwrap(); assert_eq!(worker.queue.inflight_task_count(), 1); assert_eq!(worker.queue.pending_batch_count(), 1); let now = base + BatchAsyncResolveWorker::DEFAULT_INFLIGHT_TIMEOUT + Duration::from_secs(1); set_worker_now(&mut worker, now); // case: tick first removes the expired inflight batch, so the pending // same-key batch can be scheduled in the same loop. let (keep_running, scheduled, _callbacks) = run_worker_loop_once_on_tick(&mut worker).await; assert!(keep_running); assert_eq!(scheduled.len(), 1); assert_eq!(worker.queue.pending_batch_count(), 0); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 1); assert_one_scheduled_batch(&worker.queue, &scheduled, 0, &key, &[(10, 20)]); assert!( worker .queue .inflight_by_key .get(&key) .unwrap() .iter() .all(|batch| batch.id != expired_token.id) ); // case: the late completion of the force-removed batch is reported as // missing and leaves the newer inflight batch untouched. worker .inflight_complete_tx .send((expired_token.clone(), Ok(()))) .unwrap(); let (keep_running, scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert!(scheduled.is_empty()); assert_eq!(worker.queue.inflight_task_count(), 1); let remaining_token = { let batch = worker .queue .inflight_by_key .get(&key) .unwrap() .first() .unwrap(); AsyncResolveLockInflightToken { key: key.clone(), id: batch.id, } }; assert!(worker.queue.complete_inflight(&remaining_token).is_some()); assert_eq!(worker.queue.inflight_task_count(), 0); } async fn run_complete_inflight_after_callback_case(callback_result: storage::Result<()>) { let (mut worker, task_tx) = new_test_worker(2); let now = Instant::now(); set_worker_now(&mut worker, now); let ctx = new_test_req_ctx(2002); let key = AsyncResolveLockBatchKey::from_context(&ctx); send_worker_task( &task_tx, ctx.clone(), txn_status(&[(10, 20), (11, 0), (12, 30)]), ) .await; // case: a large task is scheduled immediately and becomes inflight. let (keep_running, scheduled, mut callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert_eq!(scheduled.len(), 1); assert_eq!(callbacks.len(), 1); assert_eq!(worker.queue.inflight_task_count(), 1); assert_one_scheduled_batch( &worker.queue, &scheduled, 0, &key, &[(10, 20), (11, 0), (12, 30)], ); // case: a successful callback enqueues inflight completion, and the next // loop consumes it and removes the inflight batch. take_scheduled_callback(&mut callbacks, 0)(callback_result); let (keep_running, _scheduled, _callbacks) = run_worker_loop_once(&mut worker).await; assert!(keep_running); assert_eq!(worker.queue.pending_batch_count(), 0); assert_eq!(worker.queue.pending_txn_count(), 0); assert_eq!(worker.queue.inflight_task_count(), 0); assert!(!worker.queue.inflight_by_key.contains_key(&key)); assert_eq!(worker.queue.next_inflight_id, 1); drop(task_tx); } #[tokio::test] async fn test_worker_loop_once_completes_inflight_after_callback() { // case: a successful callback removes the inflight batch on the next loop. run_complete_inflight_after_callback_case(Ok(())).await; // case: an error callback still removes the inflight batch on the next // loop. run_complete_inflight_after_callback_case(Err(storage::errors::ErrorInner::Other( box_err!("test completion error"), ) .into())) .await; } }