From 518cbf0c2658bb43ba2297df973a3c2b1b3dca61 Mon Sep 17 00:00:00 2001 From: metamuffin Date: Thu, 10 Apr 2025 20:27:38 +0200 Subject: add semaphore module --- src/modules/mod.rs | 2 ++ src/modules/semaphore.rs | 56 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) create mode 100644 src/modules/semaphore.rs diff --git a/src/modules/mod.rs b/src/modules/mod.rs index 65e9666..71f60c9 100644 --- a/src/modules/mod.rs +++ b/src/modules/mod.rs @@ -34,6 +34,7 @@ mod ratelimit; mod redirect; mod switch; mod upgrade_insecure; +mod semaphore; pub type NodeRequest = Request>; pub type NodeResponse = Response>; @@ -63,6 +64,7 @@ pub static MODULES: &[&dyn NodeKind] = &[ &delay::DelayKind, &ratelimit::RatelimitKind, &log::LogKind, + &semaphore::SemaphoreKind, ]; pub struct NodeContext { diff --git a/src/modules/semaphore.rs b/src/modules/semaphore.rs new file mode 100644 index 0000000..b4f50b2 --- /dev/null +++ b/src/modules/semaphore.rs @@ -0,0 +1,56 @@ +/* + This file is part of gnix (https://codeberg.org/metamuffin/gnix) + which is licensed under the GNU Affero General Public License (version 3); see /COPYING. + Copyright (C) 2025 metamuffin +*/ +use super::{Node, NodeContext, NodeKind, NodeRequest, NodeResponse}; +use crate::{config::DynNode, error::ServiceError}; +use anyhow::Result; +use serde::Deserialize; +use serde_yml::Value; +use std::{ + future::Future, + pin::Pin, + sync::Arc, +}; +use tokio::sync::Semaphore; + +pub struct SemaphoreKind; + +#[derive(Deserialize)] +struct SemaphoreConfig { + permits: usize, + next: DynNode, +} + +struct SemaphoreState { + config: SemaphoreConfig, + semaphore: Semaphore, +} + +impl NodeKind for SemaphoreKind { + fn name(&self) -> &'static str { + "semaphore" + } + fn instanciate(&self, config: Value) -> Result> { + let config = serde_yml::from_value::(config)?; + Ok(Arc::new(SemaphoreState { + semaphore: Semaphore::new(config.permits), + config, + })) + } +} +impl Node for SemaphoreState { + fn handle<'a>( + &'a self, + context: &'a mut NodeContext, + request: NodeRequest, + ) -> Pin> + Send + Sync + 'a>> { + Box::pin(async move { + let _permit = self.semaphore.acquire().await; + let resp = self.config.next.handle(context, request).await?; + drop(_permit); + Ok(resp) + }) + } +} -- cgit v1.2.3-70-g09d2