use std::borrow::Cow;
use std::result::Result as StdResult;
use std::time::Duration;
use crate::connection::Connection;
use crate::error::{Error, Result};
use crate::instance::DBOS;
use crate::sysdb::types::{
AppScope, Change, NewQueue, OnExistingQueue, QueueRecord, QueueUpdate, RateLimit,
has_partition_limit,
};
use crate::sysdb::{Error as SysdbError, INTERNAL_QUEUE};
pub(crate) const DEFAULT_POLLING_INTERVAL: Duration = Duration::from_secs(1);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Queue {
name: String,
concurrency: Option<i32>,
worker_concurrency: Option<i32>,
rate_limit: Option<RateLimit>,
partition_concurrency: Option<i32>,
partition_worker_concurrency: Option<i32>,
partition_rate_limit: Option<RateLimit>,
polling_interval: Duration,
}
impl Queue {
fn from_record(record: QueueRecord) -> Self {
Self {
name: record.name,
concurrency: record.concurrency,
worker_concurrency: record.worker_concurrency,
rate_limit: record.rate_limit,
partition_concurrency: record.partition_concurrency,
partition_worker_concurrency: record.partition_worker_concurrency,
partition_rate_limit: record.partition_rate_limit,
polling_interval: record.polling_interval,
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn concurrency(&self) -> Option<i32> {
self.concurrency
}
pub fn worker_concurrency(&self) -> Option<i32> {
self.worker_concurrency
}
pub fn rate_limit(&self) -> Option<RateLimit> {
self.rate_limit
}
pub fn is_partitioned(&self) -> bool {
has_partition_limit(
self.partition_concurrency,
self.partition_worker_concurrency,
self.partition_rate_limit,
)
}
pub fn partition_concurrency(&self) -> Option<i32> {
self.partition_concurrency
}
pub fn partition_worker_concurrency(&self) -> Option<i32> {
self.partition_worker_concurrency
}
pub fn partition_rate_limit(&self) -> Option<RateLimit> {
self.partition_rate_limit
}
pub fn polling_interval(&self) -> Duration {
self.polling_interval
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueueOptions {
pub concurrency: Option<i32>,
pub worker_concurrency: Option<i32>,
pub polling_interval: Duration,
pub rate_limit: Option<RateLimit>,
pub partition_concurrency: Option<i32>,
pub partition_worker_concurrency: Option<i32>,
pub partition_rate_limit: Option<RateLimit>,
}
impl Default for QueueOptions {
fn default() -> Self {
Self {
concurrency: None,
worker_concurrency: None,
polling_interval: DEFAULT_POLLING_INTERVAL,
rate_limit: None,
partition_concurrency: None,
partition_worker_concurrency: None,
partition_rate_limit: None,
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct QueueChange {
pub concurrency: Change<Option<i32>>,
pub worker_concurrency: Change<Option<i32>>,
pub polling_interval: Change<Duration>,
pub rate_limit: Change<Option<RateLimit>>,
pub partition_concurrency: Change<Option<i32>>,
pub partition_worker_concurrency: Change<Option<i32>>,
pub partition_rate_limit: Change<Option<RateLimit>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QueueConflict {
UpdateIfLatestVersion,
AlwaysUpdate,
NeverUpdate,
}
fn validate(name: &str, options: &QueueOptions) -> Result<()> {
validate_fields(options)
.map_err(|(_, detail)| Error::Config(format!("queue `{name}`: {detail}")))
}
fn validate_fields(options: &QueueOptions) -> StdResult<(), (Cow<'static, str>, String)> {
let refuse = |field: &'static str, message: String| Err((Cow::Borrowed(field), message));
if let Some(global) = options.concurrency
&& global < 0
{
return refuse(
"concurrency",
format!("`concurrency` must not be negative, got {global}"),
);
}
if let Some(worker) = options.worker_concurrency
&& worker < 0
{
return refuse(
"worker_concurrency",
format!("`worker_concurrency` must not be negative, got {worker}"),
);
}
if let (Some(worker), Some(global)) = (options.worker_concurrency, options.concurrency)
&& worker > global
{
return refuse(
"worker_concurrency",
format!(
"`concurrency` must be greater than or equal to `worker_concurrency`, \
got {global} and {worker}"
),
);
}
if options.polling_interval.is_zero() {
return refuse(
"polling_interval",
"`polling_interval` cannot be zero".to_owned(),
);
}
if let Some(rate_limit) = options.rate_limit {
if rate_limit.limit < 1 {
return refuse(
"rate_limit.limit",
format!(
"`rate_limit.limit` must be at least 1, got {}",
rate_limit.limit
),
);
}
if rate_limit.period.is_zero() {
return refuse(
"rate_limit.period",
"`rate_limit.period` cannot be zero".to_owned(),
);
}
}
if let Some(partition) = options.partition_concurrency
&& partition < 1
{
return refuse(
"partition_concurrency",
format!("`partition_concurrency` must be at least 1, got {partition}"),
);
}
if let Some(partition) = options.partition_worker_concurrency
&& partition < 1
{
return refuse(
"partition_worker_concurrency",
format!("`partition_worker_concurrency` must be at least 1, got {partition}"),
);
}
if let Some(rate_limit) = options.partition_rate_limit {
if rate_limit.limit < 1 {
return refuse(
"partition_rate_limit.limit",
format!(
"`partition_rate_limit.limit` must be at least 1, got {}",
rate_limit.limit
),
);
}
if rate_limit.period.is_zero() {
return refuse(
"partition_rate_limit.period",
"`partition_rate_limit.period` cannot be zero".to_owned(),
);
}
}
if let (Some(partition), Some(worker)) = (
options.partition_worker_concurrency,
options.worker_concurrency,
) && partition > worker
{
return refuse(
"partition_worker_concurrency",
format!(
"`worker_concurrency` must be greater than or equal to \
`partition_worker_concurrency`, got {worker} and {partition}"
),
);
}
if let (Some(partition), Some(concurrency)) = (
options.partition_worker_concurrency,
options.partition_concurrency,
) && partition > concurrency
{
return refuse(
"partition_worker_concurrency",
format!(
"`partition_concurrency` must be greater than or equal to \
`partition_worker_concurrency`, got {concurrency} and {partition}"
),
);
}
if let (Some(partition), Some(global)) = (options.partition_concurrency, options.concurrency)
&& partition > global
{
return refuse(
"partition_concurrency",
format!(
"`concurrency` must be greater than or equal to `partition_concurrency`, \
got {global} and {partition}"
),
);
}
if let (Some(partition), Some(global)) =
(options.partition_worker_concurrency, options.concurrency)
&& partition > global
{
return refuse(
"partition_worker_concurrency",
format!(
"`concurrency` must be greater than or equal to \
`partition_worker_concurrency`, got {global} and {partition}"
),
);
}
Ok(())
}
impl DBOS {
pub async fn register_queue(
&self,
name: &str,
options: QueueOptions,
on_conflict: QueueConflict,
) -> Result<Queue> {
let executor = self.executor("register a queue")?;
executor
.connection()
.register_queue(name, options, on_conflict, Some(executor.app_version()))
.await
}
pub async fn queue(&self, name: &str) -> Result<Option<Queue>> {
self.executor("read a queue")?
.connection()
.queue(name)
.await
}
pub async fn list_queues(&self) -> Result<Vec<Queue>> {
self.executor("list queues")?
.connection()
.list_queues()
.await
}
pub async fn update_queue(&self, name: &str, change: QueueChange) -> Result<Queue> {
self.executor("update a queue")?
.connection()
.update_queue(name, change)
.await
}
pub async fn delete_queue(&self, name: &str) -> Result<()> {
self.executor("delete a queue")?
.connection()
.delete_queue(name)
.await
}
}
fn warn_partitioned(name: &str) {
tracing::warn!(
queue = name,
"the queue is now partitioned; any workflow already enqueued on it without a partition key \
will never be dequeued, so drain it before partitioning it or re-enqueue such workflows \
with a key"
);
}
impl Connection {
pub(crate) async fn register_queue(
&self,
name: &str,
options: QueueOptions,
on_conflict: QueueConflict,
app_version: Option<&str>,
) -> Result<Queue> {
if name == INTERNAL_QUEUE {
return Err(Error::Config(format!(
"the queue name `{name}` is reserved for the engine's internal queue"
)));
}
validate(name, &options)?;
let on_existing = match on_conflict {
QueueConflict::AlwaysUpdate => OnExistingQueue::Update,
QueueConflict::NeverUpdate => OnExistingQueue::Leave,
QueueConflict::UpdateIfLatestVersion => {
let Some(version) = app_version else {
return Err(Error::Config(format!(
"registering queue `{name}`: `QueueConflict::UpdateIfLatestVersion` needs \
an application version to compare against, and a client has none; ask \
for `AlwaysUpdate` or `NeverUpdate`"
)));
};
let latest = self.latest_app_version().await?;
match latest {
None => OnExistingQueue::Update,
Some(latest) if latest.version_name == version => OnExistingQueue::Update,
Some(latest) => {
tracing::debug!(
queue = name,
version,
latest = latest.version_name,
"an older version registered this queue; its stored limits stand"
);
OnExistingQueue::Leave
}
}
}
};
let upserted = self
.sysdb()
.upsert_queue(
&NewQueue {
concurrency: options.concurrency,
worker_concurrency: options.worker_concurrency,
polling_interval: options.polling_interval,
rate_limit: options.rate_limit,
partition_concurrency: options.partition_concurrency,
partition_worker_concurrency: options.partition_worker_concurrency,
partition_rate_limit: options.partition_rate_limit,
application_name: self.app_name(),
..NewQueue::new(name)
},
on_existing,
)
.await
.map_err(Error::SystemDatabase)?;
match &upserted.before {
None => tracing::info!(queue = name, "registered a queue"),
Some(before) if !before.is_partitioned() && upserted.after.is_partitioned() => {
warn_partitioned(name);
}
Some(_) => {}
}
Ok(Queue::from_record(upserted.after))
}
pub(crate) async fn queue(&self, name: &str) -> Result<Option<Queue>> {
Ok(self
.sysdb()
.get_queue(name)
.await
.map_err(Error::SystemDatabase)?
.map(Queue::from_record))
}
pub(crate) async fn list_queues(&self) -> Result<Vec<Queue>> {
Ok(self
.sysdb()
.list_queues(&AppScope::Unset)
.await
.map_err(Error::SystemDatabase)?
.into_iter()
.filter(|queue| queue.name != INTERNAL_QUEUE)
.map(Queue::from_record)
.collect())
}
pub(crate) async fn update_queue(&self, name: &str, change: QueueChange) -> Result<Queue> {
if name == INTERNAL_QUEUE {
return Err(Error::Config(format!(
"the queue name `{name}` is reserved for the engine's internal queue"
)));
}
let validate = |_stored: &QueueRecord, merged: &QueueRecord| -> StdResult<(), SysdbError> {
let options = QueueOptions {
concurrency: merged.concurrency,
worker_concurrency: merged.worker_concurrency,
polling_interval: merged.polling_interval,
rate_limit: merged.rate_limit,
partition_concurrency: merged.partition_concurrency,
partition_worker_concurrency: merged.partition_worker_concurrency,
partition_rate_limit: merged.partition_rate_limit,
};
validate_fields(&options)
.map_err(|(field, detail)| SysdbError::InvalidInput { field, detail })
};
let update = QueueUpdate {
concurrency: change.concurrency,
worker_concurrency: change.worker_concurrency,
polling_interval: change.polling_interval,
rate_limit: change.rate_limit,
partition_concurrency: change.partition_concurrency,
partition_worker_concurrency: change.partition_worker_concurrency,
partition_rate_limit: change.partition_rate_limit,
};
let updated = self
.sysdb()
.update_queue(name, &update, &validate)
.await
.map_err(|error| match error {
SysdbError::InvalidInput { detail, .. } => {
Error::Config(format!("queue `{name}`: {detail}"))
}
SysdbError::NotRegistered { .. } => {
Error::Config(format!("no queue named `{name}` is registered"))
}
other => Error::SystemDatabase(other),
})?;
if !updated.before.is_partitioned() && updated.after.is_partitioned() {
warn_partitioned(name);
} else {
tracing::info!(queue = name, "updated the queue's limits");
}
Ok(Queue::from_record(updated.after))
}
pub(crate) async fn delete_queue(&self, name: &str) -> Result<()> {
if name == INTERNAL_QUEUE {
return Err(Error::Config(format!(
"the queue name `{name}` is reserved for the engine's internal queue"
)));
}
self.sysdb()
.delete_queue(name)
.await
.map_err(Error::SystemDatabase)?;
tracing::info!(queue = name, "deleted the queue");
Ok(())
}
}