use std::{
collections::VecDeque,
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll, Waker},
};
use futures_sink::Sink;
use tower_layer::Identity;
use crate::{
backend::{Backend, BackendConfig, finalize::Ephemeral},
error::BoxDynError,
task::{
Task,
task_id::{RandomId, TaskId},
},
worker::context::WorkerContext,
};
#[derive(Debug, Clone)]
pub struct VecDequeBackend<T> {
queue: Arc<Mutex<VecDeque<Task<T>>>>,
waker: Arc<Mutex<Option<Waker>>>,
}
impl<T> Default for VecDequeBackend<T> {
fn default() -> Self {
Self {
queue: Arc::new(Mutex::new(VecDeque::new())),
waker: Arc::new(Mutex::new(None)),
}
}
}
impl<T> VecDequeBackend<T> {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
queue: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
waker: Arc::new(Mutex::new(None)),
}
}
}
#[derive(Debug, thiserror::Error, Clone)]
#[non_exhaustive]
pub enum VecDequeError {
#[error("Polling error: {0}")]
PollError(Arc<BoxDynError>),
#[error("Sending error: {0}")]
SendError(Arc<BoxDynError>),
}
impl<T> Backend for VecDequeBackend<T> {
type Task = Task<T>;
type Error = VecDequeError;
fn poll_ready(
&mut self,
cx: &mut Context<'_>,
_worker: &WorkerContext,
) -> Poll<Result<(), Self::Error>> {
if self
.queue
.lock()
.map_err(|e| VecDequeError::PollError(Arc::new(e.to_string().into())))?
.is_empty()
{
*self.waker.lock().unwrap() = Some(cx.waker().clone());
Poll::Pending
} else {
Poll::Ready(Ok(()))
}
}
fn poll_next(
&mut self,
_cx: &mut Context<'_>,
_worker: &WorkerContext,
) -> Poll<Option<Result<Self::Task, Self::Error>>> {
match self
.queue
.lock()
.map_err(|e| VecDequeError::PollError(Arc::new(e.to_string().into())))?
.pop_front()
{
Some(task) => Poll::Ready(Some(Ok(task))),
None => Poll::Ready(None),
}
}
fn poll_close(
&mut self,
_: &mut Context<'_>,
_: &WorkerContext,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
impl<T> Sink<Task<T>> for VecDequeBackend<T>
where
T: Send + Unpin + 'static,
{
type Error = VecDequeError;
fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, mut item: Task<T>) -> Result<(), Self::Error> {
let this = self.get_mut();
let mut tasks = this
.queue
.lock()
.map_err(|e| VecDequeError::SendError(Arc::new(e.to_string().into())))?;
if let Some(ref key) = item.idempotency_key() {
let exists = tasks.iter().any(|task| {
task.idempotency_key()
.as_ref()
.map(|existing| existing == key)
.unwrap_or(false)
});
if exists {
return Ok(());
}
}
if item.task_id().is_none() {
item = item
.into_builder()
.task_id(TaskId::from_string(RandomId::default()))
.build();
}
tasks.push_back(item);
if let Some(waker) = this.waker.lock().unwrap().as_ref() {
waker.wake_by_ref();
}
Ok(())
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
impl<Args> BackendConfig for VecDequeBackend<Args> {
type Args = Args;
type Id = RandomId;
type Kind = Ephemeral;
type Config = ();
type Layer = Identity;
fn config(&self) -> &Self::Config {
&()
}
fn middleware(&mut self, _worker: &mut WorkerContext) -> Self::Layer {
Identity::new()
}
}