qpismont/steep
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 5 Wiki Activity

Project website + use Apalis crate for task queue (with redis backend)
ci/woodpecker/push/tests Pipeline was successful
Details

Browse Source 
instead channel and semaphore queue + upgrade to rust 1.99
This commit is contained in:
qpismont committed 2026-10-05 20:44:57 +00:00
1 parent c546707f72
commit d74db2096c
23 files changed
+1234 -183
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/steep-server/src/api.rs
+16 -6
View File
@@ -14,7 +14,10 @@ use tracing::{info, instrument};
use tokio_util::sync::CancellationToken;
use apalis::prelude::Storage;
use crate::consts::{GITEA_EVENT_TYPE_HEADER_NAME, GITEA_SIG_HEADER_NAME, MAX_WEBHOOK_BODY_SIZE};
use crate::dedup::{Dedup, WEBHOOK_TTL};
use crate::errors::AppError;
use crate::gitea::WebhookType;
use crate::metrics;
@@ -61,15 +64,22 @@ async fn webhook(
metrics::webhook_received(event_type);
let event_id = wb.event_id();
if app_state.bot.check_and_mark(event_id).await {
let webhook_key = Dedup::webhook_key(event_id);
if !app_state
.dedup
.claim(&webhook_key, WEBHOOK_TTL)
.await
.map_err(anyhow::Error::from)?
{
metrics::webhook_duplicate(event_type);
return Err(AppError::AlreadyProcessedErr);
}
if app_state.bot_tx.try_send(wb).is_err() {
app_state.bot.unmark(event_id).await;
metrics::webhook_channel_full(event_type);
return Err(AppError::ChannelFullErr);
let mut storage = app_state.storage.clone();
if let Err(err) = storage.push(wb).await {
app_state.dedup.release(&webhook_key).await.ok();
metrics::webhook_enqueue_failed(event_type);
return Err(AppError::EnqueueErr(anyhow::Error::from(err)));
}
Ok((StatusCode::CREATED, "Task started"))
@@ -109,7 +119,7 @@ where
});
});
let webhook = parse_webhook(&type_header, &app_state.bot.name(), &body_bytes)?;
let webhook = parse_webhook(&type_header, &app_state.bot_name, &body_bytes)?;
Ok(WebhookExtract(webhook))
}
}
crates/steep-server/src/bot.rs
+65 -143
View File
@@ -1,12 +1,13 @@
use crate::{
dedup::{Dedup, REVIEW_TTL},
gitea::{GiteaAPI, WebhookType},
metrics,
open_router::OpenRouterClient,
sandbox::SandboxConfig,
};
use apalis::prelude::{Data, Error, Monitor, WorkerBuilder, WorkerBuilderExt, WorkerFactoryFn};
use apalis_redis::RedisStorage;
use serde::{Deserialize, Deserializer};
use std::{collections::HashSet, sync::Arc};
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, instrument};
@@ -120,87 +121,64 @@ where
#[derive(Clone)]
pub struct Bot {
bot_name: String,
gitea_api: GiteaAPI,
open_router_client: OpenRouterClient,
max_concurrent: usize,
open_router_model: String,
sandbox: SandboxConfig,
actions_handled: Arc<Mutex<HashSet<u64>>>,
dedup: Dedup,
}
impl Bot {
pub fn new(
bot_name: String,
gitea_api: GiteaAPI,
open_router_client: OpenRouterClient,
max_concurrent: usize,
open_router_model: String,
sandbox: SandboxConfig,
dedup: Dedup,
) -> Self {
Self {
bot_name,
gitea_api,
open_router_client,
max_concurrent,
open_router_model,
sandbox,
actions_handled: Arc::new(Mutex::new(HashSet::new())),
dedup,
}
}
pub fn name(&self) -> String {
self.bot_name.clone()
}
/// Consomme la file Redis jusqu'à l'annulation de `shutdown`.
///
/// Les jobs en cours sont laissés se terminer, et ceux qui n'ont pas encore
/// été acquittés restent dans Redis : ils seront repris au prochain démarrage.
pub async fn start(
&self,
mut rx: tokio::sync::mpsc::Receiver<WebhookType>,
storage: RedisStorage<WebhookType>,
shutdown: CancellationToken,
) -> anyhow::Result<()> {
info!("Bot started");
let sem = Arc::new(tokio::sync::Semaphore::new(self.max_concurrent));
let mut tasks = tokio::task::JoinSet::new();
let worker = WorkerBuilder::new("steep-review")
.concurrency(self.max_concurrent)
.data(self.clone())
.backend(storage)
.build_fn(handle_review);
loop {
let wb = tokio::select! {
biased;
_ = shutdown.cancelled() => break,
msg = rx.recv() => match msg {
Some(wb) => wb,
None => break,
},
};
while let Some(res) = tasks.try_join_next() {
if let Err(e) = res {
error!("Task panicked: {e}");
}
}
info!(queued = rx.len(), active = tasks.len(), "Webhook received");
let permit = sem.clone().acquire_owned().await?;
let self_clone = self.clone();
metrics::increment_task_active();
tasks.spawn(async move {
self_clone.exec(wb).await;
drop(permit);
metrics::decrement_task_active();
});
}
tasks.join_all().await;
Monitor::new()
.register(worker)
.run_with_signal(async move {
shutdown.cancelled().await;
Ok(())
})
.await?;
info!("Bot shutting down complete");
Ok(())
}
#[instrument(skip(self, webhook), fields(repo, pr))]
pub async fn exec(&self, webhook: WebhookType) {
let event_type_str = webhook.event_type_str();
pub async fn exec(&self, webhook: WebhookType) -> anyhow::Result<()> {
match &webhook {
WebhookType::Review(p) => {
tracing::Span::current().record("repo", &p.repository.full_name);
@@ -210,111 +188,55 @@ impl Bot {
let tools = crate::sandbox::tools::for_webhook(&webhook);
let exec_result = match webhook {
WebhookType::Review(review_payload) => crate::bot_actions::review::exec_review(
&self.gitea_api,
&self.open_router_client,
&self.open_router_model,
&self.sandbox,
tools,
review_payload,
),
}
.await;
match exec_result {
Ok(_) => {
metrics::task_completed(event_type_str);
info!("Task completed");
}
Err(err) => {
metrics::task_failed(event_type_str);
error!(%err, "Task error");
sentry_anyhow::capture_anyhow(&err);
match webhook {
WebhookType::Review(review_payload) => {
crate::bot_actions::review::exec_review(
&self.gitea_api,
&self.open_router_client,
&self.open_router_model,
&self.sandbox,
tools,
review_payload,
)
.await
}
}
}
pub async fn check_and_mark(&self, event_id: u64) -> bool {
let mut action_handled_lock = self.actions_handled.lock().await;
!action_handled_lock.insert(event_id)
}
pub async fn unmark(&self, event_id: u64) {
let mut action_handled_lock = self.actions_handled.lock().await;
action_handled_lock.remove(&event_id);
}
}
#[cfg(test)]
mod tests {
use super::*;
async fn handle_review(webhook: WebhookType, bot: Data<Bot>) -> Result<(), Error> {
let event_type = webhook.event_type_str();
let review_key = Dedup::review_key(webhook.event_id());
fn make_actions_handled() -> Arc<Mutex<HashSet<u64>>> {
Arc::new(Mutex::new(HashSet::new()))
match bot.dedup.is_marked(&review_key).await {
Ok(true) => {
metrics::task_skipped(event_type);
info!("Review already done, skipping");
return Ok(());
}
Ok(false) => {}
Err(err) => error!(%err, "Dedup lookup failed, reviewing anyway"),
}
async fn check_and_mark(actions_handled: &Arc<Mutex<HashSet<u64>>>, event_id: u64) -> bool {
let mut lock = actions_handled.lock().await;
!lock.insert(event_id)
metrics::increment_task_active();
let result = bot.exec(webhook).await;
metrics::decrement_task_active();
match result {
Ok(()) => {
metrics::task_completed(event_type);
info!("Task completed");
if let Err(err) = bot.dedup.mark(&review_key, REVIEW_TTL).await {
error!(%err, "Failed to mark review as done");
}
}
Err(err) => {
metrics::task_failed(event_type);
error!(%err, "Task error");
sentry_anyhow::capture_anyhow(&err);
}
}
async fn unmark(actions_handled: &Arc<Mutex<HashSet<u64>>>, event_id: u64) {
let mut lock = actions_handled.lock().await;
lock.remove(&event_id);
}
#[tokio::test]
async fn test_check_and_mark_first_call_returns_false() {
let actions_handled = make_actions_handled();
assert!(!check_and_mark(&actions_handled, 1).await);
}
#[tokio::test]
async fn test_check_and_mark_second_call_returns_true() {
let actions_handled = make_actions_handled();
check_and_mark(&actions_handled, 1).await;
assert!(check_and_mark(&actions_handled, 1).await);
}
#[tokio::test]
async fn test_check_and_mark_different_ids_are_independent() {
let actions_handled = make_actions_handled();
check_and_mark(&actions_handled, 1).await;
assert!(!check_and_mark(&actions_handled, 2).await);
}
#[tokio::test]
async fn test_unmark_allows_reprocessing() {
let actions_handled = make_actions_handled();
check_and_mark(&actions_handled, 1).await;
unmark(&actions_handled, 1).await;
assert!(!check_and_mark(&actions_handled, 1).await);
}
#[tokio::test]
async fn test_unmark_nonexistent_id_is_noop() {
let actions_handled = make_actions_handled();
unmark(&actions_handled, 99).await;
assert!(!check_and_mark(&actions_handled, 99).await);
}
#[tokio::test]
async fn test_check_and_mark_concurrent_same_id() {
let actions_handled = make_actions_handled();
let actions_handled2 = Arc::clone(&actions_handled);
let t1 = tokio::spawn(async move { check_and_mark(&actions_handled, 42).await });
let t2 = tokio::spawn(async move { check_and_mark(&actions_handled2, 42).await });
let (r1, r2) = tokio::join!(t1, t2);
let results = [r1.unwrap(), r2.unwrap()];
// exactement un seul des deux doit retourner false (non traité)
assert_eq!(results.iter().filter(|&&r| !r).count(), 1);
// l'autre doit retourner true (déjà traité)
assert_eq!(results.iter().filter(|&&r| r).count(), 1);
}
Ok(())
}
crates/steep-server/src/dedup.rs
+120
View File
@@ -0,0 +1,120 @@
use std::time::Duration;
use redis::aio::ConnectionManager;
pub const WEBHOOK_TTL: Duration = Duration::from_secs(24 * 60 * 60);
pub const REVIEW_TTL: Duration = Duration::from_secs(7 * 24 * 60 * 60);
#[derive(Clone)]
pub struct Dedup {
conn: ConnectionManager,
}
impl Dedup {
pub fn new(conn: ConnectionManager) -> Self {
Self { conn }
}
pub fn webhook_key(event_id: u64) -> String {
format!("steep:webhook:{event_id}")
}
pub fn review_key(event_id: u64) -> String {
format!("steep:review:{event_id}")
}
pub async fn claim(&self, key: &str, ttl: Duration) -> redis::RedisResult<bool> {
let set: Option<String> = redis::cmd("SET")
.arg(key)
.arg(1)
.arg("NX")
.arg("EX")
.arg(ttl.as_secs())
.query_async(&mut self.conn.clone())
.await?;
Ok(set.is_some())
}
pub async fn mark(&self, key: &str, ttl: Duration) -> redis::RedisResult<()> {
let _: String = redis::cmd("SET")
.arg(key)
.arg(1)
.arg("EX")
.arg(ttl.as_secs())
.query_async(&mut self.conn.clone())
.await?;
Ok(())
}
pub async fn is_marked(&self, key: &str) -> redis::RedisResult<bool> {
let exists: i64 = redis::cmd("EXISTS")
.arg(key)
.query_async(&mut self.conn.clone())
.await?;
Ok(exists > 0)
}
pub async fn release(&self, key: &str) -> redis::RedisResult<()> {
let _: i64 = redis::cmd("DEL")
.arg(key)
.query_async(&mut self.conn.clone())
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn webhook_key_is_namespaced() {
assert_eq!(Dedup::webhook_key(42), "steep:webhook:42");
}
#[test]
fn review_key_is_namespaced() {
assert_eq!(Dedup::review_key(42), "steep:review:42");
}
#[test]
fn the_two_keys_for_the_same_event_differ() {
assert_ne!(Dedup::webhook_key(7), Dedup::review_key(7));
}
async fn redis_dedup() -> Dedup {
let url = std::env::var("REDIS_URL").expect("REDIS_URL must be set");
Dedup::new(apalis_redis::connect(url).await.expect("connect to redis"))
}
#[tokio::test]
#[ignore = "requires a running Redis (REDIS_URL)"]
async fn claim_then_claim_is_a_duplicate() {
let dedup = redis_dedup().await;
let key = "steep:test:claim";
dedup.release(key).await.unwrap();
assert!(dedup.claim(key, Duration::from_secs(30)).await.unwrap());
assert!(!dedup.claim(key, Duration::from_secs(30)).await.unwrap());
dedup.release(key).await.unwrap();
}
#[tokio::test]
#[ignore = "requires a running Redis (REDIS_URL)"]
async fn mark_then_is_marked() {
let dedup = redis_dedup().await;
let key = "steep:test:mark";
dedup.release(key).await.unwrap();
assert!(!dedup.is_marked(key).await.unwrap());
dedup.mark(key, Duration::from_secs(30)).await.unwrap();
assert!(dedup.is_marked(key).await.unwrap());
dedup.release(key).await.unwrap();
}
}
crates/steep-server/src/env.rs
+3
View File
@@ -11,6 +11,7 @@ pub struct EnvConfig {
pub gitea_url: String,
pub gitea_token: String,
pub gitea_timeout: u64,
pub redis_url: String,
pub metrics_bind_addr: Option<String>,
pub sandbox_max_iterations: usize,
}
@@ -25,6 +26,7 @@ pub fn load_config() -> anyhow::Result<EnvConfig> {
let gitea_url = try_get_env("GITEA_URL")?;
let gitea_token = try_get_env("GITEA_TOKEN")?;
let gitea_timeout = try_get_env("GITEA_TIMEOUT")?.parse()?;
let redis_url = try_get_env("REDIS_URL")?;
let metrics_bind_addr = std::env::var("METRICS_BIND_ADDR").ok();
let sandbox_max_iterations = std::env::var("SANDBOX_MAX_ITERATIONS")
.ok()
@@ -41,6 +43,7 @@ pub fn load_config() -> anyhow::Result<EnvConfig> {
gitea_url,
gitea_token,
gitea_timeout,
redis_url,
metrics_bind_addr,
sandbox_max_iterations,
})
crates/steep-server/src/errors.rs
+5 -6
View File
@@ -1,4 +1,3 @@
use anyhow::anyhow;
use axum::response::IntoResponse;
use reqwest::StatusCode;
@@ -22,8 +21,8 @@ pub enum AppError {
#[error("WebHook have bad action")]
InvalidActionErr,
#[error("Channel full")]
ChannelFullErr,
#[error("Could not enqueue task")]
EnqueueErr(#[source] anyhow::Error),
#[error("Already processed")]
AlreadyProcessedErr,
@@ -62,11 +61,11 @@ impl IntoResponse for AppError {
"WebHook sig header is invalid".to_string(),
),
AppError::AlreadyProcessedErr => (StatusCode::OK, "Already processed".to_string()),
AppError::ChannelFullErr => {
sentry_anyhow::capture_anyhow(&anyhow!("Max concurrent tasks reached"));
AppError::EnqueueErr(err) => {
sentry_anyhow::capture_anyhow(&err);
(
StatusCode::SERVICE_UNAVAILABLE,
"Max concurrent tasks reached".to_string(),
"Could not enqueue task".to_string(),
)
}
AppError::Other(err) => {
crates/steep-server/src/gitea.rs
+6 -6
View File
@@ -1,7 +1,7 @@
use std::time::Duration;
use futures_util::stream::TryStreamExt;
use serde::Deserialize;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tokio::io::AsyncReadExt;
use tokio_util::io::StreamReader;
@@ -281,7 +281,7 @@ async fn read_capped(response: reqwest::Response) -> anyhow::Result<String> {
Ok(String::from_utf8_lossy(&buf).into_owned())
}
#[derive(Debug)]
#[derive(Debug, Serialize, Deserialize)]
pub enum WebhookType {
Review(ReviewPayload),
}
@@ -300,7 +300,7 @@ impl WebhookType {
}
}
#[derive(Deserialize, Debug)]
#[derive(Deserialize, Serialize, Debug)]
pub struct ReviewPayload {
pub action: String,
pub pull_request: PullRequest,
@@ -308,13 +308,13 @@ pub struct ReviewPayload {
pub comment: Comment,
}
#[derive(Deserialize, Debug)]
#[derive(Deserialize, Serialize, Debug)]
pub struct PullRequest {
pub number: u64,
pub title: String,
}
#[derive(Deserialize, Debug)]
#[derive(Deserialize, Serialize, Debug)]
pub struct Comment {
pub id: u64,
pub body: String,
@@ -325,7 +325,7 @@ pub struct User {
pub login: String,
}
#[derive(Deserialize, Debug)]
#[derive(Deserialize, Serialize, Debug)]
pub struct Repository {
pub full_name: String,
pub clone_url: String,
crates/steep-server/src/main.rs
+12 -5
View File
@@ -2,6 +2,7 @@ use std::time::Duration;
use crate::{
bot::Bot,
dedup::Dedup,
gitea::{GiteaAPI, WebhookType},
open_router::OpenRouterClient,
sandbox::SandboxConfig,
@@ -18,6 +19,7 @@ mod api;
mod bot;
mod bot_actions;
mod consts;
mod dedup;
mod env;
mod errors;
mod gitea;
@@ -99,19 +101,24 @@ async fn run() -> anyhow::Result<()> {
);
}
let redis_conn = apalis_redis::connect(config.redis_url.clone()).await?;
let storage: apalis_redis::RedisStorage<WebhookType> =
apalis_redis::RedisStorage::new(redis_conn.clone());
let dedup = Dedup::new(redis_conn);
let bot = Bot::new(
gitea_user.login,
gitea_api,
open_router_client,
config.bot_max_concurrent,
config.open_router_model.clone(),
sandbox,
dedup.clone(),
);
let (tx, rx) = tokio::sync::mpsc::channel::<WebhookType>(config.bot_max_concurrent * 2);
let app_state = AppState {
bot_tx: tx,
bot: bot.clone(),
storage: storage.clone(),
bot_name: gitea_user.login.clone(),
dedup,
config,
};
@@ -136,7 +143,7 @@ async fn run() -> anyhow::Result<()> {
tokio::select! {
result = async {
tokio::try_join!(
bot.start(rx, shutdown.clone()),
bot.start(storage, shutdown.clone()),
api::start(app_state, shutdown.clone()),
signal
)
crates/steep-server/src/metrics.rs
+13 -4
View File
@@ -10,8 +10,8 @@ pub fn webhook_duplicate(event_type: &str) {
counter!("steep_webhooks_duplicate_total", "event_type" => event_type.to_string()).increment(1);
}
pub fn webhook_channel_full(event_type: &str) {
counter!("steep_webhooks_channel_full_total", "event_type" => event_type.to_string())
pub fn webhook_enqueue_failed(event_type: &str) {
counter!("steep_webhooks_enqueue_failed_total", "event_type" => event_type.to_string())
.increment(1);
}
@@ -23,6 +23,10 @@ pub fn decrement_task_active() {
gauge!("steep_bot_tasks_active").decrement(1.0);
}
pub fn task_skipped(event_type: &str) {
counter!("steep_bot_tasks_skipped_total", "event_type" => event_type.to_string()).increment(1);
}
pub fn task_completed(event_type: &str) {
counter!("steep_bot_tasks_completed_total", "event_type" => event_type.to_string())
.increment(1);
@@ -48,15 +52,20 @@ pub fn describe() {
"Webhooks rejected as duplicates"
);
describe_counter!(
"steep_webhooks_channel_full_total",
"steep_webhooks_enqueue_failed_total",
Unit::Count,
"Webhooks dropped because the bot channel was full"
"Webhooks dropped because enqueueing to Redis failed"
);
describe_gauge!(
"steep_bot_tasks_active",
Unit::Count,
"Bot tasks currently in progress"
);
describe_counter!(
"steep_bot_tasks_skipped_total",
Unit::Count,
"Reviews skipped because they were already done"
);
describe_counter!(
"steep_bot_tasks_completed_total",
Unit::Count,
crates/steep-server/src/state.rs
+6 -3
View File
@@ -1,8 +1,11 @@
use crate::{bot::Bot, env::EnvConfig, gitea::WebhookType};
use apalis_redis::RedisStorage;
use crate::{dedup::Dedup, env::EnvConfig, gitea::WebhookType};
#[derive(Clone)]
pub struct AppState {
pub bot_tx: tokio::sync::mpsc::Sender<WebhookType>,
pub bot: Bot,
pub storage: RedisStorage<WebhookType>,
pub bot_name: String,
pub dedup: Dedup,
pub config: EnvConfig,
}