From 75eb26f859c229f1b269506951744650b79e59f1 Mon Sep 17 00:00:00 2001 From: james-mcnulty Date: Thu, 16 Apr 2026 15:41:04 -0700 Subject: [PATCH 1/2] Split Apart God Trait (WIP, Still Broken) --- benches/store_bench.rs | 6 +- src/fetch/mod.rs | 16 +- src/fetch/tests.rs | 133 +----- src/grpc/server.rs | 17 +- src/grpc/server_tests.rs | 18 +- src/kafka/consumer.rs | 12 +- src/kafka/deserialize_activation.rs | 22 +- src/kafka/inflight_activation_batcher.rs | 26 +- src/kafka/inflight_activation_writer.rs | 70 +-- src/main.rs | 15 +- src/push/mod.rs | 22 +- src/store/activation.rs | 50 +- src/store/adapters/postgres.rs | 507 +++++++++++---------- src/store/adapters/sqlite.rs | 551 ++++++++++++----------- src/store/tests.rs | 303 +++++++------ src/store/traits.rs | 184 ++++---- src/test_utils.rs | 46 +- src/upkeep.rs | 62 +-- 18 files changed, 1010 insertions(+), 1050 deletions(-) diff --git a/benches/store_bench.rs b/benches/store_bench.rs index 6f110e23..a868cf9f 100644 --- a/benches/store_bench.rs +++ b/benches/store_bench.rs @@ -4,7 +4,7 @@ use chrono::Utc; use criterion::{Criterion, criterion_group, criterion_main}; use rand::Rng; use taskbroker::{ - store::activation::InflightActivationStatus, + store::activation::ActivationStatus, store::adapters::sqlite::{InflightActivationStoreConfig, SqliteActivationStore}, store::traits::InflightActivationStore, test_utils::{ @@ -122,7 +122,7 @@ async fn set_status(num_activations: u32, num_workers: u32) { for task_id in 0..num_activations { if task_id % num_workers == worker_idx { store - .set_status(&format!("id_{task_id}"), InflightActivationStatus::Complete) + .set_status(&format!("id_{task_id}"), ActivationStatus::Complete) .await .unwrap(); } @@ -134,7 +134,7 @@ async fn set_status(num_activations: u32, num_workers: u32) { assert_eq!( store - .count_by_status(InflightActivationStatus::Complete) + .count_by_status(ActivationStatus::Complete) .await .unwrap(), num_activations as usize diff --git a/src/fetch/mod.rs b/src/fetch/mod.rs index 38305020..e5113747 100644 --- a/src/fetch/mod.rs +++ b/src/fetch/mod.rs @@ -9,8 +9,8 @@ use tracing::{debug, info, warn}; use crate::config::Config; use crate::push::{PushError, PushPool}; -use crate::store::activation::InflightActivation; -use crate::store::traits::InflightActivationStore; +use crate::store::activation::Activation; +use crate::store::traits::ClaimStore; use crate::store::types::BucketRange; /// This value should be a power of two. If it decreases, some ranges will no longer be queried. @@ -49,12 +49,12 @@ pub fn bucket_range_for_fetch_thread(thread_index: usize, fetch_threads: usize) #[async_trait] pub trait TaskPusher { /// Submit a single task to the push pool. - async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError>; + async fn submit_task(&self, activation: Activation) -> Result<(), PushError>; } #[async_trait] impl TaskPusher for PushPool { - async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError> { + async fn submit_task(&self, activation: Activation) -> Result<(), PushError> { self.submit(activation).await } } @@ -62,7 +62,7 @@ impl TaskPusher for PushPool { /// Wrapper around `config.fetch_threads` asynchronous tasks, each of which fetches batches of pending activations from the store, passes them to the push pool, and repeats. pub struct FetchPool { /// Inflight activation store. - store: Arc, + store: Arc, /// Pool of push threads that push activations to the worker service. pusher: Arc, @@ -73,11 +73,7 @@ pub struct FetchPool { impl FetchPool { /// Initialize a new fetch pool. - pub fn new( - store: Arc, - config: Arc, - pusher: Arc, - ) -> Self { + pub fn new(store: Arc, config: Arc, pusher: Arc) -> Self { Self { store, config, diff --git a/src/fetch/tests.rs b/src/fetch/tests.rs index ded1e681..8332b5c3 100644 --- a/src/fetch/tests.rs +++ b/src/fetch/tests.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use anyhow::{Error, anyhow}; -use chrono::{DateTime, Utc}; use tokio::sync::Mutex; use tokio::time::{Duration, sleep}; use tonic::async_trait; @@ -9,15 +8,15 @@ use tonic::async_trait; use super::*; use crate::config::Config; use crate::push::PushError; -use crate::store::activation::{InflightActivation, InflightActivationStatus}; -use crate::store::traits::InflightActivationStore; -use crate::store::types::{BucketRange, FailedTasksForwarder}; +use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::traits::ClaimStore; +use crate::store::types::BucketRange; use crate::test_utils::make_activations; /// Store stub that returns one activation once OR is always empty OR always fails. struct MockStore { /// A single (optional) pending activation. - pending: Mutex>, + pending: Mutex>, /// Should operations fail? fail: bool, @@ -31,7 +30,7 @@ impl MockStore { } } - fn one(activation: InflightActivation) -> Self { + fn one(activation: Activation) -> Self { Self { pending: Mutex::new(Some(activation)), fail: false, @@ -47,31 +46,7 @@ impl MockStore { } #[async_trait] -impl InflightActivationStore for MockStore { - fn assign_partitions(&self, _partitions: Vec) -> Result<(), Error> { - Ok(()) - } - - async fn vacuum_db(&self) -> Result<(), Error> { - unimplemented!() - } - - async fn full_vacuum_db(&self) -> Result<(), Error> { - unimplemented!() - } - - async fn db_size(&self) -> Result { - unimplemented!() - } - - async fn get_by_id(&self, _id: &str) -> Result, Error> { - unimplemented!() - } - - async fn store(&self, _batch: Vec) -> Result { - unimplemented!() - } - +impl ClaimStore for MockStore { async fn claim_activations( &self, _application: Option<&str>, @@ -79,7 +54,7 @@ impl InflightActivationStore for MockStore { _limit: Option, _bucket: Option, mark_processing: bool, - ) -> Result, Error> { + ) -> Result, Error> { if self.fail { return Err(anyhow!("mock store error")); } @@ -87,95 +62,15 @@ impl InflightActivationStore for MockStore { Ok(match self.pending.lock().await.take() { Some(mut a) => { a.status = if mark_processing { - InflightActivationStatus::Processing + ActivationStatus::Processing } else { - InflightActivationStatus::Claimed + ActivationStatus::Claimed }; vec![a] } None => vec![], }) } - - async fn mark_activation_processing(&self, _id: &str) -> Result<(), Error> { - Ok(()) - } - - async fn pending_activation_max_lag(&self, _now: &DateTime) -> f64 { - unimplemented!() - } - - async fn count_by_status(&self, _status: InflightActivationStatus) -> Result { - unimplemented!() - } - - async fn count(&self) -> Result { - unimplemented!() - } - - async fn set_status( - &self, - _id: &str, - _status: InflightActivationStatus, - ) -> Result, Error> { - unimplemented!() - } - - async fn set_processing_deadline( - &self, - _id: &str, - _deadline: Option>, - ) -> Result<(), Error> { - unimplemented!() - } - - async fn delete_activation(&self, _id: &str) -> Result<(), Error> { - unimplemented!() - } - - async fn get_retry_activations(&self) -> Result, Error> { - unimplemented!() - } - - async fn clear(&self) -> Result<(), Error> { - unimplemented!() - } - - async fn handle_claim_expiration(&self) -> Result { - unimplemented!() - } - - async fn handle_processing_deadline(&self) -> Result { - unimplemented!() - } - - async fn handle_processing_attempts(&self) -> Result { - unimplemented!() - } - - async fn handle_expires_at(&self) -> Result { - unimplemented!() - } - - async fn handle_delay_until(&self) -> Result { - unimplemented!() - } - - async fn handle_failed_tasks(&self) -> Result { - unimplemented!() - } - - async fn mark_completed(&self, _ids: Vec) -> Result { - unimplemented!() - } - - async fn remove_completed(&self) -> Result { - unimplemented!() - } - - async fn remove_killswitched(&self, _killswitched_tasks: Vec) -> Result { - unimplemented!() - } } /// Records task IDs passed to `push_task`. If `fail` is true, returns an error. @@ -196,7 +91,7 @@ impl RecordingPusher { #[async_trait] impl TaskPusher for RecordingPusher { - async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError> { + async fn submit_task(&self, activation: Activation) -> Result<(), PushError> { self.pushed_ids.lock().await.push(activation.id.clone()); if self.fail { @@ -218,7 +113,7 @@ fn test_config() -> Arc { #[tokio::test] async fn fetch_pool_delivers_activation_to_pusher() { let activation = make_activations(1).remove(0); - let store: Arc = Arc::new(MockStore::one(activation.clone())); + let store = Arc::new(MockStore::one(activation.clone())); let pusher = Arc::new(RecordingPusher::new(false)); let pool = FetchPool::new(store, test_config(), pusher.clone()); @@ -233,7 +128,7 @@ async fn fetch_pool_delivers_activation_to_pusher() { #[tokio::test] async fn fetch_pool_calls_pusher_once_when_push_errors() { let activation = make_activations(1).remove(0); - let store: Arc = Arc::new(MockStore::one(activation)); + let store = Arc::new(MockStore::one(activation)); let pusher = Arc::new(RecordingPusher::new(true)); let pool = FetchPool::new(store, test_config(), pusher.clone()); @@ -247,7 +142,7 @@ async fn fetch_pool_calls_pusher_once_when_push_errors() { #[tokio::test] async fn fetch_pool_skips_pusher_when_store_errors() { - let store: Arc = Arc::new(MockStore::error()); + let store = Arc::new(MockStore::error()); let pusher = Arc::new(RecordingPusher::new(false)); let pool = FetchPool::new(store, test_config(), pusher.clone()); @@ -261,7 +156,7 @@ async fn fetch_pool_skips_pusher_when_store_errors() { #[tokio::test] async fn fetch_pool_skips_pusher_when_no_pending() { - let store: Arc = Arc::new(MockStore::empty()); + let store = Arc::new(MockStore::empty()); let pusher = Arc::new(RecordingPusher::new(false)); let pool = FetchPool::new(store, test_config(), pusher.clone()); diff --git a/src/grpc/server.rs b/src/grpc/server.rs index d4f5161e..155acdce 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -10,12 +10,12 @@ use std::time::Instant; use tonic::{Request, Response, Status}; use crate::config::{Config, DeliveryMode}; -use crate::store::activation::InflightActivationStatus; -use crate::store::traits::InflightActivationStore; +use crate::store::activation::ActivationStatus; +use crate::store::traits::PullStore; use tracing::{error, instrument, warn}; pub struct TaskbrokerServer { - pub store: Arc, + pub store: Arc, pub config: Arc, } @@ -84,19 +84,16 @@ impl ConsumerService for TaskbrokerServer { let start_time = Instant::now(); let id = request.get_ref().id.clone(); - let status: InflightActivationStatus = - TaskActivationStatus::try_from(request.get_ref().status) - .map_err(|e| { - Status::invalid_argument(format!("Unable to deserialize status: {e:?}")) - })? - .into(); + let status: ActivationStatus = TaskActivationStatus::try_from(request.get_ref().status) + .map_err(|e| Status::invalid_argument(format!("Unable to deserialize status: {e:?}")))? + .into(); if !status.is_conclusion() { return Err(Status::invalid_argument(format!( "Invalid status, expects 3 (Failure), 4 (Retry), or 5 (Complete), but got: {status:?}" ))); } - if status == InflightActivationStatus::Failure { + if status == ActivationStatus::Failure { metrics::counter!("grpc_server.set_status.failure").increment(1); } diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index 3932de41..d073b544 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use crate::config::{Config, DeliveryMode}; use crate::grpc::server::TaskbrokerServer; -use crate::store::activation::InflightActivationStatus; +use crate::store::activation::ActivationStatus; use prost::Message; use rstest::rstest; use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService; @@ -110,7 +110,7 @@ async fn test_get_task_success(#[case] adapter: &str) { let config = create_config(); let activations = make_activations(1); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store: store.clone(), @@ -128,7 +128,7 @@ async fn test_get_task_success(#[case] adapter: &str) { assert!(task.id == "id_0"); let row = store.get_by_id("id_0").await.unwrap().expect("claimed row"); - assert_eq!(row.status, InflightActivationStatus::Processing); + assert_eq!(row.status, ActivationStatus::Processing); } #[tokio::test] @@ -147,7 +147,7 @@ async fn test_get_task_with_application_success(#[case] adapter: &str) { activations[1].activation = payload.encode_to_vec(); activations[1].application = "hammers".into(); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; let request = GetTaskRequest { @@ -175,7 +175,7 @@ async fn test_get_task_with_namespace_requires_application(#[case] adapter: &str let activations = make_activations(2); let namespace = activations[0].namespace.clone(); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; let request = GetTaskRequest { @@ -199,7 +199,7 @@ async fn test_set_task_status_success(#[case] adapter: &str) { let config = create_config(); let activations = make_activations(2); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; @@ -246,7 +246,7 @@ async fn test_set_task_status_with_application(#[case] adapter: &str) { activations[1].activation = payload.encode_to_vec(); activations[1].application = "hammers".into(); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; let request = SetTaskStatusRequest { @@ -285,7 +285,7 @@ async fn test_set_task_status_with_application_no_match(#[case] adapter: &str) { activations[1].activation = payload.encode_to_vec(); activations[1].application = "hammers".into(); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; // Request a task from an application without any activations. @@ -314,7 +314,7 @@ async fn test_set_task_status_with_namespace_requires_application(#[case] adapte let activations = make_activations(2); let namespace = activations[0].namespace.clone(); - store.store(activations).await.unwrap(); + store.write(activations).await.unwrap(); let service = TaskbrokerServer { store, config }; let request = SetTaskStatusRequest { diff --git a/src/kafka/consumer.rs b/src/kafka/consumer.rs index 8de67229..bbf452eb 100644 --- a/src/kafka/consumer.rs +++ b/src/kafka/consumer.rs @@ -1,4 +1,4 @@ -use crate::store::traits::InflightActivationStore; +use crate::store::traits::IngestStore; use anyhow::{Error, anyhow}; use futures::{ Stream, StreamExt, @@ -45,7 +45,7 @@ use tracing::{debug, error, info, instrument, warn}; pub async fn start_consumer( topics: &[&str], kafka_client_config: &ClientConfig, - activation_store: Arc, + activation_store: Arc, spawn_actors: impl FnMut( Arc>, &BTreeSet<(String, i32)>, @@ -343,7 +343,7 @@ enum ConsumerState { pub async fn handle_events( consumer: Arc>, events: UnboundedReceiver<(Event, SyncSender<()>)>, - activation_store: Arc, + activation_store: Arc, shutdown_client: oneshot::Sender<()>, mut spawn_actors: impl FnMut( Arc>, @@ -381,7 +381,7 @@ pub async fn handle_events( for (_, partition) in tpl.iter() { partitions.push(*partition); } - activation_store.assign_partitions(partitions).unwrap(); + activation_store.assign_partitions(partitions).await.unwrap(); ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl) } (ConsumerState::Ready, Event::Revoke(_)) => { @@ -396,13 +396,13 @@ pub async fn handle_events( tpl == revoked, "Revoked TPL should be equal to the subset of TPL we're consuming from" ); - activation_store.assign_partitions(vec![]).unwrap(); + activation_store.assign_partitions(vec![]).await.unwrap(); handles.shutdown(CALLBACK_DURATION).await; metrics::gauge!("arroyo.consumer.current_partitions").set(0); ConsumerState::Ready } (ConsumerState::Consuming(handles, _), Event::Shutdown) => { - activation_store.assign_partitions(vec![]).unwrap(); + activation_store.assign_partitions(vec![]).await.unwrap(); handles.shutdown(CALLBACK_DURATION).await; debug!("Signaling shutdown to client..."); shutdown_client.take(); diff --git a/src/kafka/deserialize_activation.rs b/src/kafka/deserialize_activation.rs index 2288d534..b9df63b6 100644 --- a/src/kafka/deserialize_activation.rs +++ b/src/kafka/deserialize_activation.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use crate::config::Config; use crate::fetch::MAX_FETCH_THREADS; -use crate::store::activation::{InflightActivation, InflightActivationStatus}; +use crate::store::activation::{Activation, ActivationStatus}; use anyhow::{Error, anyhow}; use chrono::{DateTime, Utc}; use prost::Message as _; @@ -32,7 +32,7 @@ pub fn bucket_from_id(id: &str) -> i16 { pub fn new( config: DeserializeActivationConfig, -) -> impl Fn(Arc) -> Result { +) -> impl Fn(Arc) -> Result { move |msg: Arc| { let Some(payload) = msg.payload() else { return Err(anyhow!("Message has no payload")); @@ -72,11 +72,11 @@ pub fn new( activation_time + delay }); - let status = delay_until.map_or(InflightActivationStatus::Pending, |delay_until| { + let status = delay_until.map_or(ActivationStatus::Pending, |delay_until| { if Utc::now() > delay_until { - InflightActivationStatus::Pending + ActivationStatus::Pending } else { - InflightActivationStatus::Delay + ActivationStatus::Delay } }); @@ -89,7 +89,7 @@ pub fn new( let bucket = bucket_from_id(&activation.id); - Ok(InflightActivation { + Ok(Activation { id: activation.id.clone(), activation: payload.to_vec(), status, @@ -122,9 +122,7 @@ mod tests { use rdkafka::{Timestamp, message::OwnedMessage}; use sentry_protos::taskbroker::v1::TaskActivation; - use crate::{ - store::activation::InflightActivationStatus, test_utils::generate_unique_namespace, - }; + use crate::{store::activation::ActivationStatus, test_utils::generate_unique_namespace}; use super::{Config, DeserializeActivationConfig, new}; @@ -263,7 +261,7 @@ mod tests { delta.num_seconds() >= 99, "Should have ~100 seconds of delay from received_at" ); - assert_eq!(inflight.status, InflightActivationStatus::Pending) + assert_eq!(inflight.status, ActivationStatus::Pending) } #[test] @@ -309,7 +307,7 @@ mod tests { delta.num_seconds() >= 99, "Should have ~100 seconds of delay from received_at" ); - assert_eq!(inflight.status, InflightActivationStatus::Delay) + assert_eq!(inflight.status, ActivationStatus::Delay) } #[test] @@ -356,6 +354,6 @@ mod tests { delta.num_seconds() <= (config.max_delayed_task_allowed_sec as f64 * 1.1) as i64, "Should have approxmiately max_delayed_task_allowed_sec of delay from received_at" ); - assert_eq!(inflight.status, InflightActivationStatus::Delay) + assert_eq!(inflight.status, ActivationStatus::Delay) } } diff --git a/src/kafka/inflight_activation_batcher.rs b/src/kafka/inflight_activation_batcher.rs index 79f50328..a6f5788e 100644 --- a/src/kafka/inflight_activation_batcher.rs +++ b/src/kafka/inflight_activation_batcher.rs @@ -1,6 +1,4 @@ -use crate::{ - config::Config, runtime_config::RuntimeConfigManager, store::activation::InflightActivation, -}; +use crate::{config::Config, runtime_config::RuntimeConfigManager, store::activation::Activation}; use chrono::Utc; use futures::future::join_all; use rdkafka::config::ClientConfig; @@ -43,7 +41,7 @@ impl ActivationBatcherConfig { } pub struct InflightActivationBatcher { - batch: Vec, + batch: Vec, batch_size: usize, forward_batch: Vec>, // payload config: ActivationBatcherConfig, @@ -81,9 +79,9 @@ impl InflightActivationBatcher { } impl Reducer for InflightActivationBatcher { - type Input = InflightActivation; + type Input = Activation; - type Output = Vec; + type Output = Vec; async fn reduce(&mut self, t: Self::Input) -> Result<(), anyhow::Error> { let runtime_config = self.runtime_config_manager.read().await; @@ -220,7 +218,7 @@ mod tests { use std::sync::Arc; use crate::{ - store::activation::InflightActivationBuilder, + store::activation::ActivationBuilder, test_utils::{TaskActivationBuilder, generate_unique_namespace}, }; @@ -244,7 +242,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = InflightActivationBuilder::new() + let inflight_activation_0 = ActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace(&namespace) @@ -267,7 +265,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = InflightActivationBuilder::new() + let inflight_activation_0 = ActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace(&namespace) @@ -294,7 +292,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = InflightActivationBuilder::new() + let inflight_activation_0 = ActivationBuilder::new() .id("0") .taskname("taskname") .namespace(&namespace) @@ -322,13 +320,13 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = InflightActivationBuilder::new() + let inflight_activation_0 = ActivationBuilder::new() .id("0") .taskname("taskname") .namespace(&namespace) .build(TaskActivationBuilder::new()); - let inflight_activation_1 = InflightActivationBuilder::new() + let inflight_activation_1 = ActivationBuilder::new() .id("1") .taskname("taskname") .namespace(&namespace) @@ -363,13 +361,13 @@ demoted_topic: taskworker-demoted"#; assert_eq!(batcher.producer_cluster, config.kafka_cluster.clone()); - let inflight_activation_0 = InflightActivationBuilder::new() + let inflight_activation_0 = ActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace("bad_namespace") .build(TaskActivationBuilder::new()); - let inflight_activation_1 = InflightActivationBuilder::new() + let inflight_activation_1 = ActivationBuilder::new() .id("1") .taskname("good_task") .namespace("good_namespace") diff --git a/src/kafka/inflight_activation_writer.rs b/src/kafka/inflight_activation_writer.rs index e6a2e20e..708d05ca 100644 --- a/src/kafka/inflight_activation_writer.rs +++ b/src/kafka/inflight_activation_writer.rs @@ -9,9 +9,11 @@ use tracing::{debug, error, instrument}; use crate::{ config::Config, - store::activation::{InflightActivation, InflightActivationStatus}, - store::traits::InflightActivationStore, - store::types::DepthCounts, + store::{ + activation::{Activation, ActivationStatus}, + traits::{CountStore, IngestStore}, + types::DepthCounts, + }, }; use super::consumer::{ @@ -42,14 +44,14 @@ impl ActivationWriterConfig { } } -pub struct InflightActivationWriter { +pub struct InflightActivationWriter { config: ActivationWriterConfig, - store: Arc, - batch: Option>, + store: Arc, + batch: Option>, } -impl InflightActivationWriter { - pub fn new(store: Arc, config: ActivationWriterConfig) -> Self { +impl InflightActivationWriter { + pub fn new(store: Arc, config: ActivationWriterConfig) -> Self { Self { config, store, @@ -58,8 +60,8 @@ impl InflightActivationWriter { } } -impl Reducer for InflightActivationWriter { - type Input = Vec; +impl Reducer for InflightActivationWriter { + type Input = Vec; type Output = (); @@ -99,7 +101,7 @@ impl Reducer for InflightActivationWriter { processing + claimed >= self.config.max_processing_activations; let exceeded_db_size = if let Some(db_max_size) = self.config.db_max_size { self.store - .db_size() + .size() .await .expect("Error getting database size") >= db_max_size @@ -110,10 +112,10 @@ impl Reducer for InflightActivationWriter { // Check if the entire batch is either pending or delay let has_delay = batch .iter() - .any(|activation| activation.status == InflightActivationStatus::Delay); + .any(|activation| activation.status == ActivationStatus::Delay); let has_pending = batch .iter() - .any(|activation| activation.status == InflightActivationStatus::Pending); + .any(|activation| activation.status == ActivationStatus::Pending); // Backpressure if any of these conditions are met: // 1. The processing limit is exceeded @@ -145,7 +147,7 @@ impl Reducer for InflightActivationWriter { let batch = self.batch.clone().unwrap(); let write_to_store_start = Instant::now(); - let res = self.store.store(batch.clone()).await; + let res = self.store.write(batch.clone()).await; match res { Ok(entries) => { self.batch.take(); @@ -200,7 +202,7 @@ mod tests { use super::{ActivationWriterConfig, InflightActivationWriter, Reducer}; use crate::{ - store::activation::{InflightActivationBuilder, InflightActivationStatus}, + store::activation::{ActivationBuilder, ActivationStatus}, test_utils::{ TaskActivationBuilder, create_test_store, generate_unique_namespace, make_activations, }, @@ -226,13 +228,13 @@ mod tests { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) .received_at(received_at) .build(TaskActivationBuilder::new()), - InflightActivationBuilder::new() + ActivationBuilder::new() .id("1") .taskname("delay_task") .namespace(&namespace) @@ -245,7 +247,7 @@ mod tests { let count_pending = writer.store.count_pending_activations().await.unwrap(); let count_delay = writer .store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); assert_eq!(count_pending + count_delay, 2); @@ -272,7 +274,7 @@ mod tests { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) @@ -307,12 +309,12 @@ mod tests { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) .received_at(received_at) - .status(InflightActivationStatus::Delay) + .status(ActivationStatus::Delay) .build(TaskActivationBuilder::new()), ]; @@ -320,7 +322,7 @@ mod tests { writer.flush().await.unwrap(); let count_delay = writer .store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); assert_eq!(count_delay, 1); @@ -347,13 +349,13 @@ mod tests { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) .received_at(received_at) .build(TaskActivationBuilder::new()), - InflightActivationBuilder::new() + ActivationBuilder::new() .id("1") .taskname("delay_task") .namespace(&namespace) @@ -367,7 +369,7 @@ mod tests { assert_eq!(count_pending, 0); let count_delay = writer .store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); assert_eq!(count_delay, 0); @@ -396,13 +398,13 @@ mod tests { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) .received_at(received_at) .build(TaskActivationBuilder::new()), - InflightActivationBuilder::new() + ActivationBuilder::new() .id("1") .taskname("pending_task") .namespace(&namespace) @@ -416,7 +418,7 @@ mod tests { assert_eq!(count_pending, 2); let count_delay = writer .store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); assert_eq!(count_delay, 0); @@ -441,25 +443,25 @@ mod tests { let received_at = DateTime::from_timestamp_nanos(0); let namespace = generate_unique_namespace(); - let existing_activation = InflightActivationBuilder::new() + let existing_activation = ActivationBuilder::new() .id("existing") .taskname("existing_task") .namespace(&namespace) .received_at(received_at) - .status(InflightActivationStatus::Processing) + .status(ActivationStatus::Processing) .build(TaskActivationBuilder::new()); store.store(vec![existing_activation]).await.unwrap(); let mut writer = InflightActivationWriter::new(store.clone(), writer_config); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("0") .taskname("pending_task") .namespace(&namespace) .received_at(received_at) .build(TaskActivationBuilder::new()), - InflightActivationBuilder::new() + ActivationBuilder::new() .id("1") .taskname("delay_task") .namespace(&namespace) @@ -476,13 +478,13 @@ mod tests { assert_eq!(count_pending, 0); let count_delay = writer .store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); assert_eq!(count_delay, 0); let count_processing = writer .store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(); // Only the existing processing activation should remain, new ones should be blocked diff --git a/src/main.rs b/src/main.rs index 9db67cf5..c5475662 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,6 +7,7 @@ use taskbroker::kafka::inflight_activation_batcher::{ ActivationBatcherConfig, InflightActivationBatcher, }; use taskbroker::push::PushPool; +use taskbroker::store::traits::Store; use taskbroker::upkeep::upkeep; use tokio::signal::unix::SignalKind; use tokio::task::JoinHandle; @@ -32,11 +33,8 @@ use taskbroker::logging; use taskbroker::metrics; use taskbroker::processing_strategy; use taskbroker::runtime_config::RuntimeConfigManager; -use taskbroker::store::adapters::postgres::{ - PostgresActivationStore, PostgresActivationStoreConfig, -}; +use taskbroker::store::adapters::postgres::{PostgresActivationStoreConfig, PostgresStore}; use taskbroker::store::adapters::sqlite::{InflightActivationStoreConfig, SqliteActivationStore}; -use taskbroker::store::traits::InflightActivationStore; use taskbroker::{Args, get_version}; use tonic_health::ServingStatus; @@ -67,7 +65,7 @@ async fn main() -> Result<(), Error> { logging::init(logging::LoggingConfig::from_config(&config)); metrics::init(metrics::MetricsConfig::from_config(&config)); - let store: Arc = match config.database_adapter { + let store: Arc = match config.database_adapter { DatabaseAdapter::Sqlite => Arc::new( SqliteActivationStore::new( &config.db_path, @@ -75,10 +73,9 @@ async fn main() -> Result<(), Error> { ) .await?, ), - DatabaseAdapter::Postgres => Arc::new( - PostgresActivationStore::new(PostgresActivationStoreConfig::from_config(&config)) - .await?, - ), + DatabaseAdapter::Postgres => { + Arc::new(PostgresStore::new(PostgresActivationStoreConfig::from_config(&config)).await?) + } }; // If this is an environment where the topics might not exist, check and create them. diff --git a/src/push/mod.rs b/src/push/mod.rs index f45efbd7..afad1e53 100644 --- a/src/push/mod.rs +++ b/src/push/mod.rs @@ -16,8 +16,8 @@ use tonic::transport::Channel; use tracing::{debug, error, info}; use crate::config::Config; -use crate::store::activation::InflightActivation; -use crate::store::traits::InflightActivationStore; +use crate::store::activation::Activation; +use crate::store::traits::PushStore; type HmacSha256 = Hmac; @@ -42,7 +42,7 @@ pub enum PushError { Timeout, /// Channel disconnected (no receivers) or another failure. - Channel(SendError), + Channel(SendError), } /// Thin interface for the worker client. It mostly serves to enable proper unit testing, but it also decouples the actual client implementation from our pushing logic. @@ -82,21 +82,21 @@ impl WorkerClient for WorkerServiceClient { /// Wrapper around `config.push_threads` asynchronous tasks, each of which receives an activation from the channel, sends it to the worker service, and repeats. pub struct PushPool { /// The sending end of a channel that accepts task activations. - sender: Sender, + sender: Sender, /// The receiving end of a channel that accepts task activations. - receiver: Receiver, + receiver: Receiver, /// Taskbroker configuration. config: Arc, /// Activation store, which we need for marking tasks as sent. - store: Arc, + store: Arc, } impl PushPool { /// Initialize a new push pool. - pub fn new(config: Arc, store: Arc) -> Self { + pub fn new(config: Arc, store: Arc) -> Self { let (sender, receiver) = flume::bounded(config.push_queue_size); Self { @@ -178,7 +178,7 @@ impl PushPool { metrics::counter!("push.delivery", "result" => "ok").increment(1); debug!(task_id = %id, "Activation sent to worker"); - if let Err(e) = store.mark_activation_processing(&id).await { + if let Err(e) = store.mark_processing(&id).await { metrics::counter!("push.mark_activation_processing", "result" => "error").increment(1); error!( @@ -222,7 +222,7 @@ impl PushPool { metrics::counter!("push.delivery", "result" => "ok").increment(1); debug!(task_id = %id, "Activation sent to worker"); - if let Err(e) = store.mark_activation_processing(&id).await { + if let Err(e) = store.mark_processing(&id).await { metrics::counter!("push.mark_activation_processing", "result" => "error").increment(1); error!( @@ -268,7 +268,7 @@ impl PushPool { } /// Send an activation to the internal asynchronous MPMC channel used by all running push threads. Times out after `config.push_queue_timeout_ms` milliseconds. - pub async fn submit(&self, activation: InflightActivation) -> Result<(), PushError> { + pub async fn submit(&self, activation: Activation) -> Result<(), PushError> { let duration = Duration::from_millis(self.config.push_queue_timeout_ms); let start = Instant::now(); @@ -298,7 +298,7 @@ impl PushPool { /// Decode task activation and push it to a worker. async fn push_task( worker: &mut W, - activation: InflightActivation, + activation: Activation, callback_url: String, timeout: Duration, grpc_shared_secret: &[String], diff --git a/src/store/activation.rs b/src/store/activation.rs index 52c2326b..f22c8012 100644 --- a/src/store/activation.rs +++ b/src/store/activation.rs @@ -8,7 +8,7 @@ use std::str::FromStr; /// The members of this enum should be a superset of the members /// of `InflightActivationStatus` in `sentry_protos`. #[derive(Clone, Copy, Debug, PartialEq, Eq, Type)] -pub enum InflightActivationStatus { +pub enum ActivationStatus { /// Unused but necessary to align with sentry-protos Unspecified, Pending, @@ -20,59 +20,57 @@ pub enum InflightActivationStatus { Delay, } -impl Display for InflightActivationStatus { +impl Display for ActivationStatus { fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { write!(f, "{:?}", self) } } -impl FromStr for InflightActivationStatus { +impl FromStr for ActivationStatus { type Err = String; fn from_str(s: &str) -> Result { if s == "Unspecified" { - Ok(InflightActivationStatus::Unspecified) + Ok(ActivationStatus::Unspecified) } else if s == "Pending" { - Ok(InflightActivationStatus::Pending) + Ok(ActivationStatus::Pending) } else if s == "Claimed" { - Ok(InflightActivationStatus::Claimed) + Ok(ActivationStatus::Claimed) } else if s == "Processing" { - Ok(InflightActivationStatus::Processing) + Ok(ActivationStatus::Processing) } else if s == "Failure" { - Ok(InflightActivationStatus::Failure) + Ok(ActivationStatus::Failure) } else if s == "Retry" { - Ok(InflightActivationStatus::Retry) + Ok(ActivationStatus::Retry) } else if s == "Complete" { - Ok(InflightActivationStatus::Complete) + Ok(ActivationStatus::Complete) } else if s == "Delay" { - Ok(InflightActivationStatus::Delay) + Ok(ActivationStatus::Delay) } else { Err(format!("Unknown inflight activation status string: {}", s)) } } } -impl InflightActivationStatus { +impl ActivationStatus { /// Is the current value a 'conclusion' status that can be supplied over GRPC. pub fn is_conclusion(&self) -> bool { matches!( self, - InflightActivationStatus::Complete - | InflightActivationStatus::Retry - | InflightActivationStatus::Failure + ActivationStatus::Complete | ActivationStatus::Retry | ActivationStatus::Failure ) } } -impl From for InflightActivationStatus { +impl From for ActivationStatus { fn from(item: TaskActivationStatus) -> Self { match item { - TaskActivationStatus::Unspecified => InflightActivationStatus::Unspecified, - TaskActivationStatus::Pending => InflightActivationStatus::Pending, - TaskActivationStatus::Processing => InflightActivationStatus::Processing, - TaskActivationStatus::Failure => InflightActivationStatus::Failure, - TaskActivationStatus::Retry => InflightActivationStatus::Retry, - TaskActivationStatus::Complete => InflightActivationStatus::Complete, + TaskActivationStatus::Unspecified => ActivationStatus::Unspecified, + TaskActivationStatus::Pending => ActivationStatus::Pending, + TaskActivationStatus::Processing => ActivationStatus::Processing, + TaskActivationStatus::Failure => ActivationStatus::Failure, + TaskActivationStatus::Retry => ActivationStatus::Retry, + TaskActivationStatus::Complete => ActivationStatus::Complete, } } } @@ -81,7 +79,7 @@ impl From for InflightActivationStatus { #[builder(pattern = "owned")] #[builder(build_fn(name = "_build"))] #[builder(field(public))] -pub struct InflightActivation { +pub struct Activation { #[builder(setter(into))] pub id: String, @@ -102,8 +100,8 @@ pub struct InflightActivation { pub activation: Vec, /// The current status of the activation - #[builder(default = InflightActivationStatus::Pending)] - pub status: InflightActivationStatus, + #[builder(default = ActivationStatus::Pending)] + pub status: ActivationStatus, /// The partition the activation was received from #[builder(default = 0)] @@ -164,7 +162,7 @@ pub struct InflightActivation { pub bucket: i16, } -impl InflightActivation { +impl Activation { /// The number of milliseconds between an activation's received timestamp /// and the provided datetime pub fn received_latency(&self, now: DateTime) -> i64 { diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 77913711..00eb72e0 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -14,8 +14,10 @@ use std::time::Instant; use tracing::{instrument, warn}; use crate::config::Config; -use crate::store::activation::{InflightActivation, InflightActivationStatus}; -use crate::store::traits::InflightActivationStore; +use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::traits::{ + ClaimStore, CountStore, IngestStore, PullStore, PushStore, Store, TestStore, UpkeepStore, +}; use crate::store::types::{BucketRange, DepthCounts, FailedTasksForwarder}; #[derive(Debug, FromRow)] @@ -42,10 +44,10 @@ struct TableRow { pub bucket: i16, } -impl TryFrom for TableRow { +impl TryFrom for TableRow { type Error = anyhow::Error; - fn try_from(value: InflightActivation) -> Result { + fn try_from(value: Activation) -> Result { Ok(Self { id: value.id, activation: value.activation, @@ -70,12 +72,12 @@ impl TryFrom for TableRow { } } -impl From for InflightActivation { +impl From for Activation { fn from(value: TableRow) -> Self { Self { id: value.id, activation: value.activation, - status: InflightActivationStatus::from_str(&value.status).unwrap(), + status: ActivationStatus::from_str(&value.status).unwrap(), partition: value.partition, offset: value.offset, added_at: value.added_at, @@ -162,14 +164,14 @@ impl PostgresActivationStoreConfig { } } -pub struct PostgresActivationStore { +pub struct PostgresStore { read_pool: PgPool, write_pool: PgPool, config: PostgresActivationStoreConfig, partitions: RwLock>, } -impl PostgresActivationStore { +impl PostgresStore { async fn acquire_write_conn_metric( &self, caller: &'static str, @@ -241,155 +243,7 @@ impl PostgresActivationStore { } #[async_trait] -impl InflightActivationStore for PostgresActivationStore { - /// Trigger incremental vacuum to reclaim free pages in the database. - /// Depending on config data, will either vacuum a set number of - /// pages or attempt to reclaim all free pages. - #[instrument(skip_all)] - async fn vacuum_db(&self) -> Result<(), Error> { - // TODO: Remove - Ok(()) - } - - /// Perform a full vacuum on the database. - async fn full_vacuum_db(&self) -> Result<(), Error> { - // TODO: Remove - Ok(()) - } - - /// Get the size of the database in bytes based on SQLite metadata queries. - async fn db_size(&self) -> Result { - let row_result: (i64,) = sqlx::query_as("SELECT pg_database_size($1) as size") - .bind(&self.config.pg_database_name) - .fetch_one(&self.read_pool) - .await?; - if row_result.0 < 0 { - return Ok(0); - } - Ok(row_result.0 as u64) - } - - /// Get an activation by id. Primarily used for testing - async fn get_by_id(&self, id: &str) -> Result, Error> { - let row_result: Option = sqlx::query_as( - " - SELECT id, - activation, - partition, - kafka_offset AS offset, - added_at, - received_at, - processing_attempts, - expires_at, - delay_until, - processing_deadline_duration, - processing_deadline, - claim_expires_at, - status, - at_most_once, - application, - namespace, - taskname, - on_attempts_exceeded, - bucket - FROM inflight_taskactivations - WHERE id = $1 - ", - ) - .bind(id) - .fetch_optional(&self.read_pool) - .await?; - - let Some(row) = row_result else { - return Ok(None); - }; - - Ok(Some(row.into())) - } - - fn assign_partitions(&self, partitions: Vec) -> Result<(), Error> { - let mut write_guard = self.partitions.write().unwrap(); - write_guard.clear(); - write_guard.extend(partitions); - Ok(()) - } - - #[instrument(skip_all)] - async fn store(&self, batch: Vec) -> Result { - if batch.is_empty() { - return Ok(0); - } - - let mut query_builder = QueryBuilder::::new( - " - INSERT INTO inflight_taskactivations - ( - id, - activation, - partition, - kafka_offset, - added_at, - received_at, - processing_attempts, - expires_at, - delay_until, - processing_deadline_duration, - processing_deadline, - claim_expires_at, - status, - at_most_once, - application, - namespace, - taskname, - on_attempts_exceeded, - bucket - ) - ", - ); - let rows = batch - .into_iter() - .map(TableRow::try_from) - .collect::, _>>()?; - let query = query_builder - .push_values(rows, |mut b, row| { - b.push_bind(row.id); - b.push_bind(row.activation); - b.push_bind(row.partition); - b.push_bind(row.offset); - b.push_bind(row.added_at); - b.push_bind(row.received_at); - b.push_bind(row.processing_attempts); - b.push_bind(row.expires_at); - b.push_bind(row.delay_until); - b.push_bind(row.processing_deadline_duration); - if let Some(deadline) = row.processing_deadline { - b.push_bind(deadline); - } else { - // Add a literal null - b.push("null"); - } - if let Some(exp) = row.claim_expires_at { - b.push_bind(exp); - } else { - b.push("null"); - } - b.push_bind(row.status); - b.push_bind(row.at_most_once); - b.push_bind(row.application); - b.push_bind(row.namespace); - b.push_bind(row.taskname); - b.push_bind(row.on_attempts_exceeded as i32); - b.push_bind(row.bucket); - }) - .push(" ON CONFLICT(id) DO NOTHING") - .build(); - - let mut conn = self.acquire_write_conn_metric("store").await?; - let result = query.execute(&mut *conn).await?; - - Ok(result.rows_affected()) - } - +impl ClaimStore for PostgresStore { #[instrument(skip_all)] async fn claim_activations( &self, @@ -398,7 +252,7 @@ impl InflightActivationStore for PostgresActivationStore { limit: Option, bucket: Option, mark_processing: bool, - ) -> Result, Error> { + ) -> Result, Error> { let now = Utc::now(); let grace_period = self.config.processing_deadline_grace_sec; @@ -410,7 +264,7 @@ impl InflightActivationStore for PostgresActivationStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push(" AND (expires_at IS NULL OR expires_at > "); query_builder.push_bind(now); query_builder.push(")"); @@ -456,7 +310,7 @@ impl InflightActivationStore for PostgresActivationStore { status = " )); - query_builder.push_bind(InflightActivationStatus::Processing.to_string()); + query_builder.push_bind(ActivationStatus::Processing.to_string()); } else { query_builder.push(format!( "UPDATE inflight_taskactivations @@ -465,7 +319,7 @@ impl InflightActivationStore for PostgresActivationStore { status = " )); - query_builder.push_bind(InflightActivationStatus::Claimed.to_string()); + query_builder.push_bind(ActivationStatus::Claimed.to_string()); } query_builder.push(" FROM selected_activations "); @@ -480,9 +334,12 @@ impl InflightActivationStore for PostgresActivationStore { Ok(rows.into_iter().map(|row| row.into()).collect()) } +} +#[async_trait] +impl PushStore for PostgresStore { #[instrument(skip_all)] - async fn mark_activation_processing(&self, id: &str) -> Result<(), Error> { + async fn mark_processing(&self, id: &str) -> Result<(), Error> { let mut conn = self .acquire_write_conn_metric("mark_activation_processing") .await?; @@ -495,9 +352,9 @@ impl InflightActivationStore for PostgresActivationStore { claim_expires_at = NULL WHERE id = $2 AND status = $3", )) - .bind(InflightActivationStatus::Processing.to_string()) + .bind(ActivationStatus::Processing.to_string()) .bind(id) - .bind(InflightActivationStatus::Claimed.to_string()) + .bind(ActivationStatus::Claimed.to_string()) .execute(&mut *conn) .await?; @@ -515,6 +372,71 @@ impl InflightActivationStore for PostgresActivationStore { Ok(()) } +} + +#[async_trait] +impl PullStore for PostgresStore { + /// Update the status of a specific activation. + #[instrument(skip_all)] + async fn set_status( + &self, + id: &str, + status: ActivationStatus, + ) -> Result, Error> { + let mut conn = self.acquire_write_conn_metric("set_status").await?; + let result: Option = sqlx::query_as( + "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *, kafka_offset AS offset", + ) + .bind(status.to_string()) + .bind(id) + .fetch_optional(&mut *conn) + .await?; + + let Some(row) = result else { + return Ok(None); + }; + + Ok(Some(row.into())) + } +} + +#[async_trait] +impl CountStore for PostgresStore { + #[instrument(skip_all)] + async fn count_by_status(&self, status: ActivationStatus) -> Result { + let mut query_builder = QueryBuilder::new( + "SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = ", + ); + query_builder.push_bind(status.to_string()); + self.add_partition_condition(&mut query_builder, false); + let result = query_builder + .build_query_as::<(i64,)>() + .fetch_one(&self.read_pool) + .await?; + Ok(result.0 as usize) + } + + async fn count(&self) -> Result { + let mut query_builder = + QueryBuilder::new("SELECT COUNT(*) as count FROM inflight_taskactivations"); + self.add_partition_condition(&mut query_builder, true); + let result = query_builder + .build_query_as::<(i64,)>() + .fetch_one(&self.read_pool) + .await?; + Ok(result.0 as usize) + } + + async fn size(&self) -> Result { + let row_result: (i64,) = sqlx::query_as("SELECT pg_database_size($1) as size") + .bind(&self.config.pg_database_name) + .fetch_one(&self.read_pool) + .await?; + if row_result.0 < 0 { + return Ok(0); + } + Ok(row_result.0 as u64) + } /// Get the age of the oldest pending activation in seconds. /// Only activations with status=pending and processing_attempts=0 are considered @@ -527,7 +449,7 @@ impl InflightActivationStore for PostgresActivationStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push(" AND processing_attempts = 0"); self.add_partition_condition(&mut query_builder, false); @@ -554,31 +476,6 @@ impl InflightActivationStore for PostgresActivationStore { } } - #[instrument(skip_all)] - async fn count_by_status(&self, status: InflightActivationStatus) -> Result { - let mut query_builder = QueryBuilder::new( - "SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = ", - ); - query_builder.push_bind(status.to_string()); - self.add_partition_condition(&mut query_builder, false); - let result = query_builder - .build_query_as::<(i64,)>() - .fetch_one(&self.read_pool) - .await?; - Ok(result.0 as usize) - } - - async fn count(&self) -> Result { - let mut query_builder = - QueryBuilder::new("SELECT COUNT(*) as count FROM inflight_taskactivations"); - self.add_partition_condition(&mut query_builder, true); - let result = query_builder - .build_query_as::<(i64,)>() - .fetch_one(&self.read_pool) - .await?; - Ok(result.0 as usize) - } - #[instrument(skip_all)] async fn count_depths(&self) -> Result { // Notice that statuses are embedded into the query for simplicity - if the enum is every changed, this must change too! @@ -604,24 +501,42 @@ impl InflightActivationStore for PostgresActivationStore { processing: row.3 as usize, }) } +} - /// Update the status of a specific activation - #[instrument(skip_all)] - async fn set_status( - &self, - id: &str, - status: InflightActivationStatus, - ) -> Result, Error> { - let mut conn = self.acquire_write_conn_metric("set_status").await?; - let result: Option = sqlx::query_as( - "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *, kafka_offset AS offset", +#[async_trait] +impl TestStore for PostgresStore { + /// Get an activation by id. Primarily used for testing + async fn get_by_id(&self, id: &str) -> Result, Error> { + let row_result: Option = sqlx::query_as( + " + SELECT id, + activation, + partition, + kafka_offset AS offset, + added_at, + received_at, + processing_attempts, + expires_at, + delay_until, + processing_deadline_duration, + processing_deadline, + claim_expires_at, + status, + at_most_once, + application, + namespace, + taskname, + on_attempts_exceeded, + bucket + FROM inflight_taskactivations + WHERE id = $1 + ", ) - .bind(status.to_string()) .bind(id) - .fetch_optional(&mut *conn) + .fetch_optional(&self.read_pool) .await?; - let Some(row) = result else { + let Some(row) = row_result else { return Ok(None); }; @@ -655,8 +570,135 @@ impl InflightActivationStore for PostgresActivationStore { Ok(()) } + // Used in tests + async fn clear(&self) -> Result<(), Error> { + let mut conn = self.acquire_write_conn_metric("clear").await?; + sqlx::query("TRUNCATE TABLE inflight_taskactivations") + .execute(&mut *conn) + .await?; + + Ok(()) + } + + // Used in tests + async fn remove_db(&self) -> Result<(), Error> { + self.read_pool.close().await; + self.write_pool.close().await; + let default_pool = create_default_postgres_pool(&self.config.pg_connection).await?; + let _ = sqlx::query(format!("DROP DATABASE {}", &self.config.pg_database_name).as_str()) + .bind(&self.config.pg_database_name) + .execute(&default_pool) + .await; + let _ = default_pool.close().await; + Ok(()) + } +} + +#[async_trait] +impl IngestStore for PostgresStore { + async fn assign_partitions(&self, partitions: Vec) -> Result<(), Error> { + let mut write_guard = self.partitions.write().unwrap(); + write_guard.clear(); + write_guard.extend(partitions); + Ok(()) + } + #[instrument(skip_all)] - async fn get_retry_activations(&self) -> Result, Error> { + async fn write(&self, batch: Vec) -> Result { + if batch.is_empty() { + return Ok(0); + } + + let mut query_builder = QueryBuilder::::new( + " + INSERT INTO inflight_taskactivations + ( + id, + activation, + partition, + kafka_offset, + added_at, + received_at, + processing_attempts, + expires_at, + delay_until, + processing_deadline_duration, + processing_deadline, + claim_expires_at, + status, + at_most_once, + application, + namespace, + taskname, + on_attempts_exceeded, + bucket + ) + ", + ); + let rows = batch + .into_iter() + .map(TableRow::try_from) + .collect::, _>>()?; + let query = query_builder + .push_values(rows, |mut b, row| { + b.push_bind(row.id); + b.push_bind(row.activation); + b.push_bind(row.partition); + b.push_bind(row.offset); + b.push_bind(row.added_at); + b.push_bind(row.received_at); + b.push_bind(row.processing_attempts); + b.push_bind(row.expires_at); + b.push_bind(row.delay_until); + b.push_bind(row.processing_deadline_duration); + if let Some(deadline) = row.processing_deadline { + b.push_bind(deadline); + } else { + // Add a literal null + b.push("null"); + } + if let Some(exp) = row.claim_expires_at { + b.push_bind(exp); + } else { + b.push("null"); + } + b.push_bind(row.status); + b.push_bind(row.at_most_once); + b.push_bind(row.application); + b.push_bind(row.namespace); + b.push_bind(row.taskname); + b.push_bind(row.on_attempts_exceeded as i32); + b.push_bind(row.bucket); + }) + .push(" ON CONFLICT(id) DO NOTHING") + .build(); + + let mut conn = self.acquire_write_conn_metric("store").await?; + let result = query.execute(&mut *conn).await?; + + Ok(result.rows_affected()) + } +} + +#[async_trait] +impl UpkeepStore for PostgresStore { + /// Trigger incremental vacuum to reclaim free pages in the database. + /// Depending on config data, will either vacuum a set number of + /// pages or attempt to reclaim all free pages. + #[instrument(skip_all)] + async fn vacuum_db(&self) -> Result<(), Error> { + // TODO: Remove + Ok(()) + } + + /// Perform a full vacuum on the database. + async fn full_vacuum_db(&self) -> Result<(), Error> { + // TODO: Remove + Ok(()) + } + + #[instrument(skip_all)] + async fn get_retry_activations(&self) -> Result, Error> { let mut query_builder = QueryBuilder::new( "SELECT id, activation, @@ -680,7 +722,7 @@ impl InflightActivationStore for PostgresActivationStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(InflightActivationStatus::Retry.to_string()); + query_builder.push_bind(ActivationStatus::Retry.to_string()); self.add_partition_condition(&mut query_builder, false); Ok(query_builder @@ -692,16 +734,6 @@ impl InflightActivationStore for PostgresActivationStore { .collect()) } - // Used in tests - async fn clear(&self) -> Result<(), Error> { - let mut conn = self.acquire_write_conn_metric("clear").await?; - sqlx::query("TRUNCATE TABLE inflight_taskactivations") - .execute(&mut *conn) - .await?; - - Ok(()) - } - /// Revert expired push claims back to pending status. #[instrument(skip_all)] async fn handle_claim_expiration(&self) -> Result { @@ -715,14 +747,14 @@ impl InflightActivationStore for PostgresActivationStore { SET claim_expires_at = null, status = ", ); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push( " WHERE claim_expires_at IS NOT NULL AND claim_expires_at < ", ); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(InflightActivationStatus::Claimed.to_string()); + query_builder.push_bind(ActivationStatus::Claimed.to_string()); self.add_partition_condition(&mut query_builder, false); let released = query_builder.build().execute(&mut *conn).await?; @@ -744,11 +776,11 @@ impl InflightActivationStore for PostgresActivationStore { "UPDATE inflight_taskactivations SET processing_deadline = null, status = ", ); - query_builder.push_bind(InflightActivationStatus::Failure.to_string()); + query_builder.push_bind(ActivationStatus::Failure.to_string()); query_builder.push(" WHERE processing_deadline < "); query_builder.push_bind(now); query_builder.push(" AND at_most_once = TRUE AND status = "); - query_builder.push_bind(InflightActivationStatus::Processing.to_string()); + query_builder.push_bind(ActivationStatus::Processing.to_string()); self.add_partition_condition(&mut query_builder, false); @@ -765,12 +797,12 @@ impl InflightActivationStore for PostgresActivationStore { "UPDATE inflight_taskactivations SET processing_deadline = null, status = ", ); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push(", processing_attempts = processing_attempts + 1"); query_builder.push(" WHERE processing_deadline < "); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(InflightActivationStatus::Processing.to_string()); + query_builder.push_bind(ActivationStatus::Processing.to_string()); self.add_partition_condition(&mut query_builder, false); let result = query_builder.build().execute(&mut *atomic).await; @@ -796,11 +828,11 @@ impl InflightActivationStore for PostgresActivationStore { "UPDATE inflight_taskactivations SET status = ", ); - query_builder.push_bind(InflightActivationStatus::Failure.to_string()); + query_builder.push_bind(ActivationStatus::Failure.to_string()); query_builder.push(" WHERE processing_attempts >= "); query_builder.push_bind(self.config.max_processing_attempts as i32); query_builder.push(" AND status = "); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); self.add_partition_condition(&mut query_builder, false); let processing_attempts_result = query_builder.build().execute(&mut *conn).await; @@ -823,7 +855,7 @@ impl InflightActivationStore for PostgresActivationStore { let mut conn = self.acquire_write_conn_metric("handle_expires_at").await?; let mut query_builder = QueryBuilder::new("DELETE FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push(" AND expires_at IS NOT NULL AND expires_at < "); query_builder.push_bind(now); self.add_partition_condition(&mut query_builder, false); @@ -847,11 +879,11 @@ impl InflightActivationStore for PostgresActivationStore { "UPDATE inflight_taskactivations SET status = ", ); - query_builder.push_bind(InflightActivationStatus::Pending.to_string()); + query_builder.push_bind(ActivationStatus::Pending.to_string()); query_builder.push(" WHERE delay_until IS NOT NULL AND delay_until < "); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(InflightActivationStatus::Delay.to_string()); + query_builder.push_bind(ActivationStatus::Delay.to_string()); self.add_partition_condition(&mut query_builder, false); let update_result = query_builder.build().execute(&mut *conn).await?; @@ -871,7 +903,7 @@ impl InflightActivationStore for PostgresActivationStore { let mut query_builder = QueryBuilder::new( "SELECT id, activation, on_attempts_exceeded FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(InflightActivationStatus::Failure.to_string()); + query_builder.push_bind(ActivationStatus::Failure.to_string()); self.add_partition_condition(&mut query_builder, false); let failed_tasks = query_builder .build_query_as::<(String, Vec, i32)>() @@ -902,7 +934,7 @@ impl InflightActivationStore for PostgresActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(InflightActivationStatus::Complete.to_string()) + .push_bind(ActivationStatus::Complete.to_string()) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -925,7 +957,7 @@ impl InflightActivationStore for PostgresActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(InflightActivationStatus::Complete.to_string()) + .push_bind(ActivationStatus::Complete.to_string()) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -946,7 +978,7 @@ impl InflightActivationStore for PostgresActivationStore { let mut conn = self.acquire_write_conn_metric("remove_completed").await?; let mut query_builder = QueryBuilder::new("DELETE FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(InflightActivationStatus::Complete.to_string()); + query_builder.push_bind(ActivationStatus::Complete.to_string()); self.add_partition_condition(&mut query_builder, false); let result = query_builder.build().execute(&mut *conn).await?; @@ -971,17 +1003,6 @@ impl InflightActivationStore for PostgresActivationStore { Ok(query.rows_affected()) } - - // Used in tests - async fn remove_db(&self) -> Result<(), Error> { - self.read_pool.close().await; - self.write_pool.close().await; - let default_pool = create_default_postgres_pool(&self.config.pg_connection).await?; - let _ = sqlx::query(format!("DROP DATABASE {}", &self.config.pg_database_name).as_str()) - .bind(&self.config.pg_database_name) - .execute(&default_pool) - .await; - let _ = default_pool.close().await; - Ok(()) - } } + +impl Store for PostgresStore {} diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index f26ddb5c..bcda5996 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -23,9 +23,11 @@ use sqlx::{ConnectOptions, FromRow, Pool, QueryBuilder, Row, Sqlite}; use std::{str::FromStr, time::Instant}; use crate::config::Config; -use crate::store::activation::{InflightActivation, InflightActivationStatus}; -use crate::store::traits::InflightActivationStore; -use crate::store::types::{BucketRange, FailedTasksForwarder}; +use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::traits::{ + ClaimStore, CountStore, IngestStore, PullStore, PushStore, Store, TestStore, UpkeepStore, +}; +use crate::store::types::{BucketRange, DepthCounts, FailedTasksForwarder}; #[derive(Debug, FromRow)] pub struct TableRow { @@ -51,10 +53,10 @@ pub struct TableRow { pub bucket: i16, } -impl TryFrom for TableRow { +impl TryFrom for TableRow { type Error = anyhow::Error; - fn try_from(value: InflightActivation) -> Result { + fn try_from(value: Activation) -> Result { Ok(Self { id: value.id, activation: value.activation, @@ -79,12 +81,12 @@ impl TryFrom for TableRow { } } -impl From for InflightActivation { +impl From for Activation { fn from(value: TableRow) -> Self { Self { id: value.id, activation: value.activation, - status: InflightActivationStatus::from_str(&value.status).unwrap(), + status: ActivationStatus::from_str(&value.status).unwrap(), partition: value.partition, offset: value.offset, added_at: value.added_at, @@ -332,45 +334,145 @@ impl SqliteActivationStore { } #[async_trait] -impl InflightActivationStore for SqliteActivationStore { - /// Trigger incremental vacuum to reclaim free pages in the database. - /// Depending on config data, will either vacuum a set number of - /// pages or attempt to reclaim all free pages. +impl ClaimStore for SqliteActivationStore { #[instrument(skip_all)] - async fn vacuum_db(&self) -> Result<(), Error> { - let timer = Instant::now(); + async fn claim_activations( + &self, + application: Option<&str>, + namespaces: Option<&[String]>, + limit: Option, + bucket: Option, + mark_processing: bool, + ) -> Result, Error> { + let now = Utc::now(); + let grace_period = self.config.processing_deadline_grace_sec; - if let Some(page_count) = self.config.vacuum_page_count { - let mut conn = self.acquire_write_conn_metric("vacuum_db").await?; - sqlx::query(format!("PRAGMA incremental_vacuum({page_count})").as_str()) - .execute(&mut *conn) - .await?; + let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations SET "); + + if mark_processing { + query_builder.push(format!( + "processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), claim_expires_at = NULL, status = " + )); + + query_builder.push_bind(ActivationStatus::Processing); } else { - let mut conn = self.acquire_write_conn_metric("vacuum_db").await?; - sqlx::query("PRAGMA incremental_vacuum") - .execute(&mut *conn) - .await?; + query_builder.push(format!( + "claim_expires_at = unixepoch('now', '+' || {:.3} || ' seconds', '+' || {grace_period} || ' seconds'), processing_deadline = NULL, status = ", + self.config.claim_lease_ms as f64 / 1000.0, + )); + + query_builder.push_bind(ActivationStatus::Claimed); } - let freelist_count: i32 = sqlx::query("PRAGMA freelist_count") - .fetch_one(&self.read_pool) - .await? - .get("freelist_count"); - metrics::histogram!("store.vacuum", "database" => "meta").record(timer.elapsed()); - metrics::gauge!("store.vacuum.freelist", "database" => "meta").set(freelist_count); - Ok(()) + query_builder.push(" WHERE id IN (SELECT id FROM inflight_taskactivations WHERE status = "); + query_builder.push_bind(ActivationStatus::Pending); + query_builder.push(" AND (expires_at IS NULL OR expires_at > "); + query_builder.push_bind(now.timestamp()); + query_builder.push(")"); + + if let Some(value) = application { + query_builder.push(" AND application ="); + query_builder.push_bind(value); + } + if let Some(namespaces) = namespaces + && !namespaces.is_empty() + { + query_builder.push(" AND namespace IN ("); + let mut separated = query_builder.separated(", "); + for namespace in namespaces.iter() { + separated.push_bind(namespace); + } + query_builder.push(")"); + } + if let Some((min, max)) = bucket { + query_builder.push(" AND bucket >= "); + query_builder.push_bind(min); + query_builder.push(" AND bucket <= "); + query_builder.push_bind(max); + } + query_builder.push(" ORDER BY added_at"); + if let Some(limit) = limit { + query_builder.push(" LIMIT "); + query_builder.push_bind(limit); + } + query_builder.push(") RETURNING *"); + + let mut conn = self.acquire_write_conn_metric("claim_activations").await?; + let rows: Vec = query_builder + .build_query_as::() + .fetch_all(&mut *conn) + .await?; + + Ok(rows.into_iter().map(|row| row.into()).collect()) } +} + +#[async_trait] +impl PushStore for SqliteActivationStore { + #[instrument(skip_all)] + async fn mark_processing(&self, id: &str) -> Result<(), Error> { + let mut conn = self + .acquire_write_conn_metric("mark_activation_processing") + .await?; + + let grace_period = self.config.processing_deadline_grace_sec; + let result = sqlx::query(&format!( + "UPDATE inflight_taskactivations SET + status = $1, + processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), + claim_expires_at = NULL + WHERE id = $2 AND status = $3", + )) + .bind(ActivationStatus::Processing) + .bind(id) + .bind(ActivationStatus::Claimed) + .execute(&mut *conn) + .await?; + + if result.rows_affected() == 0 { + metrics::counter!("push.mark_activation_processing", "result" => "not_found") + .increment(1); + + warn!( + task_id = %id, + "Activation could not be marked as sent, it may be missing or its status may have already changed" + ); + } else { + metrics::counter!("push.mark_activation_processing", "result" => "ok").increment(1); + } - /// Perform a full vacuum on the database. - async fn full_vacuum_db(&self) -> Result<(), Error> { - let mut conn = self.acquire_write_conn_metric("full_vacuum_db").await?; - sqlx::query("VACUUM").execute(&mut *conn).await?; - self.emit_db_status_metrics().await; Ok(()) } +} + +#[async_trait] +impl PullStore for SqliteActivationStore { + #[instrument(skip_all)] + async fn set_status( + &self, + id: &str, + status: ActivationStatus, + ) -> Result, Error> { + let mut conn = self.acquire_write_conn_metric("set_status").await?; + let result: Option = sqlx::query_as( + "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *", + ) + .bind(status) + .bind(id) + .fetch_optional(&mut *conn) + .await?; + + let Some(row) = result else { + return Ok(None); + }; + + Ok(Some(row.into())) + } +} - /// Get the size of the database in bytes based on SQLite metadata queries. - async fn db_size(&self) -> Result { +#[async_trait] +impl CountStore for SqliteActivationStore { + async fn size(&self) -> Result { let result: u64 = sqlx::query( "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()", ) @@ -381,8 +483,76 @@ impl InflightActivationStore for SqliteActivationStore { Ok(result) } - /// Get an activation by id. Primarily used for testing - async fn get_by_id(&self, id: &str) -> Result, Error> { + #[instrument(skip_all)] + async fn count_by_status(&self, status: ActivationStatus) -> Result { + let result = + sqlx::query("SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = $1") + .bind(status) + .fetch_one(&self.read_pool) + .await?; + Ok(result.get::("count") as usize) + } + + async fn count(&self) -> Result { + let result = sqlx::query("SELECT COUNT(*) as count FROM inflight_taskactivations") + .fetch_one(&self.read_pool) + .await?; + Ok(result.get::("count") as usize) + } + + #[instrument(skip_all)] + async fn count_depths(&self) -> Result { + let row: (i64, i64, i64, i64) = sqlx::query_as( + "SELECT IFNULL(SUM(CASE WHEN status = 'Pending' THEN 1 ELSE 0 END), 0), + IFNULL(SUM(CASE WHEN status = 'Delay' THEN 1 ELSE 0 END), 0), + IFNULL(SUM(CASE WHEN status = 'Claimed' THEN 1 ELSE 0 END), 0), + IFNULL(SUM(CASE WHEN status = 'Processing' THEN 1 ELSE 0 END), 0) + FROM inflight_taskactivations", + ) + .fetch_one(&self.read_pool) + .await?; + + Ok(DepthCounts { + pending: row.0 as usize, + delay: row.1 as usize, + claimed: row.2 as usize, + processing: row.3 as usize, + }) + } + + async fn pending_activation_max_lag(&self, now: &DateTime) -> f64 { + let result = sqlx::query( + "SELECT received_at, delay_until + FROM inflight_taskactivations + WHERE status = $1 + AND processing_attempts = 0 + ORDER BY received_at ASC + LIMIT 1 + ", + ) + .bind(ActivationStatus::Pending) + .fetch_one(&self.read_pool) + .await; + + if let Ok(row) = result { + let received_at: DateTime = row.get("received_at"); + let delay_until: Option> = row.get("delay_until"); + let millis = now.signed_duration_since(received_at).num_milliseconds() + - delay_until.map_or(0, |delay_time| { + delay_time + .signed_duration_since(received_at) + .num_milliseconds() + }); + millis as f64 / 1000.0 + } else { + 0.0 + } + } +} + +#[async_trait] +impl TestStore for SqliteActivationStore { + async fn get_by_id(&self, id: &str) -> Result, Error> { let row_result: Option = sqlx::query_as( " SELECT id, @@ -419,13 +589,51 @@ impl InflightActivationStore for SqliteActivationStore { Ok(Some(row.into())) } - fn assign_partitions(&self, partitions: Vec) -> Result<(), Error> { + #[instrument(skip_all)] + async fn set_processing_deadline( + &self, + id: &str, + deadline: Option>, + ) -> Result<(), Error> { + let mut conn = self + .acquire_write_conn_metric("set_processing_deadline") + .await?; + sqlx::query("UPDATE inflight_taskactivations SET processing_deadline = $1 WHERE id = $2") + .bind(deadline.unwrap().timestamp()) + .bind(id) + .execute(&mut *conn) + .await?; + Ok(()) + } + + #[instrument(skip_all)] + async fn delete_activation(&self, id: &str) -> Result<(), Error> { + let mut conn = self.acquire_write_conn_metric("delete_activation").await?; + sqlx::query("DELETE FROM inflight_taskactivations WHERE id = $1") + .bind(id) + .execute(&mut *conn) + .await?; + Ok(()) + } + + async fn clear(&self) -> Result<(), Error> { + let mut conn = self.acquire_write_conn_metric("clear").await?; + sqlx::query("DELETE FROM inflight_taskactivations") + .execute(&mut *conn) + .await?; + Ok(()) + } +} + +#[async_trait] +impl IngestStore for SqliteActivationStore { + async fn assign_partitions(&self, partitions: Vec) -> Result<(), Error> { warn!("assign_partitions: {:?}", partitions); Ok(()) } #[instrument(skip_all)] - async fn store(&self, batch: Vec) -> Result { + async fn write(&self, batch: Vec) -> Result { if batch.is_empty() { return Ok(0); } @@ -521,217 +729,44 @@ impl InflightActivationStore for SqliteActivationStore { rows_affected } +} +#[async_trait] +impl UpkeepStore for SqliteActivationStore { #[instrument(skip_all)] - async fn claim_activations( - &self, - application: Option<&str>, - namespaces: Option<&[String]>, - limit: Option, - bucket: Option, - mark_processing: bool, - ) -> Result, Error> { - let now = Utc::now(); - let grace_period = self.config.processing_deadline_grace_sec; - - let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations SET "); - - if mark_processing { - query_builder.push(format!( - "processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), claim_expires_at = NULL, status = " - )); - - query_builder.push_bind(InflightActivationStatus::Processing); - } else { - query_builder.push(format!( - "claim_expires_at = unixepoch('now', '+' || {:.3} || ' seconds', '+' || {grace_period} || ' seconds'), processing_deadline = NULL, status = ", - self.config.claim_lease_ms as f64 / 1000.0, - )); - - query_builder.push_bind(InflightActivationStatus::Claimed); - } - - query_builder.push(" WHERE id IN (SELECT id FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(InflightActivationStatus::Pending); - query_builder.push(" AND (expires_at IS NULL OR expires_at > "); - query_builder.push_bind(now.timestamp()); - query_builder.push(")"); - - if let Some(value) = application { - query_builder.push(" AND application ="); - query_builder.push_bind(value); - } - if let Some(namespaces) = namespaces - && !namespaces.is_empty() - { - query_builder.push(" AND namespace IN ("); - let mut separated = query_builder.separated(", "); - for namespace in namespaces.iter() { - separated.push_bind(namespace); - } - query_builder.push(")"); - } - if let Some((min, max)) = bucket { - query_builder.push(" AND bucket >= "); - query_builder.push_bind(min); - query_builder.push(" AND bucket <= "); - query_builder.push_bind(max); - } - query_builder.push(" ORDER BY added_at"); - if let Some(limit) = limit { - query_builder.push(" LIMIT "); - query_builder.push_bind(limit); - } - query_builder.push(") RETURNING *"); - - let mut conn = self.acquire_write_conn_metric("claim_activations").await?; - let rows: Vec = query_builder - .build_query_as::() - .fetch_all(&mut *conn) - .await?; - - Ok(rows.into_iter().map(|row| row.into()).collect()) - } - - #[instrument(skip_all)] - async fn mark_activation_processing(&self, id: &str) -> Result<(), Error> { - let mut conn = self - .acquire_write_conn_metric("mark_activation_processing") - .await?; - - let grace_period = self.config.processing_deadline_grace_sec; - let result = sqlx::query(&format!( - "UPDATE inflight_taskactivations SET - status = $1, - processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), - claim_expires_at = NULL - WHERE id = $2 AND status = $3", - )) - .bind(InflightActivationStatus::Processing) - .bind(id) - .bind(InflightActivationStatus::Claimed) - .execute(&mut *conn) - .await?; - - if result.rows_affected() == 0 { - metrics::counter!("push.mark_activation_processing", "result" => "not_found") - .increment(1); - - warn!( - task_id = %id, - "Activation could not be marked as sent, it may be missing or its status may have already changed" - ); - } else { - metrics::counter!("push.mark_activation_processing", "result" => "ok").increment(1); - } - - Ok(()) - } - - /// Get the age of the oldest pending activation in seconds. - /// Only activations with status=pending and processing_attempts=0 are considered - /// as we are interested in latency to the *first* attempt. - /// Tasks with delay_until set, will have their age adjusted based on their - /// delay time. No tasks = 0 lag - async fn pending_activation_max_lag(&self, now: &DateTime) -> f64 { - let result = sqlx::query( - "SELECT received_at, delay_until - FROM inflight_taskactivations - WHERE status = $1 - AND processing_attempts = 0 - ORDER BY received_at ASC - LIMIT 1 - ", - ) - .bind(InflightActivationStatus::Pending) - .fetch_one(&self.read_pool) - .await; + async fn vacuum_db(&self) -> Result<(), Error> { + let timer = Instant::now(); - if let Ok(row) = result { - let received_at: DateTime = row.get("received_at"); - let delay_until: Option> = row.get("delay_until"); - let millis = now.signed_duration_since(received_at).num_milliseconds() - - delay_until.map_or(0, |delay_time| { - delay_time - .signed_duration_since(received_at) - .num_milliseconds() - }); - millis as f64 / 1000.0 + if let Some(page_count) = self.config.vacuum_page_count { + let mut conn = self.acquire_write_conn_metric("vacuum_db").await?; + sqlx::query(format!("PRAGMA incremental_vacuum({page_count})").as_str()) + .execute(&mut *conn) + .await?; } else { - // If we couldn't find a row, there is no latency. - 0.0 - } - } - - #[instrument(skip_all)] - async fn count_by_status(&self, status: InflightActivationStatus) -> Result { - let result = - sqlx::query("SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = $1") - .bind(status) - .fetch_one(&self.read_pool) + let mut conn = self.acquire_write_conn_metric("vacuum_db").await?; + sqlx::query("PRAGMA incremental_vacuum") + .execute(&mut *conn) .await?; - Ok(result.get::("count") as usize) - } - - async fn count(&self) -> Result { - let result = sqlx::query("SELECT COUNT(*) as count FROM inflight_taskactivations") + } + let freelist_count: i32 = sqlx::query("PRAGMA freelist_count") .fetch_one(&self.read_pool) - .await?; - Ok(result.get::("count") as usize) - } - - /// Update the status of a specific activation - #[instrument(skip_all)] - async fn set_status( - &self, - id: &str, - status: InflightActivationStatus, - ) -> Result, Error> { - let mut conn = self.acquire_write_conn_metric("set_status").await?; - let result: Option = sqlx::query_as( - "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *", - ) - .bind(status) - .bind(id) - .fetch_optional(&mut *conn) - .await?; - - let Some(row) = result else { - return Ok(None); - }; - - Ok(Some(row.into())) - } + .await? + .get("freelist_count"); - #[instrument(skip_all)] - async fn set_processing_deadline( - &self, - id: &str, - deadline: Option>, - ) -> Result<(), Error> { - let mut conn = self - .acquire_write_conn_metric("set_processing_deadline") - .await?; - sqlx::query("UPDATE inflight_taskactivations SET processing_deadline = $1 WHERE id = $2") - .bind(deadline.unwrap().timestamp()) - .bind(id) - .execute(&mut *conn) - .await?; + metrics::histogram!("store.vacuum", "database" => "meta").record(timer.elapsed()); + metrics::gauge!("store.vacuum.freelist", "database" => "meta").set(freelist_count); Ok(()) } - #[instrument(skip_all)] - async fn delete_activation(&self, id: &str) -> Result<(), Error> { - let mut conn = self.acquire_write_conn_metric("delete_activation").await?; - sqlx::query("DELETE FROM inflight_taskactivations WHERE id = $1") - .bind(id) - .execute(&mut *conn) - .await?; + async fn full_vacuum_db(&self) -> Result<(), Error> { + let mut conn = self.acquire_write_conn_metric("full_vacuum_db").await?; + sqlx::query("VACUUM").execute(&mut *conn).await?; + self.emit_db_status_metrics().await; Ok(()) } #[instrument(skip_all)] - async fn get_retry_activations(&self) -> Result, Error> { + async fn get_retry_activations(&self) -> Result, Error> { Ok(sqlx::query_as( " SELECT id, @@ -757,7 +792,7 @@ impl InflightActivationStore for SqliteActivationStore { WHERE status = $1 ", ) - .bind(InflightActivationStatus::Retry) + .bind(ActivationStatus::Retry) .fetch_all(&self.read_pool) .await? .into_iter() @@ -765,14 +800,6 @@ impl InflightActivationStore for SqliteActivationStore { .collect()) } - async fn clear(&self) -> Result<(), Error> { - let mut conn = self.acquire_write_conn_metric("clear").await?; - sqlx::query("DELETE FROM inflight_taskactivations") - .execute(&mut *conn) - .await?; - Ok(()) - } - /// Expired push claims (`Claimed` + past `claim_expires_at`). #[instrument(skip_all)] async fn handle_claim_expiration(&self) -> Result { @@ -789,9 +816,9 @@ impl InflightActivationStore for SqliteActivationStore { AND claim_expires_at < $2 AND status = $3", ) - .bind(InflightActivationStatus::Pending) + .bind(ActivationStatus::Pending) .bind(now.timestamp()) - .bind(InflightActivationStatus::Claimed) + .bind(ActivationStatus::Claimed) .execute(&mut *conn) .await?; @@ -813,9 +840,9 @@ impl InflightActivationStore for SqliteActivationStore { SET processing_deadline = null, status = $1 WHERE processing_deadline < $2 AND at_most_once = TRUE AND status = $3", ) - .bind(InflightActivationStatus::Failure) + .bind(ActivationStatus::Failure) .bind(now.timestamp()) - .bind(InflightActivationStatus::Processing) + .bind(ActivationStatus::Processing) .execute(&mut *atomic) .await; @@ -831,9 +858,9 @@ impl InflightActivationStore for SqliteActivationStore { SET processing_deadline = null, status = $1, processing_attempts = processing_attempts + 1 WHERE processing_deadline < $2 AND status = $3", ) - .bind(InflightActivationStatus::Pending) + .bind(ActivationStatus::Pending) .bind(now.timestamp()) - .bind(InflightActivationStatus::Processing) + .bind(ActivationStatus::Processing) .execute(&mut *atomic) .await; @@ -859,9 +886,9 @@ impl InflightActivationStore for SqliteActivationStore { SET status = $1 WHERE processing_attempts >= $2 AND status = $3", ) - .bind(InflightActivationStatus::Failure) + .bind(ActivationStatus::Failure) .bind(self.config.max_processing_attempts as i32) - .bind(InflightActivationStatus::Pending) + .bind(ActivationStatus::Pending) .execute(&mut *conn) .await; @@ -885,7 +912,7 @@ impl InflightActivationStore for SqliteActivationStore { let query = sqlx::query( "DELETE FROM inflight_taskactivations WHERE status = $1 AND expires_at IS NOT NULL AND expires_at < $2", ) - .bind(InflightActivationStatus::Pending) + .bind(ActivationStatus::Pending) .bind(now.timestamp()) .execute(&mut *conn) .await?; @@ -909,9 +936,9 @@ impl InflightActivationStore for SqliteActivationStore { WHERE delay_until IS NOT NULL AND delay_until < $2 AND status = $3 "#, ) - .bind(InflightActivationStatus::Pending) + .bind(ActivationStatus::Pending) .bind(now.timestamp()) - .bind(InflightActivationStatus::Delay) + .bind(ActivationStatus::Delay) .execute(&mut *conn) .await?; @@ -930,7 +957,7 @@ impl InflightActivationStore for SqliteActivationStore { let failed_tasks: Vec = sqlx::query("SELECT id, activation, on_attempts_exceeded FROM inflight_taskactivations WHERE status = $1") - .bind(InflightActivationStatus::Failure) + .bind(ActivationStatus::Failure) .fetch_all(&mut *atomic) .await? .into_iter() @@ -962,7 +989,7 @@ impl InflightActivationStore for SqliteActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(InflightActivationStatus::Complete) + .push_bind(ActivationStatus::Complete) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -985,7 +1012,7 @@ impl InflightActivationStore for SqliteActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(InflightActivationStatus::Complete) + .push_bind(ActivationStatus::Complete) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -1005,7 +1032,7 @@ impl InflightActivationStore for SqliteActivationStore { async fn remove_completed(&self) -> Result { let mut conn = self.acquire_write_conn_metric("remove_completed").await?; let query = sqlx::query("DELETE FROM inflight_taskactivations WHERE status = $1") - .bind(InflightActivationStatus::Complete) + .bind(ActivationStatus::Complete) .execute(&mut *conn) .await?; @@ -1030,3 +1057,5 @@ impl InflightActivationStore for SqliteActivationStore { Ok(query.rows_affected()) } } + +impl Store for SqliteActivationStore {} diff --git a/src/store/tests.rs b/src/store/tests.rs index 061bdee1..9aecf125 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -6,14 +6,15 @@ use sqlx::{QueryBuilder, Sqlite}; use std::{collections::HashSet, fs, io::Error, path::Path, sync::Arc, time::Duration}; use tokio::{sync::broadcast, task::JoinSet}; +use crate::store::traits::{IngestStore, UpkeepStore}; use crate::{ config::Config, - store::activation::{InflightActivationBuilder, InflightActivationStatus}, + store::activation::{ActivationBuilder, ActivationStatus}, store::adapters::{ postgres::PostgresActivationStoreConfig, sqlite::{InflightActivationStoreConfig, SqliteActivationStore, create_sqlite_pool}, }, - store::traits::InflightActivationStore, + store::traits::ClaimStore, test_utils::{ StatusCount, TaskActivationBuilder, assert_counts, create_integration_config, create_integration_config_with_ssl, create_test_store, generate_temp_filename, @@ -24,41 +25,41 @@ use crate::{ #[test] fn test_inflightactivation_status_is_completion() { - let mut value = InflightActivationStatus::Unspecified; + let mut value = ActivationStatus::Unspecified; assert!(!value.is_conclusion()); - value = InflightActivationStatus::Pending; + value = ActivationStatus::Pending; assert!(!value.is_conclusion()); - value = InflightActivationStatus::Processing; + value = ActivationStatus::Processing; assert!(!value.is_conclusion()); - value = InflightActivationStatus::Retry; + value = ActivationStatus::Retry; assert!(value.is_conclusion()); - value = InflightActivationStatus::Failure; + value = ActivationStatus::Failure; assert!(value.is_conclusion()); - value = InflightActivationStatus::Complete; + value = ActivationStatus::Complete; assert!(value.is_conclusion()); } #[test] fn test_inflightactivation_status_from() { - let mut value: InflightActivationStatus = TaskActivationStatus::Pending.into(); - assert_eq!(value, InflightActivationStatus::Pending); + let mut value: ActivationStatus = TaskActivationStatus::Pending.into(); + assert_eq!(value, ActivationStatus::Pending); value = TaskActivationStatus::Processing.into(); - assert_eq!(value, InflightActivationStatus::Processing); + assert_eq!(value, ActivationStatus::Processing); value = TaskActivationStatus::Retry.into(); - assert_eq!(value, InflightActivationStatus::Retry); + assert_eq!(value, ActivationStatus::Retry); value = TaskActivationStatus::Failure.into(); - assert_eq!(value, InflightActivationStatus::Failure); + assert_eq!(value, ActivationStatus::Failure); value = TaskActivationStatus::Complete.into(); - assert_eq!(value, InflightActivationStatus::Complete); + assert_eq!(value, ActivationStatus::Complete); } #[tokio::test] @@ -89,7 +90,7 @@ async fn test_store(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let result = store.count().await; assert_eq!(result.unwrap(), 2); @@ -105,15 +106,15 @@ async fn test_count_depths(#[case] adapter: &str) { // Check counts for an empty store let pending = store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(); let delay = store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); let processing = store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(); @@ -125,31 +126,31 @@ async fn test_count_depths(#[case] adapter: &str) { // Check counts for a store with four activations with varying statuses let batch = make_activations(4); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); store - .set_status("id_0", InflightActivationStatus::Processing) + .set_status("id_0", ActivationStatus::Processing) .await .unwrap(); store - .set_status("id_1", InflightActivationStatus::Delay) + .set_status("id_1", ActivationStatus::Delay) .await .unwrap(); store - .set_status("id_2", InflightActivationStatus::Complete) + .set_status("id_2", ActivationStatus::Complete) .await .unwrap(); let pending = store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(); let delay = store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(); let processing = store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(); @@ -177,7 +178,7 @@ async fn test_store_duplicate_id_in_batch(#[case] adapter: &str) { batch[0].id = "id_0".into(); batch[1].id = "id_0".into(); - let first_result = store.store(batch).await; + let first_result = store.write(batch).await; assert!( first_result.is_ok(), "{}", @@ -197,7 +198,7 @@ async fn test_store_duplicate_id_between_batches(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let first_count = store.count().await; assert_eq!(first_count.unwrap(), 2); @@ -205,7 +206,7 @@ async fn test_store_duplicate_id_between_batches(#[case] adapter: &str) { // Old batch and new should have conflicts assert_eq!(batch[0].id, new_batch[0].id); assert_eq!(batch[1].id, new_batch[1].id); - assert!(store.store(new_batch).await.is_ok()); + assert!(store.write(new_batch).await.is_ok()); let second_count = store.count().await; assert_eq!(second_count.unwrap(), 2); @@ -220,7 +221,7 @@ async fn test_get_pending_activation(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let result = store .claim_activation_for_pull(None, None) @@ -229,7 +230,7 @@ async fn test_get_pending_activation(#[case] adapter: &str) { .expect("expected one activation"); assert_eq!(result.id, "id_0"); - assert_eq!(result.status, InflightActivationStatus::Processing); + assert_eq!(result.status, ActivationStatus::Processing); assert_eq!(result.processing_deadline_duration, 10); assert!( result.processing_deadline.unwrap().timestamp() >= Utc::now().timestamp() + 13, @@ -257,7 +258,7 @@ async fn test_get_pending_activation_bucket_filter(#[case] adapter: &str) { let mut batch = make_activations(2); batch[0].bucket = 10; batch[1].bucket = 20; - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let mut first = store .claim_activations_for_push(None, None, Some(1), Some((15, 25))) @@ -301,7 +302,7 @@ async fn test_get_pending_activation_with_race(#[case] adapter: &str) { for chunk in make_activations_with_namespace(namespace.clone(), NUM_CONCURRENT_WRITES).chunks(1024) { - store.store(chunk.to_vec()).await.unwrap(); + store.write(chunk.to_vec()).await.unwrap(); } let (tx, _) = broadcast::channel::<()>(1); @@ -347,7 +348,7 @@ async fn test_get_pending_activation_with_namespace(#[case] adapter: &str) { let mut batch = make_activations(2); batch[1].namespace = "other_namespace".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let other_namespace = "other_namespace".to_string(); // Get activation from other namespace @@ -357,7 +358,7 @@ async fn test_get_pending_activation_with_namespace(#[case] adapter: &str) { .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, InflightActivationStatus::Processing); + assert_eq!(result.status, ActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.namespace, "other_namespace"); store.remove_db().await.unwrap(); @@ -375,7 +376,7 @@ async fn test_get_pending_activation_from_multiple_namespaces(#[case] adapter: & batch[1].namespace = "ns2".into(); batch[2].namespace = "ns3".into(); batch[3].namespace = "ns4".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); // Get activation from multiple namespaces (should get oldest). // Use `claim_activations` so upkeep-style `None` application + namespaces is allowed (not `claim_activations_for_push`). @@ -388,10 +389,10 @@ async fn test_get_pending_activation_from_multiple_namespaces(#[case] adapter: & assert_eq!(result.len(), 2); assert_eq!(result[1].id, "id_2"); assert_eq!(result[1].namespace, "ns3"); - assert_eq!(result[1].status, InflightActivationStatus::Claimed); + assert_eq!(result[1].status, ActivationStatus::Claimed); assert_eq!(result[0].id, "id_1"); assert_eq!(result[0].namespace, "ns2"); - assert_eq!(result[0].status, InflightActivationStatus::Claimed); + assert_eq!(result[0].status, ActivationStatus::Claimed); store.remove_db().await.unwrap(); } @@ -404,7 +405,7 @@ async fn test_get_pending_activation_with_namespace_requires_application(#[case] let mut batch = make_activations(2); batch[1].namespace = "other_namespace".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); // This is an invalid query as we don't want to allow clients to fetch tasks from any application let other_namespace = "other_namespace".to_string(); @@ -446,7 +447,7 @@ async fn test_get_pending_activation_skip_expires(#[case] adapter: &str) { let mut batch = make_activations(1); batch[0].expires_at = Some(Utc::now() - Duration::from_secs(100)); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let result = store.claim_activation_for_pull(None, None).await; assert!(result.is_ok()); @@ -473,7 +474,7 @@ async fn test_get_pending_activation_earliest(#[case] adapter: &str) { let mut batch = make_activations(2); batch[0].added_at = Utc.with_ymd_and_hms(2024, 6, 24, 0, 0, 0).unwrap(); batch[1].added_at = Utc.with_ymd_and_hms(1998, 6, 24, 0, 0, 0).unwrap(); - let ret = store.store(batch.clone()).await; + let ret = store.write(batch.clone()).await; assert!(ret.is_ok(), "{}", ret.err().unwrap().to_string()); let result = store @@ -497,7 +498,7 @@ async fn test_get_pending_activation_fetches_application(#[case] adapter: &str) let mut batch = make_activations(1); batch[0].application = "hammers".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); // Getting an activation with no application filter should // include activations with application set. @@ -507,7 +508,7 @@ async fn test_get_pending_activation_fetches_application(#[case] adapter: &str) .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_0"); - assert_eq!(result.status, InflightActivationStatus::Processing); + assert_eq!(result.status, ActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); store.remove_db().await.unwrap(); @@ -522,7 +523,7 @@ async fn test_get_pending_activation_with_application(#[case] adapter: &str) { let mut batch = make_activations(2); batch[1].application = "hammers".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); // Get activation from a named application let result = store @@ -531,7 +532,7 @@ async fn test_get_pending_activation_with_application(#[case] adapter: &str) { .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, InflightActivationStatus::Processing); + assert_eq!(result.status, ActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); @@ -564,7 +565,7 @@ async fn test_get_pending_activation_with_application_and_namespace(#[case] adap batch[2].application = "hammers".into(); batch[2].namespace = "not-target".into(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let target_ns = "target".to_string(); // Get activation from a named application @@ -574,7 +575,7 @@ async fn test_get_pending_activation_with_application_and_namespace(#[case] adap .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, InflightActivationStatus::Processing); + assert_eq!(result.status, ActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); assert_eq!(result.namespace, "target"); @@ -599,18 +600,21 @@ async fn test_get_pending_activations_no_limit(#[case] adapter: &str) { const N: usize = 4; let batch = make_activations(N as u32); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let got = store .claim_activations_for_push(None, None, None, None) .await .unwrap(); assert_eq!(got.len(), N); - assert!( - got.iter() - .all(|a| a.status == InflightActivationStatus::Claimed) + assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); + assert_eq!( + store + .count_by_status(ActivationStatus::Pending) + .await + .unwrap(), + 0 ); - assert_eq!(store.count_pending_activations().await.unwrap(), 0); assert_counts( StatusCount { pending: 0, @@ -633,19 +637,19 @@ async fn test_get_pending_activations_limit_below_pending(#[case] adapter: &str) const X: i32 = 3; let batch = make_activations(N as u32); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let got = store .claim_activations_for_push(None, None, Some(X), None) .await .unwrap(); assert_eq!(got.len(), X as usize); - assert!( - got.iter() - .all(|a| a.status == InflightActivationStatus::Claimed) - ); + assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); assert_eq!( - store.count_pending_activations().await.unwrap(), + store + .count_by_status(ActivationStatus::Pending) + .await + .unwrap(), N - X as usize ); assert_counts( @@ -670,18 +674,21 @@ async fn test_get_pending_activations_limit_above_pending(#[case] adapter: &str) const X: i32 = 10; let batch = make_activations(Y as u32); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let got = store .claim_activations_for_push(None, None, Some(X), None) .await .unwrap(); assert_eq!(got.len(), Y); - assert!( - got.iter() - .all(|a| a.status == InflightActivationStatus::Claimed) + assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); + assert_eq!( + store + .count_by_status(ActivationStatus::Pending) + .await + .unwrap(), + 0 ); - assert_eq!(store.count_pending_activations().await.unwrap(), 0); assert_counts( StatusCount { pending: 0, @@ -702,10 +709,16 @@ async fn test_count_pending_activations(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(3); - batch[0].status = InflightActivationStatus::Processing; - assert!(store.store(batch).await.is_ok()); + batch[0].status = ActivationStatus::Processing; + assert!(store.write(batch).await.is_ok()); - assert_eq!(store.count_pending_activations().await.unwrap(), 2); + assert_eq!( + store + .count_by_status(ActivationStatus::Pending) + .await + .unwrap(), + 2 + ); assert_counts( StatusCount { pending: 2, @@ -726,7 +739,7 @@ async fn test_set_activation_status(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { pending: 2, @@ -738,7 +751,7 @@ async fn test_set_activation_status(#[case] adapter: &str) { assert!( store - .set_status("id_0", InflightActivationStatus::Failure) + .set_status("id_0", ActivationStatus::Failure) .await .is_ok() ); @@ -754,7 +767,7 @@ async fn test_set_activation_status(#[case] adapter: &str) { assert!( store - .set_status("id_0", InflightActivationStatus::Pending) + .set_status("id_0", ActivationStatus::Pending) .await .is_ok() ); @@ -768,13 +781,13 @@ async fn test_set_activation_status(#[case] adapter: &str) { .await; assert!( store - .set_status("id_0", InflightActivationStatus::Failure) + .set_status("id_0", ActivationStatus::Failure) .await .is_ok() ); assert!( store - .set_status("id_1", InflightActivationStatus::Failure) + .set_status("id_1", ActivationStatus::Failure) .await .is_ok() ); @@ -796,23 +809,21 @@ async fn test_set_activation_status(#[case] adapter: &str) { ); let result = store - .set_status("not_there", InflightActivationStatus::Complete) + .set_status("not_there", ActivationStatus::Complete) .await; assert!(result.is_ok(), "no query error"); let activation = result.unwrap(); assert!(activation.is_none(), "no activation found"); - let result = store - .set_status("id_0", InflightActivationStatus::Complete) - .await; + let result = store.set_status("id_0", ActivationStatus::Complete).await; assert!(result.is_ok(), "no query error"); let result_opt = result.unwrap(); assert!(result_opt.is_some(), "activation should be returned"); let inflight = result_opt.unwrap(); assert_eq!(inflight.id, "id_0"); - assert_eq!(inflight.status, InflightActivationStatus::Complete); + assert_eq!(inflight.status, ActivationStatus::Complete); store.remove_db().await.unwrap(); } @@ -824,7 +835,7 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { let mut batch = make_activations(2); batch[1].partition = 1; - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { pending: 1, @@ -836,7 +847,7 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { assert!( store - .set_status("id_0", InflightActivationStatus::Failure) + .set_status("id_0", ActivationStatus::Failure) .await .is_ok() ); @@ -851,7 +862,7 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { assert!( store - .set_status("id_0", InflightActivationStatus::Pending) + .set_status("id_0", ActivationStatus::Pending) .await .is_ok() ); @@ -865,13 +876,13 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { .await; assert!( store - .set_status("id_0", InflightActivationStatus::Failure) + .set_status("id_0", ActivationStatus::Failure) .await .is_ok() ); assert!( store - .set_status("id_1", InflightActivationStatus::Failure) + .set_status("id_1", ActivationStatus::Failure) .await .is_ok() ); @@ -895,23 +906,21 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { ); let result = store - .set_status("not_there", InflightActivationStatus::Complete) + .set_status("not_there", ActivationStatus::Complete) .await; assert!(result.is_ok(), "no query error"); let activation = result.unwrap(); assert!(activation.is_none(), "no activation found"); - let result = store - .set_status("id_0", InflightActivationStatus::Complete) - .await; + let result = store.set_status("id_0", ActivationStatus::Complete).await; assert!(result.is_ok(), "no query error"); let result_opt = result.unwrap(); assert!(result_opt.is_some(), "activation should be returned"); let inflight = result_opt.unwrap(); assert_eq!(inflight.id, "id_0"); - assert_eq!(inflight.status, InflightActivationStatus::Complete); + assert_eq!(inflight.status, ActivationStatus::Complete); store.remove_db().await.unwrap(); } @@ -923,7 +932,7 @@ async fn test_set_processing_deadline(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(1); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let deadline = Utc::now().round_subsecs(0); let result = store.set_processing_deadline("id_0", Some(deadline)).await; @@ -945,7 +954,7 @@ async fn test_delete_activation(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let result = store.count().await; assert_eq!(result.unwrap(), 2); @@ -972,7 +981,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { pending: 2, @@ -984,7 +993,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { assert!( store - .set_status("id_0", InflightActivationStatus::Retry) + .set_status("id_0", ActivationStatus::Retry) .await .is_ok() ); @@ -1000,7 +1009,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { assert!( store - .set_status("id_1", InflightActivationStatus::Retry) + .set_status("id_1", ActivationStatus::Retry) .await .is_ok() ); @@ -1008,7 +1017,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { let retries = store.get_retry_activations().await.unwrap(); assert_eq!(retries.len(), 2); for record in retries.iter() { - assert_eq!(record.status, InflightActivationStatus::Retry); + assert_eq!(record.status, ActivationStatus::Retry); } store.remove_db().await.unwrap(); } @@ -1021,10 +1030,10 @@ async fn test_handle_processing_deadline(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { pending: 1, @@ -1065,11 +1074,11 @@ async fn test_handle_processing_deadline_multiple_tasks(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[0].status = InflightActivationStatus::Processing; + batch[0].status = ActivationStatus::Processing; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); - batch[1].status = InflightActivationStatus::Claimed; + batch[1].status = ActivationStatus::Claimed; batch[1].processing_deadline = Some(Utc::now() + chrono::Duration::days(30)); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { processing: 1, @@ -1104,10 +1113,10 @@ async fn test_handle_processing_at_most_once(#[case] adapter: &str) { // Both records are past processing deadlines let mut batch = make_activations(2); - batch[0].status = InflightActivationStatus::Processing; + batch[0].status = ActivationStatus::Processing; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; replace_retry_state( &mut batch[1], @@ -1122,7 +1131,7 @@ async fn test_handle_processing_at_most_once(#[case] adapter: &str) { batch[1].at_most_once = true; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { processing: 2, @@ -1146,7 +1155,7 @@ async fn test_handle_processing_at_most_once(#[case] adapter: &str) { .await; let task = store.get_by_id(&batch[1].id).await.unwrap().unwrap(); - assert_eq!(task.status, InflightActivationStatus::Failure); + assert_eq!(task.status, ActivationStatus::Failure); store.remove_db().await.unwrap(); } @@ -1158,7 +1167,7 @@ async fn test_handle_processing_deadline_discard_after(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1171,7 +1180,7 @@ async fn test_handle_processing_deadline_discard_after(#[case] adapter: &str) { }), ); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { pending: 1, @@ -1204,7 +1213,7 @@ async fn test_handle_processing_deadline_deadletter_after(#[case] adapter: &str) let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1217,7 +1226,7 @@ async fn test_handle_processing_deadline_deadletter_after(#[case] adapter: &str) }), ); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { pending: 1, @@ -1250,7 +1259,7 @@ async fn test_handle_processing_deadline_no_retries_remaining(#[case] adapter: & let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1262,7 +1271,7 @@ async fn test_handle_processing_deadline_no_retries_remaining(#[case] adapter: & delay_on_retry: None, }), ); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { processing: 1, @@ -1295,13 +1304,13 @@ async fn test_handle_processing_deadline_no_retries_remaining(#[case] adapter: & async fn test_handle_claim_expiration_unsent_no_attempt_increment(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(1); - batch[0].status = InflightActivationStatus::Claimed; + batch[0].status = ActivationStatus::Claimed; batch[0].claim_expires_at = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let count = store.handle_claim_expiration().await.unwrap(); assert_eq!(count, 1); let task = store.get_by_id(&batch[0].id).await.unwrap().unwrap(); - assert_eq!(task.status, InflightActivationStatus::Pending); + assert_eq!(task.status, ActivationStatus::Pending); assert_eq!(task.processing_attempts, 0); store.remove_db().await.unwrap(); } @@ -1313,14 +1322,14 @@ async fn test_handle_claim_expiration_unsent_no_attempt_increment(#[case] adapte async fn test_handle_claim_expiration_at_most_once_reverts_to_pending(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(1); - batch[0].status = InflightActivationStatus::Claimed; + batch[0].status = ActivationStatus::Claimed; batch[0].at_most_once = true; batch[0].claim_expires_at = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); let count = store.handle_claim_expiration().await.unwrap(); assert_eq!(count, 1); let task = store.get_by_id(&batch[0].id).await.unwrap().unwrap(); - assert_eq!(task.status, InflightActivationStatus::Pending); + assert_eq!(task.status, ActivationStatus::Pending); assert_eq!(task.processing_attempts, 0); store.remove_db().await.unwrap(); } @@ -1334,18 +1343,18 @@ async fn test_processing_attempts_exceeded(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(3); - batch[0].status = InflightActivationStatus::Pending; + batch[0].status = ActivationStatus::Pending; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[0].processing_attempts = config.max_processing_attempts as i32; - batch[1].status = InflightActivationStatus::Complete; + batch[1].status = ActivationStatus::Complete; batch[1].added_at += Duration::from_secs(1); - batch[2].status = InflightActivationStatus::Pending; + batch[2].status = ActivationStatus::Pending; batch[2].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[2].processing_attempts = config.max_processing_attempts as i32; - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { complete: 1, @@ -1379,13 +1388,13 @@ async fn test_remove_completed(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut records = make_activations(3); - records[0].status = InflightActivationStatus::Complete; - records[1].status = InflightActivationStatus::Pending; + records[0].status = ActivationStatus::Complete; + records[1].status = ActivationStatus::Pending; records[1].added_at += Duration::from_secs(1); - records[2].status = InflightActivationStatus::Complete; + records[2].status = ActivationStatus::Complete; records[2].added_at += Duration::from_secs(2); - assert!(store.store(records.clone()).await.is_ok()); + assert!(store.write(records.clone()).await.is_ok()); assert_counts( StatusCount { complete: 2, @@ -1441,17 +1450,17 @@ async fn test_remove_completed_multiple_gaps(#[case] adapter: &str) { let mut records = make_activations(4); // only record 1 can be removed - records[0].status = InflightActivationStatus::Complete; - records[1].status = InflightActivationStatus::Failure; + records[0].status = ActivationStatus::Complete; + records[1].status = ActivationStatus::Failure; records[1].added_at += Duration::from_secs(1); - records[2].status = InflightActivationStatus::Complete; + records[2].status = ActivationStatus::Complete; records[2].added_at += Duration::from_secs(2); - records[3].status = InflightActivationStatus::Processing; + records[3].status = ActivationStatus::Processing; records[3].added_at += Duration::from_secs(3); - assert!(store.store(records.clone()).await.is_ok()); + assert!(store.write(records.clone()).await.is_ok()); assert_counts( StatusCount { complete: 2, @@ -1516,7 +1525,7 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { let mut records = make_activations(4); // deadletter - records[0].status = InflightActivationStatus::Failure; + records[0].status = ActivationStatus::Failure; replace_retry_state( &mut records[0], Some(RetryState { @@ -1528,7 +1537,7 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { }), ); // discard - records[1].status = InflightActivationStatus::Failure; + records[1].status = ActivationStatus::Failure; replace_retry_state( &mut records[1], Some(RetryState { @@ -1540,11 +1549,11 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { }), ); // no retry state = discard - records[2].status = InflightActivationStatus::Failure; + records[2].status = ActivationStatus::Failure; replace_retry_state(&mut records[2], None); // Another deadletter - records[3].status = InflightActivationStatus::Failure; + records[3].status = ActivationStatus::Failure; replace_retry_state( &mut records[3], Some(RetryState { @@ -1555,7 +1564,7 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { delay_on_retry: None, }), ); - assert!(store.store(records.clone()).await.is_ok()); + assert!(store.write(records.clone()).await.is_ok()); assert_counts( StatusCount { failure: 4, @@ -1605,7 +1614,7 @@ async fn test_mark_completed(#[case] adapter: &str) { let store = create_test_store(adapter).await; let records = make_activations(3); - assert!(store.store(records.clone()).await.is_ok()); + assert!(store.write(records.clone()).await.is_ok()); assert_counts( StatusCount { pending: 3, @@ -1646,7 +1655,7 @@ async fn test_handle_expires_at(#[case] adapter: &str) { batch[1].expires_at = Some(Utc::now() + (Duration::from_secs(5 * 60))); batch[2].expires_at = Some(Utc::now() - (Duration::from_secs(5 * 60))); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { pending: 3, @@ -1686,7 +1695,7 @@ async fn test_remove_killswitched(#[case] adapter: &str) { batch[2].taskname = "task_to_be_killswitched_two".to_string(); batch[4].taskname = "task_to_be_killswitched_three".to_string(); - assert!(store.store(batch.clone()).await.is_ok()); + assert!(store.write(batch.clone()).await.is_ok()); assert_counts( StatusCount { pending: 6, @@ -1729,7 +1738,7 @@ async fn test_clear(#[case] adapter: &str) { let namespace = generate_unique_namespace(); let batch = vec![ - InflightActivationBuilder::new() + ActivationBuilder::new() .id("id_0") .taskname("taskname") .namespace(&namespace) @@ -1738,7 +1747,7 @@ async fn test_clear(#[case] adapter: &str) { .build(TaskActivationBuilder::new()), ]; - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); assert_counts( StatusCount { pending: 1, @@ -1763,7 +1772,7 @@ async fn test_full_vacuum(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let result = store.full_vacuum_db().await; assert!(result.is_ok()); @@ -1778,7 +1787,7 @@ async fn test_vacuum_db_no_limit(#[case] adapter: &str) { let store = create_test_store(adapter).await; let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let result = store.vacuum_db().await; assert!(result.is_ok()); @@ -1799,7 +1808,7 @@ async fn test_vacuum_db_incremental() { .expect("could not create store"); let batch = make_activations(2); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); let result = store.vacuum_db().await; assert!(result.is_ok()); @@ -1811,16 +1820,16 @@ async fn test_vacuum_db_incremental() { #[case::postgres("postgres")] async fn test_db_size(#[case] adapter: &str) { let store = create_test_store(adapter).await; - assert!(store.db_size().await.is_ok()); + assert!(store.size().await.is_ok()); - let first_size = store.db_size().await.unwrap(); + let first_size = store.size().await.unwrap(); assert!(first_size > 0, "should have some bytes"); // Generate a large enough batch that we use another page. let batch = make_activations(50); - assert!(store.store(batch).await.is_ok()); + assert!(store.write(batch).await.is_ok()); - let second_size = store.db_size().await.unwrap(); + let second_size = store.size().await.unwrap(); assert!(second_size > first_size, "should have more bytes now"); store.remove_db().await.unwrap(); } @@ -1836,8 +1845,8 @@ async fn test_pending_activation_max_lag_no_pending(#[case] adapter: &str) { assert_eq!(0.0, store.pending_activation_max_lag(&now).await); let mut processing = make_activations(1); - processing[0].status = InflightActivationStatus::Processing; - assert!(store.store(processing).await.is_ok()); + processing[0].status = ActivationStatus::Processing; + assert!(store.write(processing).await.is_ok()); // No pending activations, max lag is 0 assert_eq!(0.0, store.pending_activation_max_lag(&now).await); @@ -1855,7 +1864,7 @@ async fn test_pending_activation_max_lag_use_oldest(#[case] adapter: &str) { let mut pending = make_activations(2); pending[0].received_at = now - Duration::from_secs(10); pending[1].received_at = now - Duration::from_secs(500); - assert!(store.store(pending).await.is_ok()); + assert!(store.write(pending).await.is_ok()); let result = store.pending_activation_max_lag(&now).await; assert!(11.0 < result, "Should not get the small record"); @@ -1875,7 +1884,7 @@ async fn test_pending_activation_max_lag_ignore_processing_attempts(#[case] adap pending[0].received_at = now - Duration::from_secs(10); pending[1].received_at = now - Duration::from_secs(500); pending[1].processing_attempts = 1; - assert!(store.store(pending).await.is_ok()); + assert!(store.write(pending).await.is_ok()); let result = store.pending_activation_max_lag(&now).await; assert_eq!(result, 10.0, "max lag: {result:?}"); @@ -1895,7 +1904,7 @@ async fn test_pending_activation_max_lag_account_for_delayed(#[case] adapter: &s // the lag of a delayed task should begin *after* the delay has passed. pending[0].received_at = now - Duration::from_secs(520); pending[0].delay_until = Some(now - Duration::from_millis(22020)); - assert!(store.store(pending).await.is_ok()); + assert!(store.write(pending).await.is_ok()); let result = store.pending_activation_max_lag(&now).await; assert!(22.00 < result, "result: {result}"); diff --git a/src/store/traits.rs b/src/store/traits.rs index 4ebf7d82..e84a879a 100644 --- a/src/store/traits.rs +++ b/src/store/traits.rs @@ -4,19 +4,58 @@ use chrono::{DateTime, Utc}; use tokio::join; use tracing::warn; -use crate::store::activation::{InflightActivation, InflightActivationStatus}; +use crate::store::activation::{Activation, ActivationStatus}; use crate::store::types::{BucketRange, DepthCounts, FailedTasksForwarder}; +/// This trait only contains methods used by the consumer component. #[async_trait] -pub trait InflightActivationStore: Send + Sync { - /// CONSUMER OPERATIONS - /// Store a batch of activations - async fn store(&self, batch: Vec) -> Result; +pub trait IngestStore: Send + Sync { + /// Write a batch of activations to the store. + async fn write(&self, batch: Vec) -> Result; - fn assign_partitions(&self, partitions: Vec) -> Result<(), Error>; + /// Change the Kafka partitions for which this instance is responsible. + async fn assign_partitions(&self, partitions: Vec) -> Result<(), Error>; +} + +/// This trait only contains methods for counting operations, used for backpressure and metrics. +#[async_trait] +pub trait CountStore: Send + Sync { + /// Get the size of the database in bytes. + async fn size(&self) -> Result; + + /// Count all activations + async fn count(&self) -> Result; + + /// Count activations by status. + async fn count_by_status(&self, status: ActivationStatus) -> Result; + + /// Queue depths for pending, delay, and processing (writer backpressure and upkeep gauges). + /// Default implementation uses separate calls, but stores may override with a single query. + async fn count_depths(&self) -> Result { + let (pending, delay, claimed, processing) = join!( + self.count_by_status(ActivationStatus::Pending), + self.count_by_status(ActivationStatus::Delay), + self.count_by_status(ActivationStatus::Claimed), + self.count_by_status(ActivationStatus::Processing), + ); + + Ok(DepthCounts { + pending: pending?, + delay: delay?, + claimed: claimed?, + processing: processing?, + }) + } - /// Get `limit` pending activations, optionally filtered by namespaces and bucket subrange. - /// If `mark_processing` is true, sets status to `Processing` and `processing_deadline`; otherwise `Claimed` and `claim_expires_at`. + /// Get the age of the oldest pending activation in seconds + async fn pending_activation_max_lag(&self, now: &DateTime) -> f64; +} + +/// This trait only contains methods for managing claims. +#[async_trait] +pub trait ClaimStore: Send + Sync { + /// Get `limit` pending activations with several optional filters. + /// If `mark_processing` is true, sets `Processing` and `processing_deadline`, otherwise `Claimed` and `claim_expires_at`. /// If no limit is provided, all matching activations will be returned. async fn claim_activations( &self, @@ -25,16 +64,16 @@ pub trait InflightActivationStore: Send + Sync { limit: Option, bucket: Option, mark_processing: bool, - ) -> Result, Error>; + ) -> Result, Error>; - /// Claims `limit` activations within the `bucket` range. Push mode uses status `Claimed` until `mark_activation_processing` moves to `Processing`. + /// Claims `limit` activations and updates status to `Claimed` until `mark_processing` moves them to `Processing`. async fn claim_activations_for_push( &self, application: Option<&str>, namespaces: Option<&[String]>, limit: Option, bucket: Option, - ) -> Result, Error> { + ) -> Result, Error> { // If a namespace filter is used, an application must also be used if namespaces.is_some() && application.is_none() { warn!( @@ -49,12 +88,12 @@ pub trait InflightActivationStore: Send + Sync { .await } - /// Claims `limit` activations with application `application` and namespace `namespace`. + /// Claims one activation with application `application` and namespace `namespace`. async fn claim_activation_for_pull( &self, application: Option<&str>, namespace: Option<&str>, - ) -> Result, Error> { + ) -> Result, Error> { // Convert single namespace to vector for internal use let namespaces = namespace.map(|ns| vec![ns.to_string()]); @@ -79,78 +118,31 @@ pub trait InflightActivationStore: Send + Sync { Ok(rows.pop()) } } +} +/// This trait only contains methods used in push mode. +#[async_trait] +pub trait PushStore: Send + Sync { /// Record successful push. - async fn mark_activation_processing(&self, id: &str) -> Result<(), Error>; + async fn mark_processing(&self, id: &str) -> Result<(), Error>; +} - /// Update the status of a specific activation +/// This trait only contains methods used in pull mode. +#[async_trait] +pub trait PullStore: ClaimStore { + /// Update the status of a specific activation. async fn set_status( &self, id: &str, - status: InflightActivationStatus, - ) -> Result, Error>; - - /// COUNT OPERATIONS - /// Get the age of the oldest pending activation in seconds - async fn pending_activation_max_lag(&self, now: &DateTime) -> f64; - - /// Count activations with Pending status - async fn count_pending_activations(&self) -> Result { - self.count_by_status(InflightActivationStatus::Pending) - .await - } - - /// Count activations by status - async fn count_by_status(&self, status: InflightActivationStatus) -> Result; - - /// Count all activations - async fn count(&self) -> Result; - - /// ACTIVATION OPERATIONS - /// Get an activation by id - async fn get_by_id(&self, id: &str) -> Result, Error>; - - /// Queue depths for pending, delay, and processing (writer backpressure and upkeep gauges). - /// Default implementation uses separate calls, but stores may override with a single query. - async fn count_depths(&self) -> Result { - let (pending, delay, claimed, processing) = join!( - self.count_by_status(InflightActivationStatus::Pending), - self.count_by_status(InflightActivationStatus::Delay), - self.count_by_status(InflightActivationStatus::Claimed), - self.count_by_status(InflightActivationStatus::Processing), - ); - - Ok(DepthCounts { - pending: pending?, - delay: delay?, - claimed: claimed?, - processing: processing?, - }) - } - - /// Set the processing deadline for a specific activation - async fn set_processing_deadline( - &self, - id: &str, - deadline: Option>, - ) -> Result<(), Error>; - - /// Delete an activation by id - async fn delete_activation(&self, id: &str) -> Result<(), Error>; - - /// DATABASE OPERATIONS - /// Trigger incremental vacuum to reclaim free pages in the database - async fn vacuum_db(&self) -> Result<(), Error>; - - /// Perform a full vacuum on the database - async fn full_vacuum_db(&self) -> Result<(), Error>; - - /// Get the size of the database in bytes - async fn db_size(&self) -> Result; + status: ActivationStatus, + ) -> Result, Error>; +} - /// UPKEEP OPERATIONS +/// This trait only contains methods used by upkeep. +#[async_trait] +pub trait UpkeepStore { /// Get all activations with status Retry - async fn get_retry_activations(&self) -> Result, Error>; + async fn get_retry_activations(&self) -> Result, Error>; /// Revert expired push claims back to pending status. async fn handle_claim_expiration(&self) -> Result; @@ -167,19 +159,47 @@ pub trait InflightActivationStore: Send + Sync { /// Update delayed tasks past their delay_until deadline to Pending async fn handle_delay_until(&self) -> Result; - /// Process failed tasks for discard or deadletter + /// Process failed tasks for discard or deadletter. async fn handle_failed_tasks(&self) -> Result; - /// Mark tasks as complete by id + /// Mark tasks as completd by ID. async fn mark_completed(&self, ids: Vec) -> Result; - /// Remove completed tasks + /// Remove completed tasks. async fn remove_completed(&self) -> Result; - /// Remove killswitched tasks + /// Remove killswitched tasks. async fn remove_killswitched(&self, killswitched_tasks: Vec) -> Result; - /// TEST OPERATIONS + /// Trigger incremental vacuum to reclaim free pages in the database. + async fn vacuum_db(&self) -> Result<(), Error>; + + /// Perform a full vacuum on the database. + async fn full_vacuum_db(&self) -> Result<(), Error>; +} + +/// Umbrella interface that represents ALL possible operations on any runtime store. +#[async_trait] +pub trait Store: + IngestStore + ClaimStore + UpkeepStore + CountStore + PushStore + PullStore +{ +} + +#[async_trait] +pub trait TestStore: Store { + /// Get an activation by id + async fn get_by_id(&self, id: &str) -> Result, Error>; + + /// Set the processing deadline for a specific activation + async fn set_processing_deadline( + &self, + id: &str, + deadline: Option>, + ) -> Result<(), Error>; + + /// Delete an activation by id + async fn delete_activation(&self, id: &str) -> Result<(), Error>; + /// Clear all activations from the store async fn clear(&self) -> Result<(), Error>; diff --git a/src/test_utils.rs b/src/test_utils.rs index 3949dbdb..9dda1c54 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -15,12 +15,12 @@ use uuid::Uuid; use crate::{ config::Config, store::{ - activation::{InflightActivation, InflightActivationBuilder, InflightActivationStatus}, + activation::{Activation, ActivationBuilder, ActivationStatus}, adapters::{ - postgres::{PostgresActivationStore, PostgresActivationStoreConfig}, + postgres::{PostgresActivationStoreConfig, PostgresStore}, sqlite::{InflightActivationStoreConfig, SqliteActivationStore}, }, - traits::InflightActivationStore, + traits::{Store, TestStore}, }, }; @@ -134,12 +134,12 @@ impl Default for TaskActivationBuilder { } } -impl InflightActivationBuilder { +impl ActivationBuilder { pub fn new() -> Self { Self::default() } - pub fn build(mut self, builder: TaskActivationBuilder) -> InflightActivation { + pub fn build(mut self, builder: TaskActivationBuilder) -> Activation { // Grab required fields let id = self.id.as_ref().expect("field 'id' is required"); @@ -230,13 +230,13 @@ pub fn generate_unique_namespace() -> String { } /// Create a collection of `count` pending unsaved activations in a particular `namespace`. If you do not want to provide a namespace, use `make_activations`. -pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec { - let mut records: Vec = vec![]; +pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec { + let mut records: Vec = vec![]; for i in 0..count { let now = Utc::now(); - let item = InflightActivationBuilder::new() + let item = ActivationBuilder::new() .id(format!("id_{i}")) .taskname("taskname") .namespace(&namespace) @@ -252,7 +252,7 @@ pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec Vec { +pub fn make_activations(count: u32) -> Vec { let namespace = generate_unique_namespace(); make_activations_with_namespace(namespace, count) } @@ -263,7 +263,7 @@ pub fn create_config() -> Arc { } /// Create an InflightActivationStore instance -pub async fn create_test_store(adapter: &str) -> Arc { +pub async fn create_test_store(adapter: &str) -> Arc { match adapter { "sqlite" => Arc::new( SqliteActivationStore::new( @@ -272,16 +272,16 @@ pub async fn create_test_store(adapter: &str) -> Arc, + ) as Arc, "postgres" => { let store = Arc::new( - PostgresActivationStore::new(PostgresActivationStoreConfig::from_config( + PostgresStore::new(PostgresActivationStoreConfig::from_config( &create_integration_config(), )) .await .unwrap(), - ) as Arc; - store.assign_partitions(vec![0]).unwrap(); + ) as Arc; + store.assign_partitions(vec![0]).await.unwrap(); store } _ => panic!("Invalid adapter: {}", adapter), @@ -426,7 +426,7 @@ pub async fn consume_topic( results } -pub fn replace_retry_state(inflight: &mut InflightActivation, retry: Option) { +pub fn replace_retry_state(inflight: &mut Activation, retry: Option) { let mut activation = TaskActivation::decode(&inflight.activation as &[u8]).unwrap(); activation.retry_state = retry; inflight.activation = activation.encode_to_vec(); @@ -451,11 +451,11 @@ pub struct StatusCount { } /// Assert the state of all counts in the inflight activation store. -pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivationStore) { +pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.pending, store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), "difference in pending count", @@ -463,7 +463,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.claimed, store - .count_by_status(InflightActivationStatus::Claimed) + .count_by_status(ActivationStatus::Claimed) .await .unwrap(), "difference in claimed count", @@ -471,7 +471,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.processing, store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(), "difference in processing count", @@ -479,7 +479,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.retry, store - .count_by_status(InflightActivationStatus::Retry) + .count_by_status(ActivationStatus::Retry) .await .unwrap(), "difference in retry count", @@ -487,7 +487,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.delayed, store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(), "difference in delay count", @@ -495,7 +495,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.complete, store - .count_by_status(InflightActivationStatus::Complete) + .count_by_status(ActivationStatus::Complete) .await .unwrap(), "difference in complete count", @@ -503,7 +503,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn InflightActivation assert_eq!( expected.failure, store - .count_by_status(InflightActivationStatus::Failure) + .count_by_status(ActivationStatus::Failure) .await .unwrap(), "difference in failure count", diff --git a/src/upkeep.rs b/src/upkeep.rs index 9132d3f5..911a129a 100644 --- a/src/upkeep.rs +++ b/src/upkeep.rs @@ -18,13 +18,13 @@ use uuid::Uuid; use crate::SERVICE_NAME; use crate::config::Config; use crate::runtime_config::RuntimeConfigManager; -use crate::store::traits::InflightActivationStore; +use crate::store::traits::UpkeepStore; /// The upkeep task that periodically performs upkeep /// on the inflight store pub async fn upkeep( config: Arc, - store: Arc, + store: Arc, startup_time: DateTime, runtime_config_manager: Arc, health_reporter: HealthReporter, @@ -109,7 +109,7 @@ impl UpkeepResults { )] pub async fn do_upkeep( config: Arc, - store: Arc, + store: Arc, producer: Arc, startup_time: DateTime, runtime_config_manager: Arc, @@ -536,7 +536,7 @@ mod tests { use crate::{ config::Config, runtime_config::RuntimeConfigManager, - store::activation::InflightActivationStatus, + store::activation::ActivationStatus, test_utils::{ StatusCount, assert_counts, consume_topic, create_config, create_integration_config, create_integration_config_with_topic, create_producer, create_test_store, @@ -664,7 +664,7 @@ mod tests { .expect(""); activation.parameters = r#"{"a":"b"}"#.into(); activation.delay = Some(30); - records[0].status = InflightActivationStatus::Retry; + records[0].status = ActivationStatus::Retry; records[0].delay_until = Some(Utc::now() + Duration::from_secs(30)); records[0].activation = activation.encode_to_vec(); @@ -724,7 +724,7 @@ mod tests { let mut batch = make_activations(2); // Make a task with a future processing deadline - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc::now() + TimeDelta::minutes(5)); assert!(store.store(batch.clone()).await.is_ok()); @@ -741,7 +741,7 @@ mod tests { // Should retain the processing record assert_eq!( store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(), 1 @@ -760,7 +760,7 @@ mod tests { let mut batch = make_activations(2); // Make a task past with a processing deadline in the past - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); assert!(store.store(batch.clone()).await.is_ok()); @@ -768,7 +768,7 @@ mod tests { // Should start off with one in processing assert_eq!( store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(), 1 @@ -816,7 +816,7 @@ mod tests { let mut batch = make_activations(2); // Make a task past with a processing deadline in the past - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[1].processing_attempts = 0; @@ -825,14 +825,14 @@ mod tests { // Should start off with one in processing assert_eq!( store - .count_by_status(InflightActivationStatus::Processing) + .count_by_status(ActivationStatus::Processing) .await .unwrap(), 1 ); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 1 @@ -885,7 +885,7 @@ mod tests { delay_on_retry: None, }), ); - batch[1].status = InflightActivationStatus::Processing; + batch[1].status = ActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[1].at_most_once = true; @@ -930,7 +930,7 @@ mod tests { // Because 1 is complete and has a higher offset than 0, index 2 can be discarded batch[0].processing_attempts = config.max_processing_attempts as i32; - batch[1].status = InflightActivationStatus::Complete; + batch[1].status = ActivationStatus::Complete; batch[1].added_at += Duration::from_secs(1); batch[2].processing_attempts = config.max_processing_attempts as i32; @@ -952,7 +952,7 @@ mod tests { assert_eq!(result_context.completed, 3); // all three are removed as completed assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 0, @@ -960,7 +960,7 @@ mod tests { ); assert_eq!( store - .count_by_status(InflightActivationStatus::Complete) + .count_by_status(ActivationStatus::Complete) .await .unwrap(), 0, @@ -992,7 +992,7 @@ mod tests { delay_on_retry: None, }), ); - records[0].status = InflightActivationStatus::Failure; + records[0].status = ActivationStatus::Failure; records[1].added_at += Duration::from_secs(1); assert!(store.store(records.clone()).await.is_ok()); @@ -1035,7 +1035,7 @@ mod tests { let mut last_vacuum = Instant::now(); let mut batch = make_activations(2); - batch[0].status = InflightActivationStatus::Failure; + batch[0].status = ActivationStatus::Failure; batch[1].added_at += Duration::from_secs(1); assert!(store.store(batch).await.is_ok()); @@ -1058,7 +1058,7 @@ mod tests { ); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 1, @@ -1081,7 +1081,7 @@ mod tests { let mut batch = make_activations(4); batch[0].expires_at = Some(Utc::now() - Duration::from_secs(100)); - batch[1].status = InflightActivationStatus::Complete; + batch[1].status = ActivationStatus::Complete; batch[2].expires_at = Some(Utc::now() - Duration::from_secs(100)); // Ensure the fourth task is in the future @@ -1103,7 +1103,7 @@ mod tests { assert_eq!(result_context.completed, 1); // 1 complete assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 1, @@ -1111,7 +1111,7 @@ mod tests { ); assert_eq!( store - .count_by_status(InflightActivationStatus::Complete) + .count_by_status(ActivationStatus::Complete) .await .unwrap(), 0, @@ -1150,23 +1150,23 @@ mod tests { let mut batch = make_activations(2); - batch[0].status = InflightActivationStatus::Delay; + batch[0].status = ActivationStatus::Delay; batch[0].delay_until = Some(Utc::now() - Duration::from_secs(1)); - batch[1].status = InflightActivationStatus::Delay; + batch[1].status = ActivationStatus::Delay; batch[1].delay_until = Some(Utc::now() + Duration::from_secs(1)); assert!(store.store(batch.clone()).await.is_ok()); assert_eq!( store - .count_by_status(InflightActivationStatus::Delay) + .count_by_status(ActivationStatus::Delay) .await .unwrap(), 2 ); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 0 @@ -1183,7 +1183,7 @@ mod tests { assert_eq!(result_context.delay_elapsed, 1); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 1 @@ -1215,7 +1215,7 @@ mod tests { assert_eq!(result_context.delay_elapsed, 1); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 1 @@ -1269,7 +1269,7 @@ demoted_namespaces: assert_eq!(result_context.forwarded, 2); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 4, @@ -1277,7 +1277,7 @@ demoted_namespaces: ); assert_eq!( store - .count_by_status(InflightActivationStatus::Complete) + .count_by_status(ActivationStatus::Complete) .await .unwrap(), 2, @@ -1328,7 +1328,7 @@ demoted_namespaces: assert_eq!(result_context.killswitched, 3); assert_eq!( store - .count_by_status(InflightActivationStatus::Pending) + .count_by_status(ActivationStatus::Pending) .await .unwrap(), 3 From ff552a621341fd04e9d6fcff1d81d6c6fbc17b7c Mon Sep 17 00:00:00 2001 From: james-mcnulty Date: Thu, 16 Apr 2026 15:45:30 -0700 Subject: [PATCH 2/2] Undo Naming Changes --- benches/store_bench.rs | 6 +- src/fetch/mod.rs | 6 +- src/fetch/tests.rs | 14 +- src/grpc/server.rs | 13 +- src/grpc/server_tests.rs | 4 +- src/kafka/deserialize_activation.rs | 22 +-- src/kafka/inflight_activation_batcher.rs | 26 ++-- src/kafka/inflight_activation_writer.rs | 50 +++---- src/push/mod.rs | 12 +- src/store/activation.rs | 50 ++++--- src/store/adapters/postgres.rs | 68 ++++----- src/store/adapters/sqlite.rs | 68 ++++----- src/store/tests.rs | 179 ++++++++++++----------- src/store/traits.rs | 28 ++-- src/test_utils.rs | 30 ++-- src/upkeep.rs | 56 +++---- 16 files changed, 327 insertions(+), 305 deletions(-) diff --git a/benches/store_bench.rs b/benches/store_bench.rs index a868cf9f..6f110e23 100644 --- a/benches/store_bench.rs +++ b/benches/store_bench.rs @@ -4,7 +4,7 @@ use chrono::Utc; use criterion::{Criterion, criterion_group, criterion_main}; use rand::Rng; use taskbroker::{ - store::activation::ActivationStatus, + store::activation::InflightActivationStatus, store::adapters::sqlite::{InflightActivationStoreConfig, SqliteActivationStore}, store::traits::InflightActivationStore, test_utils::{ @@ -122,7 +122,7 @@ async fn set_status(num_activations: u32, num_workers: u32) { for task_id in 0..num_activations { if task_id % num_workers == worker_idx { store - .set_status(&format!("id_{task_id}"), ActivationStatus::Complete) + .set_status(&format!("id_{task_id}"), InflightActivationStatus::Complete) .await .unwrap(); } @@ -134,7 +134,7 @@ async fn set_status(num_activations: u32, num_workers: u32) { assert_eq!( store - .count_by_status(ActivationStatus::Complete) + .count_by_status(InflightActivationStatus::Complete) .await .unwrap(), num_activations as usize diff --git a/src/fetch/mod.rs b/src/fetch/mod.rs index e5113747..90026ed5 100644 --- a/src/fetch/mod.rs +++ b/src/fetch/mod.rs @@ -9,7 +9,7 @@ use tracing::{debug, info, warn}; use crate::config::Config; use crate::push::{PushError, PushPool}; -use crate::store::activation::Activation; +use crate::store::activation::InflightActivation; use crate::store::traits::ClaimStore; use crate::store::types::BucketRange; @@ -49,12 +49,12 @@ pub fn bucket_range_for_fetch_thread(thread_index: usize, fetch_threads: usize) #[async_trait] pub trait TaskPusher { /// Submit a single task to the push pool. - async fn submit_task(&self, activation: Activation) -> Result<(), PushError>; + async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError>; } #[async_trait] impl TaskPusher for PushPool { - async fn submit_task(&self, activation: Activation) -> Result<(), PushError> { + async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError> { self.submit(activation).await } } diff --git a/src/fetch/tests.rs b/src/fetch/tests.rs index 8332b5c3..ede79e4c 100644 --- a/src/fetch/tests.rs +++ b/src/fetch/tests.rs @@ -8,7 +8,7 @@ use tonic::async_trait; use super::*; use crate::config::Config; use crate::push::PushError; -use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::activation::{InflightActivation, InflightActivationStatus}; use crate::store::traits::ClaimStore; use crate::store::types::BucketRange; use crate::test_utils::make_activations; @@ -16,7 +16,7 @@ use crate::test_utils::make_activations; /// Store stub that returns one activation once OR is always empty OR always fails. struct MockStore { /// A single (optional) pending activation. - pending: Mutex>, + pending: Mutex>, /// Should operations fail? fail: bool, @@ -30,7 +30,7 @@ impl MockStore { } } - fn one(activation: Activation) -> Self { + fn one(activation: InflightActivation) -> Self { Self { pending: Mutex::new(Some(activation)), fail: false, @@ -54,7 +54,7 @@ impl ClaimStore for MockStore { _limit: Option, _bucket: Option, mark_processing: bool, - ) -> Result, Error> { + ) -> Result, Error> { if self.fail { return Err(anyhow!("mock store error")); } @@ -62,9 +62,9 @@ impl ClaimStore for MockStore { Ok(match self.pending.lock().await.take() { Some(mut a) => { a.status = if mark_processing { - ActivationStatus::Processing + InflightActivationStatus::Processing } else { - ActivationStatus::Claimed + InflightActivationStatus::Claimed }; vec![a] } @@ -91,7 +91,7 @@ impl RecordingPusher { #[async_trait] impl TaskPusher for RecordingPusher { - async fn submit_task(&self, activation: Activation) -> Result<(), PushError> { + async fn submit_task(&self, activation: InflightActivation) -> Result<(), PushError> { self.pushed_ids.lock().await.push(activation.id.clone()); if self.fail { diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 155acdce..e15005cb 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -10,7 +10,7 @@ use std::time::Instant; use tonic::{Request, Response, Status}; use crate::config::{Config, DeliveryMode}; -use crate::store::activation::ActivationStatus; +use crate::store::activation::InflightActivationStatus; use crate::store::traits::PullStore; use tracing::{error, instrument, warn}; @@ -84,16 +84,19 @@ impl ConsumerService for TaskbrokerServer { let start_time = Instant::now(); let id = request.get_ref().id.clone(); - let status: ActivationStatus = TaskActivationStatus::try_from(request.get_ref().status) - .map_err(|e| Status::invalid_argument(format!("Unable to deserialize status: {e:?}")))? - .into(); + let status: InflightActivationStatus = + TaskActivationStatus::try_from(request.get_ref().status) + .map_err(|e| { + Status::invalid_argument(format!("Unable to deserialize status: {e:?}")) + })? + .into(); if !status.is_conclusion() { return Err(Status::invalid_argument(format!( "Invalid status, expects 3 (Failure), 4 (Retry), or 5 (Complete), but got: {status:?}" ))); } - if status == ActivationStatus::Failure { + if status == InflightActivationStatus::Failure { metrics::counter!("grpc_server.set_status.failure").increment(1); } diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index d073b544..81461c79 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use crate::config::{Config, DeliveryMode}; use crate::grpc::server::TaskbrokerServer; -use crate::store::activation::ActivationStatus; +use crate::store::activation::InflightActivationStatus; use prost::Message; use rstest::rstest; use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService; @@ -128,7 +128,7 @@ async fn test_get_task_success(#[case] adapter: &str) { assert!(task.id == "id_0"); let row = store.get_by_id("id_0").await.unwrap().expect("claimed row"); - assert_eq!(row.status, ActivationStatus::Processing); + assert_eq!(row.status, InflightActivationStatus::Processing); } #[tokio::test] diff --git a/src/kafka/deserialize_activation.rs b/src/kafka/deserialize_activation.rs index b9df63b6..2288d534 100644 --- a/src/kafka/deserialize_activation.rs +++ b/src/kafka/deserialize_activation.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use crate::config::Config; use crate::fetch::MAX_FETCH_THREADS; -use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::activation::{InflightActivation, InflightActivationStatus}; use anyhow::{Error, anyhow}; use chrono::{DateTime, Utc}; use prost::Message as _; @@ -32,7 +32,7 @@ pub fn bucket_from_id(id: &str) -> i16 { pub fn new( config: DeserializeActivationConfig, -) -> impl Fn(Arc) -> Result { +) -> impl Fn(Arc) -> Result { move |msg: Arc| { let Some(payload) = msg.payload() else { return Err(anyhow!("Message has no payload")); @@ -72,11 +72,11 @@ pub fn new( activation_time + delay }); - let status = delay_until.map_or(ActivationStatus::Pending, |delay_until| { + let status = delay_until.map_or(InflightActivationStatus::Pending, |delay_until| { if Utc::now() > delay_until { - ActivationStatus::Pending + InflightActivationStatus::Pending } else { - ActivationStatus::Delay + InflightActivationStatus::Delay } }); @@ -89,7 +89,7 @@ pub fn new( let bucket = bucket_from_id(&activation.id); - Ok(Activation { + Ok(InflightActivation { id: activation.id.clone(), activation: payload.to_vec(), status, @@ -122,7 +122,9 @@ mod tests { use rdkafka::{Timestamp, message::OwnedMessage}; use sentry_protos::taskbroker::v1::TaskActivation; - use crate::{store::activation::ActivationStatus, test_utils::generate_unique_namespace}; + use crate::{ + store::activation::InflightActivationStatus, test_utils::generate_unique_namespace, + }; use super::{Config, DeserializeActivationConfig, new}; @@ -261,7 +263,7 @@ mod tests { delta.num_seconds() >= 99, "Should have ~100 seconds of delay from received_at" ); - assert_eq!(inflight.status, ActivationStatus::Pending) + assert_eq!(inflight.status, InflightActivationStatus::Pending) } #[test] @@ -307,7 +309,7 @@ mod tests { delta.num_seconds() >= 99, "Should have ~100 seconds of delay from received_at" ); - assert_eq!(inflight.status, ActivationStatus::Delay) + assert_eq!(inflight.status, InflightActivationStatus::Delay) } #[test] @@ -354,6 +356,6 @@ mod tests { delta.num_seconds() <= (config.max_delayed_task_allowed_sec as f64 * 1.1) as i64, "Should have approxmiately max_delayed_task_allowed_sec of delay from received_at" ); - assert_eq!(inflight.status, ActivationStatus::Delay) + assert_eq!(inflight.status, InflightActivationStatus::Delay) } } diff --git a/src/kafka/inflight_activation_batcher.rs b/src/kafka/inflight_activation_batcher.rs index a6f5788e..79f50328 100644 --- a/src/kafka/inflight_activation_batcher.rs +++ b/src/kafka/inflight_activation_batcher.rs @@ -1,4 +1,6 @@ -use crate::{config::Config, runtime_config::RuntimeConfigManager, store::activation::Activation}; +use crate::{ + config::Config, runtime_config::RuntimeConfigManager, store::activation::InflightActivation, +}; use chrono::Utc; use futures::future::join_all; use rdkafka::config::ClientConfig; @@ -41,7 +43,7 @@ impl ActivationBatcherConfig { } pub struct InflightActivationBatcher { - batch: Vec, + batch: Vec, batch_size: usize, forward_batch: Vec>, // payload config: ActivationBatcherConfig, @@ -79,9 +81,9 @@ impl InflightActivationBatcher { } impl Reducer for InflightActivationBatcher { - type Input = Activation; + type Input = InflightActivation; - type Output = Vec; + type Output = Vec; async fn reduce(&mut self, t: Self::Input) -> Result<(), anyhow::Error> { let runtime_config = self.runtime_config_manager.read().await; @@ -218,7 +220,7 @@ mod tests { use std::sync::Arc; use crate::{ - store::activation::ActivationBuilder, + store::activation::InflightActivationBuilder, test_utils::{TaskActivationBuilder, generate_unique_namespace}, }; @@ -242,7 +244,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = ActivationBuilder::new() + let inflight_activation_0 = InflightActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace(&namespace) @@ -265,7 +267,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = ActivationBuilder::new() + let inflight_activation_0 = InflightActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace(&namespace) @@ -292,7 +294,7 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = ActivationBuilder::new() + let inflight_activation_0 = InflightActivationBuilder::new() .id("0") .taskname("taskname") .namespace(&namespace) @@ -320,13 +322,13 @@ demoted_namespaces: let namespace = generate_unique_namespace(); - let inflight_activation_0 = ActivationBuilder::new() + let inflight_activation_0 = InflightActivationBuilder::new() .id("0") .taskname("taskname") .namespace(&namespace) .build(TaskActivationBuilder::new()); - let inflight_activation_1 = ActivationBuilder::new() + let inflight_activation_1 = InflightActivationBuilder::new() .id("1") .taskname("taskname") .namespace(&namespace) @@ -361,13 +363,13 @@ demoted_topic: taskworker-demoted"#; assert_eq!(batcher.producer_cluster, config.kafka_cluster.clone()); - let inflight_activation_0 = ActivationBuilder::new() + let inflight_activation_0 = InflightActivationBuilder::new() .id("0") .taskname("task_to_be_filtered") .namespace("bad_namespace") .build(TaskActivationBuilder::new()); - let inflight_activation_1 = ActivationBuilder::new() + let inflight_activation_1 = InflightActivationBuilder::new() .id("1") .taskname("good_task") .namespace("good_namespace") diff --git a/src/kafka/inflight_activation_writer.rs b/src/kafka/inflight_activation_writer.rs index 708d05ca..820ef037 100644 --- a/src/kafka/inflight_activation_writer.rs +++ b/src/kafka/inflight_activation_writer.rs @@ -10,7 +10,7 @@ use tracing::{debug, error, instrument}; use crate::{ config::Config, store::{ - activation::{Activation, ActivationStatus}, + activation::{InflightActivation, InflightActivationStatus}, traits::{CountStore, IngestStore}, types::DepthCounts, }, @@ -47,7 +47,7 @@ impl ActivationWriterConfig { pub struct InflightActivationWriter { config: ActivationWriterConfig, store: Arc, - batch: Option>, + batch: Option>, } impl InflightActivationWriter { @@ -61,7 +61,7 @@ impl InflightActivationWriter { } impl Reducer for InflightActivationWriter { - type Input = Vec; + type Input = Vec; type Output = (); @@ -112,10 +112,10 @@ impl Reducer for InflightActivationWriter; @@ -42,7 +42,7 @@ pub enum PushError { Timeout, /// Channel disconnected (no receivers) or another failure. - Channel(SendError), + Channel(SendError), } /// Thin interface for the worker client. It mostly serves to enable proper unit testing, but it also decouples the actual client implementation from our pushing logic. @@ -82,10 +82,10 @@ impl WorkerClient for WorkerServiceClient { /// Wrapper around `config.push_threads` asynchronous tasks, each of which receives an activation from the channel, sends it to the worker service, and repeats. pub struct PushPool { /// The sending end of a channel that accepts task activations. - sender: Sender, + sender: Sender, /// The receiving end of a channel that accepts task activations. - receiver: Receiver, + receiver: Receiver, /// Taskbroker configuration. config: Arc, @@ -268,7 +268,7 @@ impl PushPool { } /// Send an activation to the internal asynchronous MPMC channel used by all running push threads. Times out after `config.push_queue_timeout_ms` milliseconds. - pub async fn submit(&self, activation: Activation) -> Result<(), PushError> { + pub async fn submit(&self, activation: InflightActivation) -> Result<(), PushError> { let duration = Duration::from_millis(self.config.push_queue_timeout_ms); let start = Instant::now(); @@ -298,7 +298,7 @@ impl PushPool { /// Decode task activation and push it to a worker. async fn push_task( worker: &mut W, - activation: Activation, + activation: InflightActivation, callback_url: String, timeout: Duration, grpc_shared_secret: &[String], diff --git a/src/store/activation.rs b/src/store/activation.rs index f22c8012..52c2326b 100644 --- a/src/store/activation.rs +++ b/src/store/activation.rs @@ -8,7 +8,7 @@ use std::str::FromStr; /// The members of this enum should be a superset of the members /// of `InflightActivationStatus` in `sentry_protos`. #[derive(Clone, Copy, Debug, PartialEq, Eq, Type)] -pub enum ActivationStatus { +pub enum InflightActivationStatus { /// Unused but necessary to align with sentry-protos Unspecified, Pending, @@ -20,57 +20,59 @@ pub enum ActivationStatus { Delay, } -impl Display for ActivationStatus { +impl Display for InflightActivationStatus { fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { write!(f, "{:?}", self) } } -impl FromStr for ActivationStatus { +impl FromStr for InflightActivationStatus { type Err = String; fn from_str(s: &str) -> Result { if s == "Unspecified" { - Ok(ActivationStatus::Unspecified) + Ok(InflightActivationStatus::Unspecified) } else if s == "Pending" { - Ok(ActivationStatus::Pending) + Ok(InflightActivationStatus::Pending) } else if s == "Claimed" { - Ok(ActivationStatus::Claimed) + Ok(InflightActivationStatus::Claimed) } else if s == "Processing" { - Ok(ActivationStatus::Processing) + Ok(InflightActivationStatus::Processing) } else if s == "Failure" { - Ok(ActivationStatus::Failure) + Ok(InflightActivationStatus::Failure) } else if s == "Retry" { - Ok(ActivationStatus::Retry) + Ok(InflightActivationStatus::Retry) } else if s == "Complete" { - Ok(ActivationStatus::Complete) + Ok(InflightActivationStatus::Complete) } else if s == "Delay" { - Ok(ActivationStatus::Delay) + Ok(InflightActivationStatus::Delay) } else { Err(format!("Unknown inflight activation status string: {}", s)) } } } -impl ActivationStatus { +impl InflightActivationStatus { /// Is the current value a 'conclusion' status that can be supplied over GRPC. pub fn is_conclusion(&self) -> bool { matches!( self, - ActivationStatus::Complete | ActivationStatus::Retry | ActivationStatus::Failure + InflightActivationStatus::Complete + | InflightActivationStatus::Retry + | InflightActivationStatus::Failure ) } } -impl From for ActivationStatus { +impl From for InflightActivationStatus { fn from(item: TaskActivationStatus) -> Self { match item { - TaskActivationStatus::Unspecified => ActivationStatus::Unspecified, - TaskActivationStatus::Pending => ActivationStatus::Pending, - TaskActivationStatus::Processing => ActivationStatus::Processing, - TaskActivationStatus::Failure => ActivationStatus::Failure, - TaskActivationStatus::Retry => ActivationStatus::Retry, - TaskActivationStatus::Complete => ActivationStatus::Complete, + TaskActivationStatus::Unspecified => InflightActivationStatus::Unspecified, + TaskActivationStatus::Pending => InflightActivationStatus::Pending, + TaskActivationStatus::Processing => InflightActivationStatus::Processing, + TaskActivationStatus::Failure => InflightActivationStatus::Failure, + TaskActivationStatus::Retry => InflightActivationStatus::Retry, + TaskActivationStatus::Complete => InflightActivationStatus::Complete, } } } @@ -79,7 +81,7 @@ impl From for ActivationStatus { #[builder(pattern = "owned")] #[builder(build_fn(name = "_build"))] #[builder(field(public))] -pub struct Activation { +pub struct InflightActivation { #[builder(setter(into))] pub id: String, @@ -100,8 +102,8 @@ pub struct Activation { pub activation: Vec, /// The current status of the activation - #[builder(default = ActivationStatus::Pending)] - pub status: ActivationStatus, + #[builder(default = InflightActivationStatus::Pending)] + pub status: InflightActivationStatus, /// The partition the activation was received from #[builder(default = 0)] @@ -162,7 +164,7 @@ pub struct Activation { pub bucket: i16, } -impl Activation { +impl InflightActivation { /// The number of milliseconds between an activation's received timestamp /// and the provided datetime pub fn received_latency(&self, now: DateTime) -> i64 { diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 00eb72e0..dce1edb6 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -14,7 +14,7 @@ use std::time::Instant; use tracing::{instrument, warn}; use crate::config::Config; -use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::activation::{InflightActivation, InflightActivationStatus}; use crate::store::traits::{ ClaimStore, CountStore, IngestStore, PullStore, PushStore, Store, TestStore, UpkeepStore, }; @@ -44,10 +44,10 @@ struct TableRow { pub bucket: i16, } -impl TryFrom for TableRow { +impl TryFrom for TableRow { type Error = anyhow::Error; - fn try_from(value: Activation) -> Result { + fn try_from(value: InflightActivation) -> Result { Ok(Self { id: value.id, activation: value.activation, @@ -72,12 +72,12 @@ impl TryFrom for TableRow { } } -impl From for Activation { +impl From for InflightActivation { fn from(value: TableRow) -> Self { Self { id: value.id, activation: value.activation, - status: ActivationStatus::from_str(&value.status).unwrap(), + status: InflightActivationStatus::from_str(&value.status).unwrap(), partition: value.partition, offset: value.offset, added_at: value.added_at, @@ -252,7 +252,7 @@ impl ClaimStore for PostgresStore { limit: Option, bucket: Option, mark_processing: bool, - ) -> Result, Error> { + ) -> Result, Error> { let now = Utc::now(); let grace_period = self.config.processing_deadline_grace_sec; @@ -264,7 +264,7 @@ impl ClaimStore for PostgresStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push(" AND (expires_at IS NULL OR expires_at > "); query_builder.push_bind(now); query_builder.push(")"); @@ -310,7 +310,7 @@ impl ClaimStore for PostgresStore { status = " )); - query_builder.push_bind(ActivationStatus::Processing.to_string()); + query_builder.push_bind(InflightActivationStatus::Processing.to_string()); } else { query_builder.push(format!( "UPDATE inflight_taskactivations @@ -319,7 +319,7 @@ impl ClaimStore for PostgresStore { status = " )); - query_builder.push_bind(ActivationStatus::Claimed.to_string()); + query_builder.push_bind(InflightActivationStatus::Claimed.to_string()); } query_builder.push(" FROM selected_activations "); @@ -352,9 +352,9 @@ impl PushStore for PostgresStore { claim_expires_at = NULL WHERE id = $2 AND status = $3", )) - .bind(ActivationStatus::Processing.to_string()) + .bind(InflightActivationStatus::Processing.to_string()) .bind(id) - .bind(ActivationStatus::Claimed.to_string()) + .bind(InflightActivationStatus::Claimed.to_string()) .execute(&mut *conn) .await?; @@ -381,8 +381,8 @@ impl PullStore for PostgresStore { async fn set_status( &self, id: &str, - status: ActivationStatus, - ) -> Result, Error> { + status: InflightActivationStatus, + ) -> Result, Error> { let mut conn = self.acquire_write_conn_metric("set_status").await?; let result: Option = sqlx::query_as( "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *, kafka_offset AS offset", @@ -403,7 +403,7 @@ impl PullStore for PostgresStore { #[async_trait] impl CountStore for PostgresStore { #[instrument(skip_all)] - async fn count_by_status(&self, status: ActivationStatus) -> Result { + async fn count_by_status(&self, status: InflightActivationStatus) -> Result { let mut query_builder = QueryBuilder::new( "SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = ", ); @@ -449,7 +449,7 @@ impl CountStore for PostgresStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push(" AND processing_attempts = 0"); self.add_partition_condition(&mut query_builder, false); @@ -506,7 +506,7 @@ impl CountStore for PostgresStore { #[async_trait] impl TestStore for PostgresStore { /// Get an activation by id. Primarily used for testing - async fn get_by_id(&self, id: &str) -> Result, Error> { + async fn get_by_id(&self, id: &str) -> Result, Error> { let row_result: Option = sqlx::query_as( " SELECT id, @@ -604,7 +604,7 @@ impl IngestStore for PostgresStore { } #[instrument(skip_all)] - async fn write(&self, batch: Vec) -> Result { + async fn write(&self, batch: Vec) -> Result { if batch.is_empty() { return Ok(0); } @@ -698,7 +698,7 @@ impl UpkeepStore for PostgresStore { } #[instrument(skip_all)] - async fn get_retry_activations(&self) -> Result, Error> { + async fn get_retry_activations(&self) -> Result, Error> { let mut query_builder = QueryBuilder::new( "SELECT id, activation, @@ -722,7 +722,7 @@ impl UpkeepStore for PostgresStore { FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(ActivationStatus::Retry.to_string()); + query_builder.push_bind(InflightActivationStatus::Retry.to_string()); self.add_partition_condition(&mut query_builder, false); Ok(query_builder @@ -747,14 +747,14 @@ impl UpkeepStore for PostgresStore { SET claim_expires_at = null, status = ", ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push( " WHERE claim_expires_at IS NOT NULL AND claim_expires_at < ", ); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Claimed.to_string()); + query_builder.push_bind(InflightActivationStatus::Claimed.to_string()); self.add_partition_condition(&mut query_builder, false); let released = query_builder.build().execute(&mut *conn).await?; @@ -776,11 +776,11 @@ impl UpkeepStore for PostgresStore { "UPDATE inflight_taskactivations SET processing_deadline = null, status = ", ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push_bind(InflightActivationStatus::Failure.to_string()); query_builder.push(" WHERE processing_deadline < "); query_builder.push_bind(now); query_builder.push(" AND at_most_once = TRUE AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); + query_builder.push_bind(InflightActivationStatus::Processing.to_string()); self.add_partition_condition(&mut query_builder, false); @@ -797,12 +797,12 @@ impl UpkeepStore for PostgresStore { "UPDATE inflight_taskactivations SET processing_deadline = null, status = ", ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push(", processing_attempts = processing_attempts + 1"); query_builder.push(" WHERE processing_deadline < "); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); + query_builder.push_bind(InflightActivationStatus::Processing.to_string()); self.add_partition_condition(&mut query_builder, false); let result = query_builder.build().execute(&mut *atomic).await; @@ -828,11 +828,11 @@ impl UpkeepStore for PostgresStore { "UPDATE inflight_taskactivations SET status = ", ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push_bind(InflightActivationStatus::Failure.to_string()); query_builder.push(" WHERE processing_attempts >= "); query_builder.push_bind(self.config.max_processing_attempts as i32); query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); self.add_partition_condition(&mut query_builder, false); let processing_attempts_result = query_builder.build().execute(&mut *conn).await; @@ -855,7 +855,7 @@ impl UpkeepStore for PostgresStore { let mut conn = self.acquire_write_conn_metric("handle_expires_at").await?; let mut query_builder = QueryBuilder::new("DELETE FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push(" AND expires_at IS NOT NULL AND expires_at < "); query_builder.push_bind(now); self.add_partition_condition(&mut query_builder, false); @@ -879,11 +879,11 @@ impl UpkeepStore for PostgresStore { "UPDATE inflight_taskactivations SET status = ", ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push_bind(InflightActivationStatus::Pending.to_string()); query_builder.push(" WHERE delay_until IS NOT NULL AND delay_until < "); query_builder.push_bind(now); query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Delay.to_string()); + query_builder.push_bind(InflightActivationStatus::Delay.to_string()); self.add_partition_condition(&mut query_builder, false); let update_result = query_builder.build().execute(&mut *conn).await?; @@ -903,7 +903,7 @@ impl UpkeepStore for PostgresStore { let mut query_builder = QueryBuilder::new( "SELECT id, activation, on_attempts_exceeded FROM inflight_taskactivations WHERE status = ", ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push_bind(InflightActivationStatus::Failure.to_string()); self.add_partition_condition(&mut query_builder, false); let failed_tasks = query_builder .build_query_as::<(String, Vec, i32)>() @@ -934,7 +934,7 @@ impl UpkeepStore for PostgresStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(ActivationStatus::Complete.to_string()) + .push_bind(InflightActivationStatus::Complete.to_string()) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -957,7 +957,7 @@ impl UpkeepStore for PostgresStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(ActivationStatus::Complete.to_string()) + .push_bind(InflightActivationStatus::Complete.to_string()) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -978,7 +978,7 @@ impl UpkeepStore for PostgresStore { let mut conn = self.acquire_write_conn_metric("remove_completed").await?; let mut query_builder = QueryBuilder::new("DELETE FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(ActivationStatus::Complete.to_string()); + query_builder.push_bind(InflightActivationStatus::Complete.to_string()); self.add_partition_condition(&mut query_builder, false); let result = query_builder.build().execute(&mut *conn).await?; diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index bcda5996..f414e54c 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -23,7 +23,7 @@ use sqlx::{ConnectOptions, FromRow, Pool, QueryBuilder, Row, Sqlite}; use std::{str::FromStr, time::Instant}; use crate::config::Config; -use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::activation::{InflightActivation, InflightActivationStatus}; use crate::store::traits::{ ClaimStore, CountStore, IngestStore, PullStore, PushStore, Store, TestStore, UpkeepStore, }; @@ -53,10 +53,10 @@ pub struct TableRow { pub bucket: i16, } -impl TryFrom for TableRow { +impl TryFrom for TableRow { type Error = anyhow::Error; - fn try_from(value: Activation) -> Result { + fn try_from(value: InflightActivation) -> Result { Ok(Self { id: value.id, activation: value.activation, @@ -81,12 +81,12 @@ impl TryFrom for TableRow { } } -impl From for Activation { +impl From for InflightActivation { fn from(value: TableRow) -> Self { Self { id: value.id, activation: value.activation, - status: ActivationStatus::from_str(&value.status).unwrap(), + status: InflightActivationStatus::from_str(&value.status).unwrap(), partition: value.partition, offset: value.offset, added_at: value.added_at, @@ -343,7 +343,7 @@ impl ClaimStore for SqliteActivationStore { limit: Option, bucket: Option, mark_processing: bool, - ) -> Result, Error> { + ) -> Result, Error> { let now = Utc::now(); let grace_period = self.config.processing_deadline_grace_sec; @@ -354,18 +354,18 @@ impl ClaimStore for SqliteActivationStore { "processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), claim_expires_at = NULL, status = " )); - query_builder.push_bind(ActivationStatus::Processing); + query_builder.push_bind(InflightActivationStatus::Processing); } else { query_builder.push(format!( "claim_expires_at = unixepoch('now', '+' || {:.3} || ' seconds', '+' || {grace_period} || ' seconds'), processing_deadline = NULL, status = ", self.config.claim_lease_ms as f64 / 1000.0, )); - query_builder.push_bind(ActivationStatus::Claimed); + query_builder.push_bind(InflightActivationStatus::Claimed); } query_builder.push(" WHERE id IN (SELECT id FROM inflight_taskactivations WHERE status = "); - query_builder.push_bind(ActivationStatus::Pending); + query_builder.push_bind(InflightActivationStatus::Pending); query_builder.push(" AND (expires_at IS NULL OR expires_at > "); query_builder.push_bind(now.timestamp()); query_builder.push(")"); @@ -423,9 +423,9 @@ impl PushStore for SqliteActivationStore { claim_expires_at = NULL WHERE id = $2 AND status = $3", )) - .bind(ActivationStatus::Processing) + .bind(InflightActivationStatus::Processing) .bind(id) - .bind(ActivationStatus::Claimed) + .bind(InflightActivationStatus::Claimed) .execute(&mut *conn) .await?; @@ -451,8 +451,8 @@ impl PullStore for SqliteActivationStore { async fn set_status( &self, id: &str, - status: ActivationStatus, - ) -> Result, Error> { + status: InflightActivationStatus, + ) -> Result, Error> { let mut conn = self.acquire_write_conn_metric("set_status").await?; let result: Option = sqlx::query_as( "UPDATE inflight_taskactivations SET status = $1 WHERE id = $2 RETURNING *", @@ -484,7 +484,7 @@ impl CountStore for SqliteActivationStore { } #[instrument(skip_all)] - async fn count_by_status(&self, status: ActivationStatus) -> Result { + async fn count_by_status(&self, status: InflightActivationStatus) -> Result { let result = sqlx::query("SELECT COUNT(*) as count FROM inflight_taskactivations WHERE status = $1") .bind(status) @@ -530,7 +530,7 @@ impl CountStore for SqliteActivationStore { LIMIT 1 ", ) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .fetch_one(&self.read_pool) .await; @@ -552,7 +552,7 @@ impl CountStore for SqliteActivationStore { #[async_trait] impl TestStore for SqliteActivationStore { - async fn get_by_id(&self, id: &str) -> Result, Error> { + async fn get_by_id(&self, id: &str) -> Result, Error> { let row_result: Option = sqlx::query_as( " SELECT id, @@ -633,7 +633,7 @@ impl IngestStore for SqliteActivationStore { } #[instrument(skip_all)] - async fn write(&self, batch: Vec) -> Result { + async fn write(&self, batch: Vec) -> Result { if batch.is_empty() { return Ok(0); } @@ -766,7 +766,7 @@ impl UpkeepStore for SqliteActivationStore { } #[instrument(skip_all)] - async fn get_retry_activations(&self) -> Result, Error> { + async fn get_retry_activations(&self) -> Result, Error> { Ok(sqlx::query_as( " SELECT id, @@ -792,7 +792,7 @@ impl UpkeepStore for SqliteActivationStore { WHERE status = $1 ", ) - .bind(ActivationStatus::Retry) + .bind(InflightActivationStatus::Retry) .fetch_all(&self.read_pool) .await? .into_iter() @@ -816,9 +816,9 @@ impl UpkeepStore for SqliteActivationStore { AND claim_expires_at < $2 AND status = $3", ) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .bind(now.timestamp()) - .bind(ActivationStatus::Claimed) + .bind(InflightActivationStatus::Claimed) .execute(&mut *conn) .await?; @@ -840,9 +840,9 @@ impl UpkeepStore for SqliteActivationStore { SET processing_deadline = null, status = $1 WHERE processing_deadline < $2 AND at_most_once = TRUE AND status = $3", ) - .bind(ActivationStatus::Failure) + .bind(InflightActivationStatus::Failure) .bind(now.timestamp()) - .bind(ActivationStatus::Processing) + .bind(InflightActivationStatus::Processing) .execute(&mut *atomic) .await; @@ -858,9 +858,9 @@ impl UpkeepStore for SqliteActivationStore { SET processing_deadline = null, status = $1, processing_attempts = processing_attempts + 1 WHERE processing_deadline < $2 AND status = $3", ) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .bind(now.timestamp()) - .bind(ActivationStatus::Processing) + .bind(InflightActivationStatus::Processing) .execute(&mut *atomic) .await; @@ -886,9 +886,9 @@ impl UpkeepStore for SqliteActivationStore { SET status = $1 WHERE processing_attempts >= $2 AND status = $3", ) - .bind(ActivationStatus::Failure) + .bind(InflightActivationStatus::Failure) .bind(self.config.max_processing_attempts as i32) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .execute(&mut *conn) .await; @@ -912,7 +912,7 @@ impl UpkeepStore for SqliteActivationStore { let query = sqlx::query( "DELETE FROM inflight_taskactivations WHERE status = $1 AND expires_at IS NOT NULL AND expires_at < $2", ) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .bind(now.timestamp()) .execute(&mut *conn) .await?; @@ -936,9 +936,9 @@ impl UpkeepStore for SqliteActivationStore { WHERE delay_until IS NOT NULL AND delay_until < $2 AND status = $3 "#, ) - .bind(ActivationStatus::Pending) + .bind(InflightActivationStatus::Pending) .bind(now.timestamp()) - .bind(ActivationStatus::Delay) + .bind(InflightActivationStatus::Delay) .execute(&mut *conn) .await?; @@ -957,7 +957,7 @@ impl UpkeepStore for SqliteActivationStore { let failed_tasks: Vec = sqlx::query("SELECT id, activation, on_attempts_exceeded FROM inflight_taskactivations WHERE status = $1") - .bind(ActivationStatus::Failure) + .bind(InflightActivationStatus::Failure) .fetch_all(&mut *atomic) .await? .into_iter() @@ -989,7 +989,7 @@ impl UpkeepStore for SqliteActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(ActivationStatus::Complete) + .push_bind(InflightActivationStatus::Complete) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -1012,7 +1012,7 @@ impl UpkeepStore for SqliteActivationStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations "); query_builder .push("SET status = ") - .push_bind(ActivationStatus::Complete) + .push_bind(InflightActivationStatus::Complete) .push(" WHERE id IN ("); let mut separated = query_builder.separated(", "); @@ -1032,7 +1032,7 @@ impl UpkeepStore for SqliteActivationStore { async fn remove_completed(&self) -> Result { let mut conn = self.acquire_write_conn_metric("remove_completed").await?; let query = sqlx::query("DELETE FROM inflight_taskactivations WHERE status = $1") - .bind(ActivationStatus::Complete) + .bind(InflightActivationStatus::Complete) .execute(&mut *conn) .await?; diff --git a/src/store/tests.rs b/src/store/tests.rs index 9aecf125..dcabb585 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -9,7 +9,7 @@ use tokio::{sync::broadcast, task::JoinSet}; use crate::store::traits::{IngestStore, UpkeepStore}; use crate::{ config::Config, - store::activation::{ActivationBuilder, ActivationStatus}, + store::activation::{InflightActivationBuilder, InflightActivationStatus}, store::adapters::{ postgres::PostgresActivationStoreConfig, sqlite::{InflightActivationStoreConfig, SqliteActivationStore, create_sqlite_pool}, @@ -25,41 +25,41 @@ use crate::{ #[test] fn test_inflightactivation_status_is_completion() { - let mut value = ActivationStatus::Unspecified; + let mut value = InflightActivationStatus::Unspecified; assert!(!value.is_conclusion()); - value = ActivationStatus::Pending; + value = InflightActivationStatus::Pending; assert!(!value.is_conclusion()); - value = ActivationStatus::Processing; + value = InflightActivationStatus::Processing; assert!(!value.is_conclusion()); - value = ActivationStatus::Retry; + value = InflightActivationStatus::Retry; assert!(value.is_conclusion()); - value = ActivationStatus::Failure; + value = InflightActivationStatus::Failure; assert!(value.is_conclusion()); - value = ActivationStatus::Complete; + value = InflightActivationStatus::Complete; assert!(value.is_conclusion()); } #[test] fn test_inflightactivation_status_from() { - let mut value: ActivationStatus = TaskActivationStatus::Pending.into(); - assert_eq!(value, ActivationStatus::Pending); + let mut value: InflightActivationStatus = TaskActivationStatus::Pending.into(); + assert_eq!(value, InflightActivationStatus::Pending); value = TaskActivationStatus::Processing.into(); - assert_eq!(value, ActivationStatus::Processing); + assert_eq!(value, InflightActivationStatus::Processing); value = TaskActivationStatus::Retry.into(); - assert_eq!(value, ActivationStatus::Retry); + assert_eq!(value, InflightActivationStatus::Retry); value = TaskActivationStatus::Failure.into(); - assert_eq!(value, ActivationStatus::Failure); + assert_eq!(value, InflightActivationStatus::Failure); value = TaskActivationStatus::Complete.into(); - assert_eq!(value, ActivationStatus::Complete); + assert_eq!(value, InflightActivationStatus::Complete); } #[tokio::test] @@ -106,15 +106,15 @@ async fn test_count_depths(#[case] adapter: &str) { // Check counts for an empty store let pending = store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(); let delay = store - .count_by_status(ActivationStatus::Delay) + .count_by_status(InflightActivationStatus::Delay) .await .unwrap(); let processing = store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(); @@ -129,28 +129,28 @@ async fn test_count_depths(#[case] adapter: &str) { assert!(store.write(batch).await.is_ok()); store - .set_status("id_0", ActivationStatus::Processing) + .set_status("id_0", InflightActivationStatus::Processing) .await .unwrap(); store - .set_status("id_1", ActivationStatus::Delay) + .set_status("id_1", InflightActivationStatus::Delay) .await .unwrap(); store - .set_status("id_2", ActivationStatus::Complete) + .set_status("id_2", InflightActivationStatus::Complete) .await .unwrap(); let pending = store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(); let delay = store - .count_by_status(ActivationStatus::Delay) + .count_by_status(InflightActivationStatus::Delay) .await .unwrap(); let processing = store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(); @@ -230,7 +230,7 @@ async fn test_get_pending_activation(#[case] adapter: &str) { .expect("expected one activation"); assert_eq!(result.id, "id_0"); - assert_eq!(result.status, ActivationStatus::Processing); + assert_eq!(result.status, InflightActivationStatus::Processing); assert_eq!(result.processing_deadline_duration, 10); assert!( result.processing_deadline.unwrap().timestamp() >= Utc::now().timestamp() + 13, @@ -358,7 +358,7 @@ async fn test_get_pending_activation_with_namespace(#[case] adapter: &str) { .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, ActivationStatus::Processing); + assert_eq!(result.status, InflightActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.namespace, "other_namespace"); store.remove_db().await.unwrap(); @@ -389,10 +389,10 @@ async fn test_get_pending_activation_from_multiple_namespaces(#[case] adapter: & assert_eq!(result.len(), 2); assert_eq!(result[1].id, "id_2"); assert_eq!(result[1].namespace, "ns3"); - assert_eq!(result[1].status, ActivationStatus::Claimed); + assert_eq!(result[1].status, InflightActivationStatus::Claimed); assert_eq!(result[0].id, "id_1"); assert_eq!(result[0].namespace, "ns2"); - assert_eq!(result[0].status, ActivationStatus::Claimed); + assert_eq!(result[0].status, InflightActivationStatus::Claimed); store.remove_db().await.unwrap(); } @@ -508,7 +508,7 @@ async fn test_get_pending_activation_fetches_application(#[case] adapter: &str) .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_0"); - assert_eq!(result.status, ActivationStatus::Processing); + assert_eq!(result.status, InflightActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); store.remove_db().await.unwrap(); @@ -532,7 +532,7 @@ async fn test_get_pending_activation_with_application(#[case] adapter: &str) { .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, ActivationStatus::Processing); + assert_eq!(result.status, InflightActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); @@ -575,7 +575,7 @@ async fn test_get_pending_activation_with_application_and_namespace(#[case] adap .unwrap() .expect("expected one activation"); assert_eq!(result.id, "id_1"); - assert_eq!(result.status, ActivationStatus::Processing); + assert_eq!(result.status, InflightActivationStatus::Processing); assert!(result.processing_deadline.unwrap() > Utc::now()); assert_eq!(result.application, "hammers"); assert_eq!(result.namespace, "target"); @@ -607,10 +607,13 @@ async fn test_get_pending_activations_no_limit(#[case] adapter: &str) { .await .unwrap(); assert_eq!(got.len(), N); - assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); + assert!( + got.iter() + .all(|a| a.status == InflightActivationStatus::Claimed) + ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 0 @@ -644,10 +647,13 @@ async fn test_get_pending_activations_limit_below_pending(#[case] adapter: &str) .await .unwrap(); assert_eq!(got.len(), X as usize); - assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); + assert!( + got.iter() + .all(|a| a.status == InflightActivationStatus::Claimed) + ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), N - X as usize @@ -681,10 +687,13 @@ async fn test_get_pending_activations_limit_above_pending(#[case] adapter: &str) .await .unwrap(); assert_eq!(got.len(), Y); - assert!(got.iter().all(|a| a.status == ActivationStatus::Claimed)); + assert!( + got.iter() + .all(|a| a.status == InflightActivationStatus::Claimed) + ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 0 @@ -709,12 +718,12 @@ async fn test_count_pending_activations(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(3); - batch[0].status = ActivationStatus::Processing; + batch[0].status = InflightActivationStatus::Processing; assert!(store.write(batch).await.is_ok()); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 2 @@ -751,7 +760,7 @@ async fn test_set_activation_status(#[case] adapter: &str) { assert!( store - .set_status("id_0", ActivationStatus::Failure) + .set_status("id_0", InflightActivationStatus::Failure) .await .is_ok() ); @@ -767,7 +776,7 @@ async fn test_set_activation_status(#[case] adapter: &str) { assert!( store - .set_status("id_0", ActivationStatus::Pending) + .set_status("id_0", InflightActivationStatus::Pending) .await .is_ok() ); @@ -781,13 +790,13 @@ async fn test_set_activation_status(#[case] adapter: &str) { .await; assert!( store - .set_status("id_0", ActivationStatus::Failure) + .set_status("id_0", InflightActivationStatus::Failure) .await .is_ok() ); assert!( store - .set_status("id_1", ActivationStatus::Failure) + .set_status("id_1", InflightActivationStatus::Failure) .await .is_ok() ); @@ -809,21 +818,23 @@ async fn test_set_activation_status(#[case] adapter: &str) { ); let result = store - .set_status("not_there", ActivationStatus::Complete) + .set_status("not_there", InflightActivationStatus::Complete) .await; assert!(result.is_ok(), "no query error"); let activation = result.unwrap(); assert!(activation.is_none(), "no activation found"); - let result = store.set_status("id_0", ActivationStatus::Complete).await; + let result = store + .set_status("id_0", InflightActivationStatus::Complete) + .await; assert!(result.is_ok(), "no query error"); let result_opt = result.unwrap(); assert!(result_opt.is_some(), "activation should be returned"); let inflight = result_opt.unwrap(); assert_eq!(inflight.id, "id_0"); - assert_eq!(inflight.status, ActivationStatus::Complete); + assert_eq!(inflight.status, InflightActivationStatus::Complete); store.remove_db().await.unwrap(); } @@ -847,7 +858,7 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { assert!( store - .set_status("id_0", ActivationStatus::Failure) + .set_status("id_0", InflightActivationStatus::Failure) .await .is_ok() ); @@ -862,7 +873,7 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { assert!( store - .set_status("id_0", ActivationStatus::Pending) + .set_status("id_0", InflightActivationStatus::Pending) .await .is_ok() ); @@ -876,13 +887,13 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { .await; assert!( store - .set_status("id_0", ActivationStatus::Failure) + .set_status("id_0", InflightActivationStatus::Failure) .await .is_ok() ); assert!( store - .set_status("id_1", ActivationStatus::Failure) + .set_status("id_1", InflightActivationStatus::Failure) .await .is_ok() ); @@ -906,21 +917,23 @@ async fn test_set_activation_status_with_partitions(#[case] adapter: &str) { ); let result = store - .set_status("not_there", ActivationStatus::Complete) + .set_status("not_there", InflightActivationStatus::Complete) .await; assert!(result.is_ok(), "no query error"); let activation = result.unwrap(); assert!(activation.is_none(), "no activation found"); - let result = store.set_status("id_0", ActivationStatus::Complete).await; + let result = store + .set_status("id_0", InflightActivationStatus::Complete) + .await; assert!(result.is_ok(), "no query error"); let result_opt = result.unwrap(); assert!(result_opt.is_some(), "activation should be returned"); let inflight = result_opt.unwrap(); assert_eq!(inflight.id, "id_0"); - assert_eq!(inflight.status, ActivationStatus::Complete); + assert_eq!(inflight.status, InflightActivationStatus::Complete); store.remove_db().await.unwrap(); } @@ -993,7 +1006,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { assert!( store - .set_status("id_0", ActivationStatus::Retry) + .set_status("id_0", InflightActivationStatus::Retry) .await .is_ok() ); @@ -1009,7 +1022,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { assert!( store - .set_status("id_1", ActivationStatus::Retry) + .set_status("id_1", InflightActivationStatus::Retry) .await .is_ok() ); @@ -1017,7 +1030,7 @@ async fn test_get_retry_activations(#[case] adapter: &str) { let retries = store.get_retry_activations().await.unwrap(); assert_eq!(retries.len(), 2); for record in retries.iter() { - assert_eq!(record.status, ActivationStatus::Retry); + assert_eq!(record.status, InflightActivationStatus::Retry); } store.remove_db().await.unwrap(); } @@ -1030,7 +1043,7 @@ async fn test_handle_processing_deadline(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); assert!(store.write(batch.clone()).await.is_ok()); @@ -1074,9 +1087,9 @@ async fn test_handle_processing_deadline_multiple_tasks(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[0].status = ActivationStatus::Processing; + batch[0].status = InflightActivationStatus::Processing; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); - batch[1].status = ActivationStatus::Claimed; + batch[1].status = InflightActivationStatus::Claimed; batch[1].processing_deadline = Some(Utc::now() + chrono::Duration::days(30)); assert!(store.write(batch).await.is_ok()); assert_counts( @@ -1113,10 +1126,10 @@ async fn test_handle_processing_at_most_once(#[case] adapter: &str) { // Both records are past processing deadlines let mut batch = make_activations(2); - batch[0].status = ActivationStatus::Processing; + batch[0].status = InflightActivationStatus::Processing; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; replace_retry_state( &mut batch[1], @@ -1155,7 +1168,7 @@ async fn test_handle_processing_at_most_once(#[case] adapter: &str) { .await; let task = store.get_by_id(&batch[1].id).await.unwrap().unwrap(); - assert_eq!(task.status, ActivationStatus::Failure); + assert_eq!(task.status, InflightActivationStatus::Failure); store.remove_db().await.unwrap(); } @@ -1167,7 +1180,7 @@ async fn test_handle_processing_deadline_discard_after(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1213,7 +1226,7 @@ async fn test_handle_processing_deadline_deadletter_after(#[case] adapter: &str) let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1259,7 +1272,7 @@ async fn test_handle_processing_deadline_no_retries_remaining(#[case] adapter: & let store = create_test_store(adapter).await; let mut batch = make_activations(2); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); replace_retry_state( &mut batch[1], @@ -1304,13 +1317,13 @@ async fn test_handle_processing_deadline_no_retries_remaining(#[case] adapter: & async fn test_handle_claim_expiration_unsent_no_attempt_increment(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(1); - batch[0].status = ActivationStatus::Claimed; + batch[0].status = InflightActivationStatus::Claimed; batch[0].claim_expires_at = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); assert!(store.write(batch.clone()).await.is_ok()); let count = store.handle_claim_expiration().await.unwrap(); assert_eq!(count, 1); let task = store.get_by_id(&batch[0].id).await.unwrap().unwrap(); - assert_eq!(task.status, ActivationStatus::Pending); + assert_eq!(task.status, InflightActivationStatus::Pending); assert_eq!(task.processing_attempts, 0); store.remove_db().await.unwrap(); } @@ -1322,14 +1335,14 @@ async fn test_handle_claim_expiration_unsent_no_attempt_increment(#[case] adapte async fn test_handle_claim_expiration_at_most_once_reverts_to_pending(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(1); - batch[0].status = ActivationStatus::Claimed; + batch[0].status = InflightActivationStatus::Claimed; batch[0].at_most_once = true; batch[0].claim_expires_at = Some(Utc.with_ymd_and_hms(2020, 1, 1, 1, 1, 1).unwrap()); assert!(store.write(batch.clone()).await.is_ok()); let count = store.handle_claim_expiration().await.unwrap(); assert_eq!(count, 1); let task = store.get_by_id(&batch[0].id).await.unwrap().unwrap(); - assert_eq!(task.status, ActivationStatus::Pending); + assert_eq!(task.status, InflightActivationStatus::Pending); assert_eq!(task.processing_attempts, 0); store.remove_db().await.unwrap(); } @@ -1343,14 +1356,14 @@ async fn test_processing_attempts_exceeded(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut batch = make_activations(3); - batch[0].status = ActivationStatus::Pending; + batch[0].status = InflightActivationStatus::Pending; batch[0].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[0].processing_attempts = config.max_processing_attempts as i32; - batch[1].status = ActivationStatus::Complete; + batch[1].status = InflightActivationStatus::Complete; batch[1].added_at += Duration::from_secs(1); - batch[2].status = ActivationStatus::Pending; + batch[2].status = InflightActivationStatus::Pending; batch[2].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[2].processing_attempts = config.max_processing_attempts as i32; @@ -1388,10 +1401,10 @@ async fn test_remove_completed(#[case] adapter: &str) { let store = create_test_store(adapter).await; let mut records = make_activations(3); - records[0].status = ActivationStatus::Complete; - records[1].status = ActivationStatus::Pending; + records[0].status = InflightActivationStatus::Complete; + records[1].status = InflightActivationStatus::Pending; records[1].added_at += Duration::from_secs(1); - records[2].status = ActivationStatus::Complete; + records[2].status = InflightActivationStatus::Complete; records[2].added_at += Duration::from_secs(2); assert!(store.write(records.clone()).await.is_ok()); @@ -1450,14 +1463,14 @@ async fn test_remove_completed_multiple_gaps(#[case] adapter: &str) { let mut records = make_activations(4); // only record 1 can be removed - records[0].status = ActivationStatus::Complete; - records[1].status = ActivationStatus::Failure; + records[0].status = InflightActivationStatus::Complete; + records[1].status = InflightActivationStatus::Failure; records[1].added_at += Duration::from_secs(1); - records[2].status = ActivationStatus::Complete; + records[2].status = InflightActivationStatus::Complete; records[2].added_at += Duration::from_secs(2); - records[3].status = ActivationStatus::Processing; + records[3].status = InflightActivationStatus::Processing; records[3].added_at += Duration::from_secs(3); assert!(store.write(records.clone()).await.is_ok()); @@ -1525,7 +1538,7 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { let mut records = make_activations(4); // deadletter - records[0].status = ActivationStatus::Failure; + records[0].status = InflightActivationStatus::Failure; replace_retry_state( &mut records[0], Some(RetryState { @@ -1537,7 +1550,7 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { }), ); // discard - records[1].status = ActivationStatus::Failure; + records[1].status = InflightActivationStatus::Failure; replace_retry_state( &mut records[1], Some(RetryState { @@ -1549,11 +1562,11 @@ async fn test_handle_failed_tasks(#[case] adapter: &str) { }), ); // no retry state = discard - records[2].status = ActivationStatus::Failure; + records[2].status = InflightActivationStatus::Failure; replace_retry_state(&mut records[2], None); // Another deadletter - records[3].status = ActivationStatus::Failure; + records[3].status = InflightActivationStatus::Failure; replace_retry_state( &mut records[3], Some(RetryState { @@ -1738,7 +1751,7 @@ async fn test_clear(#[case] adapter: &str) { let namespace = generate_unique_namespace(); let batch = vec![ - ActivationBuilder::new() + InflightActivationBuilder::new() .id("id_0") .taskname("taskname") .namespace(&namespace) @@ -1845,7 +1858,7 @@ async fn test_pending_activation_max_lag_no_pending(#[case] adapter: &str) { assert_eq!(0.0, store.pending_activation_max_lag(&now).await); let mut processing = make_activations(1); - processing[0].status = ActivationStatus::Processing; + processing[0].status = InflightActivationStatus::Processing; assert!(store.write(processing).await.is_ok()); // No pending activations, max lag is 0 diff --git a/src/store/traits.rs b/src/store/traits.rs index e84a879a..314645ba 100644 --- a/src/store/traits.rs +++ b/src/store/traits.rs @@ -4,14 +4,14 @@ use chrono::{DateTime, Utc}; use tokio::join; use tracing::warn; -use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::activation::{InflightActivation, InflightActivationStatus}; use crate::store::types::{BucketRange, DepthCounts, FailedTasksForwarder}; /// This trait only contains methods used by the consumer component. #[async_trait] pub trait IngestStore: Send + Sync { /// Write a batch of activations to the store. - async fn write(&self, batch: Vec) -> Result; + async fn write(&self, batch: Vec) -> Result; /// Change the Kafka partitions for which this instance is responsible. async fn assign_partitions(&self, partitions: Vec) -> Result<(), Error>; @@ -27,16 +27,16 @@ pub trait CountStore: Send + Sync { async fn count(&self) -> Result; /// Count activations by status. - async fn count_by_status(&self, status: ActivationStatus) -> Result; + async fn count_by_status(&self, status: InflightActivationStatus) -> Result; /// Queue depths for pending, delay, and processing (writer backpressure and upkeep gauges). /// Default implementation uses separate calls, but stores may override with a single query. async fn count_depths(&self) -> Result { let (pending, delay, claimed, processing) = join!( - self.count_by_status(ActivationStatus::Pending), - self.count_by_status(ActivationStatus::Delay), - self.count_by_status(ActivationStatus::Claimed), - self.count_by_status(ActivationStatus::Processing), + self.count_by_status(InflightActivationStatus::Pending), + self.count_by_status(InflightActivationStatus::Delay), + self.count_by_status(InflightActivationStatus::Claimed), + self.count_by_status(InflightActivationStatus::Processing), ); Ok(DepthCounts { @@ -64,7 +64,7 @@ pub trait ClaimStore: Send + Sync { limit: Option, bucket: Option, mark_processing: bool, - ) -> Result, Error>; + ) -> Result, Error>; /// Claims `limit` activations and updates status to `Claimed` until `mark_processing` moves them to `Processing`. async fn claim_activations_for_push( @@ -73,7 +73,7 @@ pub trait ClaimStore: Send + Sync { namespaces: Option<&[String]>, limit: Option, bucket: Option, - ) -> Result, Error> { + ) -> Result, Error> { // If a namespace filter is used, an application must also be used if namespaces.is_some() && application.is_none() { warn!( @@ -93,7 +93,7 @@ pub trait ClaimStore: Send + Sync { &self, application: Option<&str>, namespace: Option<&str>, - ) -> Result, Error> { + ) -> Result, Error> { // Convert single namespace to vector for internal use let namespaces = namespace.map(|ns| vec![ns.to_string()]); @@ -134,15 +134,15 @@ pub trait PullStore: ClaimStore { async fn set_status( &self, id: &str, - status: ActivationStatus, - ) -> Result, Error>; + status: InflightActivationStatus, + ) -> Result, Error>; } /// This trait only contains methods used by upkeep. #[async_trait] pub trait UpkeepStore { /// Get all activations with status Retry - async fn get_retry_activations(&self) -> Result, Error>; + async fn get_retry_activations(&self) -> Result, Error>; /// Revert expired push claims back to pending status. async fn handle_claim_expiration(&self) -> Result; @@ -188,7 +188,7 @@ pub trait Store: #[async_trait] pub trait TestStore: Store { /// Get an activation by id - async fn get_by_id(&self, id: &str) -> Result, Error>; + async fn get_by_id(&self, id: &str) -> Result, Error>; /// Set the processing deadline for a specific activation async fn set_processing_deadline( diff --git a/src/test_utils.rs b/src/test_utils.rs index 9dda1c54..4f6eab90 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -15,7 +15,7 @@ use uuid::Uuid; use crate::{ config::Config, store::{ - activation::{Activation, ActivationBuilder, ActivationStatus}, + activation::{InflightActivation, InflightActivationBuilder, InflightActivationStatus}, adapters::{ postgres::{PostgresActivationStoreConfig, PostgresStore}, sqlite::{InflightActivationStoreConfig, SqliteActivationStore}, @@ -134,12 +134,12 @@ impl Default for TaskActivationBuilder { } } -impl ActivationBuilder { +impl InflightActivationBuilder { pub fn new() -> Self { Self::default() } - pub fn build(mut self, builder: TaskActivationBuilder) -> Activation { + pub fn build(mut self, builder: TaskActivationBuilder) -> InflightActivation { // Grab required fields let id = self.id.as_ref().expect("field 'id' is required"); @@ -230,13 +230,13 @@ pub fn generate_unique_namespace() -> String { } /// Create a collection of `count` pending unsaved activations in a particular `namespace`. If you do not want to provide a namespace, use `make_activations`. -pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec { - let mut records: Vec = vec![]; +pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec { + let mut records: Vec = vec![]; for i in 0..count { let now = Utc::now(); - let item = ActivationBuilder::new() + let item = InflightActivationBuilder::new() .id(format!("id_{i}")) .taskname("taskname") .namespace(&namespace) @@ -252,7 +252,7 @@ pub fn make_activations_with_namespace(namespace: String, count: u32) -> Vec Vec { +pub fn make_activations(count: u32) -> Vec { let namespace = generate_unique_namespace(); make_activations_with_namespace(namespace, count) } @@ -426,7 +426,7 @@ pub async fn consume_topic( results } -pub fn replace_retry_state(inflight: &mut Activation, retry: Option) { +pub fn replace_retry_state(inflight: &mut InflightActivation, retry: Option) { let mut activation = TaskActivation::decode(&inflight.activation as &[u8]).unwrap(); activation.retry_state = retry; inflight.activation = activation.encode_to_vec(); @@ -455,7 +455,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.pending, store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), "difference in pending count", @@ -463,7 +463,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.claimed, store - .count_by_status(ActivationStatus::Claimed) + .count_by_status(InflightActivationStatus::Claimed) .await .unwrap(), "difference in claimed count", @@ -471,7 +471,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.processing, store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(), "difference in processing count", @@ -479,7 +479,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.retry, store - .count_by_status(ActivationStatus::Retry) + .count_by_status(InflightActivationStatus::Retry) .await .unwrap(), "difference in retry count", @@ -487,7 +487,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.delayed, store - .count_by_status(ActivationStatus::Delay) + .count_by_status(InflightActivationStatus::Delay) .await .unwrap(), "difference in delay count", @@ -495,7 +495,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.complete, store - .count_by_status(ActivationStatus::Complete) + .count_by_status(InflightActivationStatus::Complete) .await .unwrap(), "difference in complete count", @@ -503,7 +503,7 @@ pub async fn assert_counts(expected: StatusCount, store: &dyn Store) { assert_eq!( expected.failure, store - .count_by_status(ActivationStatus::Failure) + .count_by_status(InflightActivationStatus::Failure) .await .unwrap(), "difference in failure count", diff --git a/src/upkeep.rs b/src/upkeep.rs index 911a129a..b3e04294 100644 --- a/src/upkeep.rs +++ b/src/upkeep.rs @@ -536,7 +536,7 @@ mod tests { use crate::{ config::Config, runtime_config::RuntimeConfigManager, - store::activation::ActivationStatus, + store::activation::InflightActivationStatus, test_utils::{ StatusCount, assert_counts, consume_topic, create_config, create_integration_config, create_integration_config_with_topic, create_producer, create_test_store, @@ -664,7 +664,7 @@ mod tests { .expect(""); activation.parameters = r#"{"a":"b"}"#.into(); activation.delay = Some(30); - records[0].status = ActivationStatus::Retry; + records[0].status = InflightActivationStatus::Retry; records[0].delay_until = Some(Utc::now() + Duration::from_secs(30)); records[0].activation = activation.encode_to_vec(); @@ -724,7 +724,7 @@ mod tests { let mut batch = make_activations(2); // Make a task with a future processing deadline - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc::now() + TimeDelta::minutes(5)); assert!(store.store(batch.clone()).await.is_ok()); @@ -741,7 +741,7 @@ mod tests { // Should retain the processing record assert_eq!( store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(), 1 @@ -760,7 +760,7 @@ mod tests { let mut batch = make_activations(2); // Make a task past with a processing deadline in the past - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); assert!(store.store(batch.clone()).await.is_ok()); @@ -768,7 +768,7 @@ mod tests { // Should start off with one in processing assert_eq!( store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(), 1 @@ -816,7 +816,7 @@ mod tests { let mut batch = make_activations(2); // Make a task past with a processing deadline in the past - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[1].processing_attempts = 0; @@ -825,14 +825,14 @@ mod tests { // Should start off with one in processing assert_eq!( store - .count_by_status(ActivationStatus::Processing) + .count_by_status(InflightActivationStatus::Processing) .await .unwrap(), 1 ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 1 @@ -885,7 +885,7 @@ mod tests { delay_on_retry: None, }), ); - batch[1].status = ActivationStatus::Processing; + batch[1].status = InflightActivationStatus::Processing; batch[1].processing_deadline = Some(Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap()); batch[1].at_most_once = true; @@ -930,7 +930,7 @@ mod tests { // Because 1 is complete and has a higher offset than 0, index 2 can be discarded batch[0].processing_attempts = config.max_processing_attempts as i32; - batch[1].status = ActivationStatus::Complete; + batch[1].status = InflightActivationStatus::Complete; batch[1].added_at += Duration::from_secs(1); batch[2].processing_attempts = config.max_processing_attempts as i32; @@ -952,7 +952,7 @@ mod tests { assert_eq!(result_context.completed, 3); // all three are removed as completed assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 0, @@ -960,7 +960,7 @@ mod tests { ); assert_eq!( store - .count_by_status(ActivationStatus::Complete) + .count_by_status(InflightActivationStatus::Complete) .await .unwrap(), 0, @@ -992,7 +992,7 @@ mod tests { delay_on_retry: None, }), ); - records[0].status = ActivationStatus::Failure; + records[0].status = InflightActivationStatus::Failure; records[1].added_at += Duration::from_secs(1); assert!(store.store(records.clone()).await.is_ok()); @@ -1035,7 +1035,7 @@ mod tests { let mut last_vacuum = Instant::now(); let mut batch = make_activations(2); - batch[0].status = ActivationStatus::Failure; + batch[0].status = InflightActivationStatus::Failure; batch[1].added_at += Duration::from_secs(1); assert!(store.store(batch).await.is_ok()); @@ -1058,7 +1058,7 @@ mod tests { ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 1, @@ -1081,7 +1081,7 @@ mod tests { let mut batch = make_activations(4); batch[0].expires_at = Some(Utc::now() - Duration::from_secs(100)); - batch[1].status = ActivationStatus::Complete; + batch[1].status = InflightActivationStatus::Complete; batch[2].expires_at = Some(Utc::now() - Duration::from_secs(100)); // Ensure the fourth task is in the future @@ -1103,7 +1103,7 @@ mod tests { assert_eq!(result_context.completed, 1); // 1 complete assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 1, @@ -1111,7 +1111,7 @@ mod tests { ); assert_eq!( store - .count_by_status(ActivationStatus::Complete) + .count_by_status(InflightActivationStatus::Complete) .await .unwrap(), 0, @@ -1150,23 +1150,23 @@ mod tests { let mut batch = make_activations(2); - batch[0].status = ActivationStatus::Delay; + batch[0].status = InflightActivationStatus::Delay; batch[0].delay_until = Some(Utc::now() - Duration::from_secs(1)); - batch[1].status = ActivationStatus::Delay; + batch[1].status = InflightActivationStatus::Delay; batch[1].delay_until = Some(Utc::now() + Duration::from_secs(1)); assert!(store.store(batch.clone()).await.is_ok()); assert_eq!( store - .count_by_status(ActivationStatus::Delay) + .count_by_status(InflightActivationStatus::Delay) .await .unwrap(), 2 ); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 0 @@ -1183,7 +1183,7 @@ mod tests { assert_eq!(result_context.delay_elapsed, 1); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 1 @@ -1215,7 +1215,7 @@ mod tests { assert_eq!(result_context.delay_elapsed, 1); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 1 @@ -1269,7 +1269,7 @@ demoted_namespaces: assert_eq!(result_context.forwarded, 2); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 4, @@ -1277,7 +1277,7 @@ demoted_namespaces: ); assert_eq!( store - .count_by_status(ActivationStatus::Complete) + .count_by_status(InflightActivationStatus::Complete) .await .unwrap(), 2, @@ -1328,7 +1328,7 @@ demoted_namespaces: assert_eq!(result_context.killswitched, 3); assert_eq!( store - .count_by_status(ActivationStatus::Pending) + .count_by_status(InflightActivationStatus::Pending) .await .unwrap(), 3