mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
perf: queue request candidate persistence
This commit is contained in:
elky
@@ -0,0 +1,751 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use aether_runtime::{MetricKind, MetricSample};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::{interval, MissedTickBehavior};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
const MODE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_WRITE_MODE";
|
||||
const QUEUE_CAPACITY_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_CAPACITY";
|
||||
const BATCH_SIZE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_BATCH_SIZE";
|
||||
const FLUSH_INTERVAL_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FLUSH_INTERVAL_MS";
|
||||
const WORKERS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_WORKERS";
|
||||
const QUEUE_FULL_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FULL";
|
||||
|
||||
const DEFAULT_QUEUE_CAPACITY: usize = 65_536;
|
||||
const DEFAULT_BATCH_SIZE: usize = 512;
|
||||
const DEFAULT_FLUSH_INTERVAL_MS: u64 = 50;
|
||||
const DEFAULT_WORKERS: usize = 2;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RequestCandidateWriteMode {
|
||||
Sync,
|
||||
Async,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RequestCandidateQueueFullPolicy {
|
||||
Drop,
|
||||
Sync,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct RequestCandidateQueueConfig {
|
||||
pub(crate) mode: RequestCandidateWriteMode,
|
||||
pub(crate) capacity: usize,
|
||||
pub(crate) batch_size: usize,
|
||||
pub(crate) flush_interval: Duration,
|
||||
pub(crate) workers: usize,
|
||||
pub(crate) full_policy: RequestCandidateQueueFullPolicy,
|
||||
}
|
||||
|
||||
impl Default for RequestCandidateQueueConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
mode: RequestCandidateWriteMode::Sync,
|
||||
capacity: DEFAULT_QUEUE_CAPACITY,
|
||||
batch_size: DEFAULT_BATCH_SIZE,
|
||||
flush_interval: Duration::from_millis(DEFAULT_FLUSH_INTERVAL_MS),
|
||||
workers: DEFAULT_WORKERS,
|
||||
full_policy: RequestCandidateQueueFullPolicy::Drop,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RequestCandidateQueueConfig {
|
||||
pub(crate) fn from_env() -> Self {
|
||||
let mut config = Self::default();
|
||||
config.mode = match env_string(MODE_ENV).as_deref() {
|
||||
Some("async") | Some("queued") | Some("queue") => RequestCandidateWriteMode::Async,
|
||||
_ => RequestCandidateWriteMode::Sync,
|
||||
};
|
||||
config.capacity = env_usize(QUEUE_CAPACITY_ENV, DEFAULT_QUEUE_CAPACITY).max(1);
|
||||
config.batch_size = env_usize(BATCH_SIZE_ENV, DEFAULT_BATCH_SIZE).max(1);
|
||||
config.flush_interval =
|
||||
Duration::from_millis(env_u64(FLUSH_INTERVAL_MS_ENV, DEFAULT_FLUSH_INTERVAL_MS).max(1));
|
||||
config.workers = env_usize(WORKERS_ENV, DEFAULT_WORKERS).clamp(1, 32);
|
||||
config.full_policy = match env_string(QUEUE_FULL_ENV).as_deref() {
|
||||
Some("sync") | Some("fallback_sync") | Some("fallback-sync") => {
|
||||
RequestCandidateQueueFullPolicy::Sync
|
||||
}
|
||||
_ => RequestCandidateQueueFullPolicy::Drop,
|
||||
};
|
||||
config
|
||||
}
|
||||
|
||||
pub(crate) fn async_enabled(&self) -> bool {
|
||||
matches!(self.mode, RequestCandidateWriteMode::Async)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RequestCandidateQueueMetrics {
|
||||
queued_current: AtomicUsize,
|
||||
pending_current: AtomicUsize,
|
||||
enqueued_total: AtomicU64,
|
||||
dropped_total: AtomicU64,
|
||||
flushed_total: AtomicU64,
|
||||
flush_failed_total: AtomicU64,
|
||||
flush_batches_total: AtomicU64,
|
||||
flush_sql_ops_total: AtomicU64,
|
||||
compacted_total: AtomicU64,
|
||||
sync_fallback_total: AtomicU64,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RequestCandidateQueueRuntime {
|
||||
senders: Vec<mpsc::Sender<UpsertRequestCandidateRecord>>,
|
||||
repository: Arc<dyn RequestCandidateWriteRepository>,
|
||||
config: RequestCandidateQueueConfig,
|
||||
metrics: Arc<RequestCandidateQueueMetrics>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RequestCandidateQueueRuntime {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("RequestCandidateQueueRuntime")
|
||||
.field("config", &self.config)
|
||||
.field(
|
||||
"queued_current",
|
||||
&self.metrics.queued_current.load(Ordering::Acquire),
|
||||
)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl RequestCandidateQueueRuntime {
|
||||
pub(crate) fn spawn(
|
||||
repository: Arc<dyn RequestCandidateWriteRepository>,
|
||||
mut config: RequestCandidateQueueConfig,
|
||||
) -> Arc<Self> {
|
||||
config.workers = config.workers.min(config.capacity).max(1);
|
||||
let mut senders = Vec::with_capacity(config.workers);
|
||||
let mut receivers = Vec::with_capacity(config.workers);
|
||||
for worker_index in 0..config.workers {
|
||||
let capacity = worker_queue_capacity(config.capacity, config.workers, worker_index);
|
||||
let (sender, receiver) = mpsc::channel(capacity);
|
||||
senders.push(sender);
|
||||
receivers.push(receiver);
|
||||
}
|
||||
let runtime = Arc::new(Self {
|
||||
senders,
|
||||
repository,
|
||||
config,
|
||||
metrics: Arc::new(RequestCandidateQueueMetrics::default()),
|
||||
});
|
||||
runtime.spawn_workers(receivers);
|
||||
runtime
|
||||
}
|
||||
|
||||
pub(crate) async fn enqueue_or_fallback(
|
||||
&self,
|
||||
record: UpsertRequestCandidateRecord,
|
||||
) -> Result<(), aether_data::DataLayerError> {
|
||||
let worker_index = self.worker_index_for(&record);
|
||||
let Some(sender) = self.senders.get(worker_index) else {
|
||||
self.metrics
|
||||
.sync_fallback_total
|
||||
.fetch_add(1, Ordering::AcqRel);
|
||||
return self.repository.upsert(record).await.map(|_| ());
|
||||
};
|
||||
self.metrics.queued_current.fetch_add(1, Ordering::AcqRel);
|
||||
self.metrics.pending_current.fetch_add(1, Ordering::AcqRel);
|
||||
match sender.try_send(record) {
|
||||
Ok(()) => {
|
||||
self.metrics.enqueued_total.fetch_add(1, Ordering::AcqRel);
|
||||
Ok(())
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Full(record)) => {
|
||||
decrement_atomic_usize(&self.metrics.queued_current);
|
||||
decrement_atomic_usize(&self.metrics.pending_current);
|
||||
warn!(
|
||||
event_name = "request_candidate_queue_full",
|
||||
log_type = "event",
|
||||
full_policy = ?self.config.full_policy,
|
||||
worker_index,
|
||||
queued = self.metrics.queued_current.load(Ordering::Acquire),
|
||||
capacity = self.config.capacity,
|
||||
"gateway request candidate async queue is full"
|
||||
);
|
||||
match self.config.full_policy {
|
||||
RequestCandidateQueueFullPolicy::Drop => {
|
||||
self.metrics.dropped_total.fetch_add(1, Ordering::AcqRel);
|
||||
Ok(())
|
||||
}
|
||||
RequestCandidateQueueFullPolicy::Sync => {
|
||||
self.metrics
|
||||
.sync_fallback_total
|
||||
.fetch_add(1, Ordering::AcqRel);
|
||||
self.repository.upsert(record).await.map(|_| ())
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Closed(record)) => {
|
||||
decrement_atomic_usize(&self.metrics.queued_current);
|
||||
decrement_atomic_usize(&self.metrics.pending_current);
|
||||
self.metrics
|
||||
.sync_fallback_total
|
||||
.fetch_add(1, Ordering::AcqRel);
|
||||
self.repository.upsert(record).await.map(|_| ())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn metric_samples(&self) -> Vec<MetricSample> {
|
||||
vec![
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_depth",
|
||||
"Current number of request candidate records waiting in the async persistence queue.",
|
||||
MetricKind::Gauge,
|
||||
self.metrics.queued_current.load(Ordering::Acquire) as u64,
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_pending_depth",
|
||||
"Current number of request candidate records accepted into the async persistence queue but not yet flushed.",
|
||||
MetricKind::Gauge,
|
||||
self.metrics.pending_current.load(Ordering::Acquire) as u64,
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_capacity",
|
||||
"Configured request candidate async persistence queue capacity.",
|
||||
MetricKind::Gauge,
|
||||
self.config.capacity as u64,
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_enqueued_total",
|
||||
"Total request candidate records accepted into the async persistence queue.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.enqueued_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_dropped_total",
|
||||
"Total request candidate records dropped because the async persistence queue was full.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.dropped_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_flushed_total",
|
||||
"Total request candidate records flushed by async persistence workers.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.flushed_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_flush_failed_total",
|
||||
"Total request candidate records that failed during async persistence flush.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.flush_failed_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_flush_batches_total",
|
||||
"Total async request candidate persistence flush batches.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.flush_batches_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_flush_sql_ops_total",
|
||||
"Total repository upsert operations issued by async request candidate persistence workers after compaction.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.flush_sql_ops_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_compacted_total",
|
||||
"Total request candidate records compacted before async persistence because a later queued record covered the same slot and status.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.compacted_total.load(Ordering::Acquire),
|
||||
),
|
||||
MetricSample::new(
|
||||
"request_candidate_queue_sync_fallback_total",
|
||||
"Total request candidate records synchronously persisted after async queue fallback.",
|
||||
MetricKind::Counter,
|
||||
self.metrics.sync_fallback_total.load(Ordering::Acquire),
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
fn spawn_workers(
|
||||
self: &Arc<Self>,
|
||||
receivers: Vec<mpsc::Receiver<UpsertRequestCandidateRecord>>,
|
||||
) {
|
||||
for (worker_index, receiver) in receivers.into_iter().enumerate() {
|
||||
let repository = Arc::clone(&self.repository);
|
||||
let config = self.config.clone();
|
||||
let metrics = Arc::clone(&self.metrics);
|
||||
tokio::spawn(async move {
|
||||
run_worker(repository, config, metrics, worker_index, receiver).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn worker_index_for(&self, record: &UpsertRequestCandidateRecord) -> usize {
|
||||
let worker_count = self.senders.len();
|
||||
if worker_count <= 1 {
|
||||
return 0;
|
||||
}
|
||||
(request_candidate_slot_hash(record) % worker_count as u64) as usize
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_worker(
|
||||
repository: Arc<dyn RequestCandidateWriteRepository>,
|
||||
config: RequestCandidateQueueConfig,
|
||||
metrics: Arc<RequestCandidateQueueMetrics>,
|
||||
worker_index: usize,
|
||||
mut receiver: mpsc::Receiver<UpsertRequestCandidateRecord>,
|
||||
) {
|
||||
let mut ticker = interval(config.flush_interval);
|
||||
ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
|
||||
let mut batch = Vec::with_capacity(config.batch_size);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = ticker.tick() => {
|
||||
if !batch.is_empty() {
|
||||
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
|
||||
}
|
||||
}
|
||||
received = receiver.recv() => {
|
||||
match received {
|
||||
Some(record) => {
|
||||
decrement_atomic_usize(&metrics.queued_current);
|
||||
batch.push(record);
|
||||
if batch.len() >= config.batch_size {
|
||||
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
if !batch.is_empty() {
|
||||
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn flush_batch(
|
||||
repository: &Arc<dyn RequestCandidateWriteRepository>,
|
||||
metrics: &RequestCandidateQueueMetrics,
|
||||
worker_index: usize,
|
||||
batch: &mut Vec<UpsertRequestCandidateRecord>,
|
||||
) {
|
||||
let records = std::mem::take(batch);
|
||||
if records.is_empty() {
|
||||
return;
|
||||
}
|
||||
let source_count = records.len();
|
||||
let records = compact_same_status_records(records);
|
||||
let compacted = source_count.saturating_sub(records.len());
|
||||
if compacted > 0 {
|
||||
metrics
|
||||
.compacted_total
|
||||
.fetch_add(compacted as u64, Ordering::AcqRel);
|
||||
}
|
||||
metrics.flush_batches_total.fetch_add(1, Ordering::AcqRel);
|
||||
let mut failed = 0_u64;
|
||||
for record in records {
|
||||
metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel);
|
||||
if let Err(err) = repository.upsert(record.record).await {
|
||||
failed = failed.saturating_add(record.source_count as u64);
|
||||
decrement_atomic_usize_by(&metrics.pending_current, record.source_count);
|
||||
warn!(
|
||||
event_name = "request_candidate_async_flush_failed",
|
||||
log_type = "event",
|
||||
worker_index,
|
||||
error = ?err,
|
||||
"gateway failed to asynchronously persist request candidate"
|
||||
);
|
||||
} else {
|
||||
metrics
|
||||
.flushed_total
|
||||
.fetch_add(record.source_count as u64, Ordering::AcqRel);
|
||||
decrement_atomic_usize_by(&metrics.pending_current, record.source_count);
|
||||
}
|
||||
}
|
||||
if failed > 0 {
|
||||
metrics
|
||||
.flush_failed_total
|
||||
.fetch_add(failed, Ordering::AcqRel);
|
||||
}
|
||||
debug!(
|
||||
event_name = "request_candidate_async_flush_completed",
|
||||
log_type = "event",
|
||||
worker_index,
|
||||
failed,
|
||||
"gateway completed request candidate async flush batch"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CompactedRequestCandidateRecord {
|
||||
record: UpsertRequestCandidateRecord,
|
||||
source_count: usize,
|
||||
}
|
||||
|
||||
fn compact_same_status_records(
|
||||
records: Vec<UpsertRequestCandidateRecord>,
|
||||
) -> Vec<CompactedRequestCandidateRecord> {
|
||||
let mut latest_slot_status = HashMap::<(String, u32, u32), (u8, usize)>::new();
|
||||
let mut compacted = Vec::<CompactedRequestCandidateRecord>::with_capacity(records.len());
|
||||
for record in records {
|
||||
let slot = (
|
||||
record.request_id.clone(),
|
||||
record.candidate_index,
|
||||
record.retry_index,
|
||||
);
|
||||
let status = request_candidate_status_discriminant(record.status);
|
||||
match latest_slot_status.get(&slot).copied() {
|
||||
Some((latest_status, index)) if latest_status == status => {
|
||||
merge_request_candidate_record(&mut compacted[index].record, record);
|
||||
compacted[index].source_count = compacted[index].source_count.saturating_add(1);
|
||||
}
|
||||
_ => {
|
||||
latest_slot_status.insert(slot, (status, compacted.len()));
|
||||
compacted.push(CompactedRequestCandidateRecord {
|
||||
record,
|
||||
source_count: 1,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
compacted
|
||||
}
|
||||
|
||||
fn request_candidate_slot_hash(record: &UpsertRequestCandidateRecord) -> u64 {
|
||||
const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
|
||||
const FNV_PRIME: u64 = 1_099_511_628_211;
|
||||
|
||||
let mut hash = FNV_OFFSET;
|
||||
for byte in record.request_id.as_bytes() {
|
||||
hash ^= u64::from(*byte);
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
for byte in record.candidate_index.to_le_bytes() {
|
||||
hash ^= u64::from(byte);
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
for byte in record.retry_index.to_le_bytes() {
|
||||
hash ^= u64::from(byte);
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
hash
|
||||
}
|
||||
|
||||
fn worker_queue_capacity(total_capacity: usize, workers: usize, worker_index: usize) -> usize {
|
||||
let workers = workers.max(1);
|
||||
let base = total_capacity / workers;
|
||||
let remainder = total_capacity % workers;
|
||||
(base + usize::from(worker_index < remainder)).max(1)
|
||||
}
|
||||
|
||||
fn merge_request_candidate_record(
|
||||
target: &mut UpsertRequestCandidateRecord,
|
||||
incoming: UpsertRequestCandidateRecord,
|
||||
) {
|
||||
if !incoming.id.trim().is_empty() {
|
||||
target.id = incoming.id;
|
||||
}
|
||||
target.request_id = incoming.request_id;
|
||||
target.candidate_index = incoming.candidate_index;
|
||||
target.retry_index = incoming.retry_index;
|
||||
target.status = incoming.status;
|
||||
|
||||
take_if_some(&mut target.user_id, incoming.user_id);
|
||||
take_if_some(&mut target.api_key_id, incoming.api_key_id);
|
||||
take_if_some(&mut target.username, incoming.username);
|
||||
take_if_some(&mut target.api_key_name, incoming.api_key_name);
|
||||
take_if_some(&mut target.provider_id, incoming.provider_id);
|
||||
take_if_some(&mut target.endpoint_id, incoming.endpoint_id);
|
||||
take_if_some(&mut target.key_id, incoming.key_id);
|
||||
take_if_some(&mut target.skip_reason, incoming.skip_reason);
|
||||
take_if_some(&mut target.is_cached, incoming.is_cached);
|
||||
take_if_some(&mut target.status_code, incoming.status_code);
|
||||
take_if_some(&mut target.error_type, incoming.error_type);
|
||||
take_if_some(&mut target.error_message, incoming.error_message);
|
||||
take_if_some(&mut target.latency_ms, incoming.latency_ms);
|
||||
take_if_some(
|
||||
&mut target.concurrent_requests,
|
||||
incoming.concurrent_requests,
|
||||
);
|
||||
merge_json_value(&mut target.extra_data, incoming.extra_data);
|
||||
merge_json_value(
|
||||
&mut target.required_capabilities,
|
||||
incoming.required_capabilities,
|
||||
);
|
||||
|
||||
if target.created_at_unix_ms.is_none() {
|
||||
target.created_at_unix_ms = incoming.created_at_unix_ms;
|
||||
}
|
||||
take_if_some(&mut target.started_at_unix_ms, incoming.started_at_unix_ms);
|
||||
take_if_some(
|
||||
&mut target.finished_at_unix_ms,
|
||||
incoming.finished_at_unix_ms,
|
||||
);
|
||||
}
|
||||
|
||||
fn request_candidate_status_discriminant(status: RequestCandidateStatus) -> u8 {
|
||||
match status {
|
||||
RequestCandidateStatus::Available => 0,
|
||||
RequestCandidateStatus::Unused => 1,
|
||||
RequestCandidateStatus::Pending => 2,
|
||||
RequestCandidateStatus::Streaming => 3,
|
||||
RequestCandidateStatus::Success => 4,
|
||||
RequestCandidateStatus::Failed => 5,
|
||||
RequestCandidateStatus::Cancelled => 6,
|
||||
RequestCandidateStatus::Skipped => 7,
|
||||
}
|
||||
}
|
||||
|
||||
fn take_if_some<T>(target: &mut Option<T>, incoming: Option<T>) {
|
||||
if incoming.is_some() {
|
||||
*target = incoming;
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_json_value(target: &mut Option<serde_json::Value>, incoming: Option<serde_json::Value>) {
|
||||
match (target.as_mut(), incoming) {
|
||||
(Some(serde_json::Value::Object(target)), Some(serde_json::Value::Object(incoming))) => {
|
||||
target.extend(incoming);
|
||||
}
|
||||
(_, Some(incoming)) => {
|
||||
*target = Some(incoming);
|
||||
}
|
||||
(_, None) => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn env_string(key: &str) -> Option<String> {
|
||||
std::env::var(key)
|
||||
.ok()
|
||||
.map(|value| value.trim().to_ascii_lowercase())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn env_usize(key: &str, default: usize) -> usize {
|
||||
std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn env_u64(key: &str, default: u64) -> u64 {
|
||||
std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<u64>().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn decrement_atomic_usize(value: &AtomicUsize) {
|
||||
decrement_atomic_usize_by(value, 1);
|
||||
}
|
||||
|
||||
fn decrement_atomic_usize_by(value: &AtomicUsize, amount: usize) {
|
||||
let _ = value.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
|
||||
Some(current.saturating_sub(amount))
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
compact_same_status_records, RequestCandidateQueueConfig, RequestCandidateQueueRuntime,
|
||||
};
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::DataLayerError;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Default)]
|
||||
struct DelayedPendingRequestCandidateRepository {
|
||||
inner: InMemoryRequestCandidateRepository,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl RequestCandidateWriteRepository for DelayedPendingRequestCandidateRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
if candidate.status == RequestCandidateStatus::Pending {
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
}
|
||||
self.inner.upsert(candidate).await
|
||||
}
|
||||
|
||||
async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
self.inner
|
||||
.delete_created_before(created_before_unix_secs, limit)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn record(
|
||||
request_id: &str,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
status: RequestCandidateStatus,
|
||||
) -> UpsertRequestCandidateRecord {
|
||||
UpsertRequestCandidateRecord {
|
||||
id: format!("{request_id}-{candidate_index}-{retry_index}-{status:?}"),
|
||||
request_id: request_id.to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
status,
|
||||
skip_reason: None,
|
||||
is_cached: None,
|
||||
status_code: None,
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
concurrent_requests: None,
|
||||
extra_data: None,
|
||||
required_capabilities: None,
|
||||
created_at_unix_ms: None,
|
||||
started_at_unix_ms: None,
|
||||
finished_at_unix_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_merges_same_slot_and_status_without_dropping_state_transitions() {
|
||||
let mut first_success = record("req", 0, 0, RequestCandidateStatus::Success);
|
||||
first_success.provider_id = Some("provider-a".to_string());
|
||||
first_success.extra_data = Some(serde_json::json!({"first": true}));
|
||||
let mut second_success = record("req", 0, 0, RequestCandidateStatus::Success);
|
||||
second_success.latency_ms = Some(123);
|
||||
second_success.extra_data = Some(serde_json::json!({"second": true}));
|
||||
|
||||
let compacted = compact_same_status_records(vec![
|
||||
record("req", 0, 0, RequestCandidateStatus::Pending),
|
||||
first_success,
|
||||
record("req", 0, 1, RequestCandidateStatus::Failed),
|
||||
second_success,
|
||||
]);
|
||||
|
||||
assert_eq!(compacted.len(), 3);
|
||||
assert_eq!(compacted[0].record.status, RequestCandidateStatus::Pending);
|
||||
assert_eq!(compacted[0].source_count, 1);
|
||||
assert_eq!(compacted[1].record.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(compacted[1].source_count, 2);
|
||||
assert_eq!(
|
||||
compacted[1].record.provider_id.as_deref(),
|
||||
Some("provider-a")
|
||||
);
|
||||
assert_eq!(compacted[1].record.latency_ms, Some(123));
|
||||
assert_eq!(
|
||||
compacted[1].record.extra_data,
|
||||
Some(serde_json::json!({"first": true, "second": true}))
|
||||
);
|
||||
assert_eq!(compacted[2].record.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(compacted[2].source_count, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_preserves_same_slot_status_order_across_transitions() {
|
||||
let compacted = compact_same_status_records(vec![
|
||||
record("req", 0, 0, RequestCandidateStatus::Success),
|
||||
record("req", 0, 0, RequestCandidateStatus::Failed),
|
||||
record("req", 0, 0, RequestCandidateStatus::Success),
|
||||
]);
|
||||
|
||||
assert_eq!(compacted.len(), 3);
|
||||
assert_eq!(compacted[0].record.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(compacted[1].record.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(compacted[2].record.status, RequestCandidateStatus::Success);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_queue_flushes_enqueued_records() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let runtime = RequestCandidateQueueRuntime::spawn(
|
||||
repository.clone(),
|
||||
RequestCandidateQueueConfig {
|
||||
mode: super::RequestCandidateWriteMode::Async,
|
||||
capacity: 16,
|
||||
batch_size: 2,
|
||||
flush_interval: Duration::from_millis(10),
|
||||
workers: 1,
|
||||
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
|
||||
},
|
||||
);
|
||||
|
||||
runtime
|
||||
.enqueue_or_fallback(record("req", 0, 0, RequestCandidateStatus::Success))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for _ in 0..50 {
|
||||
let rows = repository.list_by_request_id("req").await.unwrap();
|
||||
if rows.len() == 1 {
|
||||
assert_eq!(rows[0].status, RequestCandidateStatus::Success);
|
||||
return;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
panic!("async request candidate queue did not flush record in time");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_queue_preserves_same_slot_order_with_multiple_workers() {
|
||||
let repository = Arc::new(DelayedPendingRequestCandidateRepository::default());
|
||||
let runtime = RequestCandidateQueueRuntime::spawn(
|
||||
repository.clone(),
|
||||
RequestCandidateQueueConfig {
|
||||
mode: super::RequestCandidateWriteMode::Async,
|
||||
capacity: 16,
|
||||
batch_size: 1,
|
||||
flush_interval: Duration::from_millis(100),
|
||||
workers: 2,
|
||||
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
|
||||
},
|
||||
);
|
||||
|
||||
runtime
|
||||
.enqueue_or_fallback(record("req-order", 0, 0, RequestCandidateStatus::Pending))
|
||||
.await
|
||||
.unwrap();
|
||||
runtime
|
||||
.enqueue_or_fallback(record("req-order", 0, 0, RequestCandidateStatus::Success))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for _ in 0..50 {
|
||||
if runtime.metrics.pending_current.load(Ordering::Acquire) == 0 {
|
||||
let rows = repository
|
||||
.inner
|
||||
.list_by_request_id("req-order")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].status, RequestCandidateStatus::Success);
|
||||
return;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
panic!("async request candidate queue did not finish ordered same-slot writes in time");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user