#![allow(non_snake_case)]
use super::{RiGatewayRequest, RiGatewayResponse};
use super::radix_tree::RiRadixTree;
use crate::core::RiResult;
use crate::core::lock::RwLockExtensions;
use crate::gateway::middleware::RiMiddleware;
use std::collections::HashMap as FxHashMap;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
#[cfg(feature = "pyo3")]
use pyo3::PyResult;
pub type RiRouteHandler = Arc<
dyn Fn(RiGatewayRequest) -> Pin<Box<dyn Future<Output = RiResult<RiGatewayResponse>> + Send>>
+ Send
+ Sync,
>;
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
#[derive(Clone)]
pub struct RiRoute {
pub method: String,
pub path: String,
pub handler: RiRouteHandler,
pub middleware: Vec<Arc<dyn RiMiddleware>>,
}
impl fmt::Debug for RiRoute {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RiRoute")
.field("method", &self.method)
.field("path", &self.path)
.field("handler", &"<handler>")
.field("middleware_count", &self.middleware.len())
.finish()
}
}
impl RiRoute {
pub fn new(
method: String,
path: String,
handler: RiRouteHandler,
) -> Self {
Self {
method,
path,
handler,
middleware: Vec::new(),
}
}
pub fn with_middleware(mut self, middleware: Arc<dyn RiMiddleware>) -> Self {
self.middleware.push(middleware);
self
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiRoute {
#[new]
fn py_new(method: String, path: String) -> PyResult<Self> {
use crate::gateway::{RiGatewayRequest, RiGatewayResponse};
let handler = Arc::new(|_req: RiGatewayRequest| -> Pin<Box<dyn Future<Output = Result<RiGatewayResponse, crate::core::RiError>> + Send>> {
Box::pin(async move {
Ok(RiGatewayResponse {
status_code: 200,
headers: FxHashMap::default(),
body: b"Hello from Ri Python!".to_vec(),
request_id: String::new(),
})
})
});
Ok(Self {
method,
path,
handler,
middleware: Vec::new(),
})
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiRouter {
trees: std::sync::RwLock<FxHashMap<String, RiRadixTree>>,
routes: std::sync::RwLock<Vec<RiRoute>>,
route_cache: std::sync::RwLock<FxHashMap<String, RiRoute>>,
}
impl Default for RiRouter {
fn default() -> Self {
Self::new()
}
}
impl RiRouter {
pub fn new() -> Self {
Self {
trees: std::sync::RwLock::new(FxHashMap::default()),
routes: std::sync::RwLock::new(Vec::new()),
route_cache: std::sync::RwLock::new(FxHashMap::default()),
}
}
#[allow(dead_code)]
fn get_or_create_tree(&self, method: &str) -> RiRadixTree {
let trees = match self.trees.read_safe("trees for get_or_create") {
Ok(t) => t,
Err(_) => return RiRadixTree::new(),
};
if let Some(_tree) = trees.get(method) {
return RiRadixTree::new();
}
drop(trees);
let mut trees_mut = match self.trees.write_safe("trees for create") {
Ok(t) => t,
Err(_) => return RiRadixTree::new(),
};
let tree = RiRadixTree::new();
trees_mut.insert(method.to_string(), RiRadixTree::new());
tree
}
pub fn add_route(&self, route: RiRoute) {
let method = route.method.clone();
let mut trees = match self.trees.write_safe("trees for add_route") {
Ok(t) => t,
Err(e) => {
log::error!("Failed to acquire trees write lock: {}", e);
return;
}
};
let tree = trees.entry(method).or_insert_with(RiRadixTree::new);
tree.insert(route.clone());
drop(trees);
let mut routes = match self.routes.write_safe("routes for add_route") {
Ok(r) => r,
Err(e) => {
log::error!("Failed to acquire routes write lock: {}", e);
return;
}
};
routes.push(route);
let mut cache = match self.route_cache.write_safe("cache for add_route") {
Ok(c) => c,
Err(e) => {
log::error!("Failed to acquire cache write lock: {}", e);
return;
}
};
cache.clear();
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiRouter {
#[new]
fn py_new() -> Self {
Self::new()
}
fn add_get_route(&self, path: String) {
let route = RiRoute::py_new("GET".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_post_route(&self, path: String) {
let route = RiRoute::py_new("POST".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_put_route(&self, path: String) {
let route = RiRoute::py_new("PUT".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_delete_route(&self, path: String) {
let route = RiRoute::py_new("DELETE".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_patch_route(&self, path: String) {
let route = RiRoute::py_new("PATCH".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_options_route(&self, path: String) {
let route = RiRoute::py_new("OPTIONS".to_string(), path).expect("Failed to create route");
self.add_route(route);
}
fn add_custom_route(&self, method: String, path: String) {
let route = RiRoute::py_new(method, path).expect("Failed to create route");
self.add_route(route);
}
fn get_route_count(&self) -> usize {
self.route_count()
}
fn clear_all_routes(&self) {
self.clear_routes();
}
}
impl RiRouter {
pub fn get(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("GET".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub fn post(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("POST".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub fn put(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("PUT".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub fn delete(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("DELETE".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub fn patch(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("PATCH".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub fn options(&self, path: &str, handler: RiRouteHandler) {
let route = RiRoute::new("OPTIONS".to_string(), path.to_string(), handler);
self.add_route(route);
}
pub async fn route(&self, request: &RiGatewayRequest) -> RiResult<RiRouteHandler> {
let cache_key = format!("{}:{}", request.method, request.path);
{
let cache = match self.route_cache.read_safe("cache for route lookup") {
Ok(c) => c,
Err(_) => return Err(crate::core::RiError::InvalidState("Failed to acquire cache read lock".to_string())),
};
if let Some(cached_route) = cache.get(&cache_key) {
return Ok(cached_route.handler.clone());
}
}
let trees = match self.trees.read_safe("trees for route lookup") {
Ok(t) => t,
Err(_) => return Err(crate::core::RiError::InvalidState("Failed to acquire trees read lock".to_string())),
};
if let Some(tree) = trees.get(&request.method) {
if let Some(route_match) = tree.find(&request.path) {
let mut cache = match self.route_cache.write_safe("cache for route insert") {
Ok(c) => c,
Err(_) => return Ok(route_match.route.handler.clone()),
};
cache.insert(cache_key.clone(), route_match.route.clone());
return Ok(route_match.route.handler.clone());
}
}
Err(crate::core::RiError::Other(format!(
"No route found for {} {}",
request.method, request.path
)))
}
#[allow(dead_code)]
fn matches_route(&self, route_method: &str, route_path: &str, request_method: &str, request_path: &str) -> bool {
if route_method != request_method {
return false;
}
let trees = match self.trees.read_safe("trees for matches_route") {
Ok(t) => t,
Err(_) => return false,
};
if let Some(tree) = trees.get(route_method) {
if let Some(route_match) = tree.find(request_path) {
return route_match.route.path == route_path;
}
}
false
}
pub fn mount(&self, prefix: &str, router: &RiRouter) {
let routes = match router.routes.read_safe("routes for mount") {
Ok(r) => r,
Err(_) => return,
};
for route in routes.iter() {
let mounted_path = if prefix.ends_with('/') && route.path.starts_with('/') {
format!("{}{}", prefix, &route.path[1..])
} else if !prefix.ends_with('/') && !route.path.starts_with('/') {
format!("{}/{}", prefix, route.path)
} else {
format!("{}{}", prefix, route.path)
};
let mut mounted_route = route.clone();
mounted_route.path = mounted_path;
self.add_route(mounted_route);
}
}
pub fn clear_routes(&self) {
let trees = match self.trees.write_safe("trees for clear") {
Ok(t) => t,
Err(e) => {
log::error!("Failed to acquire trees write lock: {}", e);
return;
}
};
for tree in trees.values() {
tree.clear();
}
drop(trees);
let mut routes = match self.routes.write_safe("routes for clear") {
Ok(r) => r,
Err(e) => {
log::error!("Failed to acquire routes write lock: {}", e);
return;
}
};
let mut cache = match self.route_cache.write_safe("cache for clear") {
Ok(c) => c,
Err(e) => {
log::error!("Failed to acquire cache write lock: {}", e);
return;
}
};
routes.clear();
cache.clear();
}
pub fn route_count(&self) -> usize {
match self.routes.read_safe("routes for count") {
Ok(routes) => routes.len(),
Err(_) => 0,
}
}
}