Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Rust: TaskQueueReceive Receive to ReceiveAll #418

Merged
merged 2 commits into from
May 6, 2022
Merged

Rust: TaskQueueReceive Receive to ReceiveAll #418

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6444e18
Rust: TaskQueueReceive Receive to ReceiveAll
jaymell May 4, 2022
20786ef
Improve test readability
jaymell May 6, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions server/svix-server/src/queue/memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,10 @@ pub struct MemoryQueueConsumer {

#[async_trait]
impl TaskQueueReceive for MemoryQueueConsumer {
async fn receive(&mut self) -> Result<TaskQueueDelivery> {
async fn receive_all(&mut self) -> Result<Vec<TaskQueueDelivery>> {
if let Some(delivery) = self.rx.recv().await {
tracing::trace!("MemoryQueue: event recv <");
Ok(delivery)
Ok(vec![delivery])
} else {
Err(Error::Queue("Failed to fetch from queue".to_owned()))
}
Expand Down
31 changes: 22 additions & 9 deletions server/svix-server/src/queue/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,8 @@ impl TaskQueueProducer {
pub struct TaskQueueConsumer(Box<dyn TaskQueueReceive + Send + Sync>);

impl TaskQueueConsumer {
pub async fn receive(&mut self) -> Result<TaskQueueDelivery> {
self.0.receive().await
pub async fn receive_all(&mut self) -> Result<Vec<TaskQueueDelivery>> {
self.0.receive_all().await
}
}

Expand Down Expand Up @@ -112,7 +112,7 @@ impl Clone for Box<dyn TaskQueueSend> {

#[async_trait]
trait TaskQueueReceive {
async fn receive(&mut self) -> Result<TaskQueueDelivery>;
async fn receive_all(&mut self) -> Result<Vec<TaskQueueDelivery>>;
}

#[cfg(test)]
Expand Down Expand Up @@ -143,7 +143,7 @@ mod tests {
/// equal to the mock message with the given message_id.
async fn assert_recv(rx: &mut TaskQueueConsumer, message_id: &str) {
assert_eq!(
rx.receive().await.unwrap().task,
rx.receive_all().await.unwrap().get(0).unwrap().task,
mock_message(message_id.to_owned())
)
}
Expand Down Expand Up @@ -211,13 +211,20 @@ mod tests {
.await
.is_ok());

let recv = rx_mem.receive().await.unwrap();
assert_eq!(&recv.task, &mock_message("1".to_owned()));
let recv = rx_mem
.receive_all()
.await
.unwrap()
.into_iter()
.nth(0)
.unwrap();

assert_eq!(recv.task, mock_message("1".to_owned()));
svix-james marked this conversation as resolved.
Show resolved Hide resolved

assert!(tx_mem.ack(recv).await.is_ok());

tokio::select! {
_ = rx_mem.receive() => {
_ = rx_mem.receive_all() => {
panic!("`rx_mem` received second message");
}

Expand All @@ -236,13 +243,19 @@ mod tests {
.await
.is_ok());

let recv = rx_mem.receive().await.unwrap();
let recv = rx_mem
.receive_all()
.await
.unwrap()
.into_iter()
.nth(0)
.unwrap();
assert_eq!(&recv.task, &mock_message("1".to_owned()));

assert!(tx_mem.nack(recv).await.is_ok());

tokio::select! {
_ = rx_mem.receive() => {}
_ = rx_mem.receive_all() => {}

// FIXME: Find out correct timeout duration
_ = tokio::time::sleep(Duration::from_millis(500)) => {
Expand Down
4 changes: 2 additions & 2 deletions server/svix-server/src/queue/redis.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ pub struct RedisQueueConsumer {

#[async_trait]
impl TaskQueueReceive for RedisQueueConsumer {
async fn receive(&mut self) -> Result<TaskQueueDelivery> {
async fn receive_all(&mut self) -> Result<Vec<TaskQueueDelivery>> {
let mut pool = self.pool.get().await.unwrap();
let mut cmd = redis::cmd("BLMOVE");
cmd.arg(MAIN)
Expand All @@ -228,7 +228,7 @@ impl TaskQueueReceive for RedisQueueConsumer {
.await
.map_err(|x| Error::Queue(x.to_string()))?;
tracing::trace!("RedisQueue: event recv <");
Ok(from_redis_key(&key))
Ok(vec![from_redis_key(&key)])
}
}

Expand Down
47 changes: 25 additions & 22 deletions server/svix-server/src/worker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,30 +300,33 @@ pub async fn worker_loop(
mut queue_rx: TaskQueueConsumer,
) -> Result<()> {
loop {
let pool = pool.clone();
match queue_rx.receive().await {
Ok(delivery) => {
let queue_task = delivery.task.clone();
let queue_tx = queue_tx.clone();
let cfg = cfg.clone();
match queue_task {
QueueTask::MessageV1(msg) => {
let cache = cache.clone();
tokio::spawn(async move {
if let Err(err) = dispatch(cfg, &pool, cache, &queue_tx, msg).await {
tracing::error!("Error executing task: {}", err);
queue_tx.nack(delivery).await.expect(
"Error sending 'nack' to Redis after task execution error",
);
} else {
// No unwrap
queue_tx.ack(delivery).await.expect(
match queue_rx.receive_all().await {
Ok(batch) => {
for delivery in batch {
let pool = pool.clone();
let queue_task = delivery.task.clone();
let queue_tx = queue_tx.clone();
let cfg = cfg.clone();
match queue_task {
QueueTask::MessageV1(msg) => {
let cache = cache.clone();
tokio::spawn(async move {
if let Err(err) = dispatch(cfg, &pool, cache, &queue_tx, msg).await
{
tracing::error!("Error executing task: {}", err);
queue_tx.nack(delivery).await.expect(
"Error sending 'nack' to Redis after task execution error",
);
} else {
// No unwrap
queue_tx.ack(delivery).await.expect(
"Error sending 'ack' to Redis after successful task execution",
);
}
});
}
};
}
});
}
};
}
}
Err(err) => {
tracing::error!("Error receiving task: {}", err);
Expand Down