Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
25 changes: 20 additions & 5 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,13 @@
name = "gpu-share-vm-manager"
version = "0.1.0"
edition = "2021"
resolver = "2"
default-run = "gpu-share-vm-manager"

[dependencies]

tokio = { version = "1.36", features = ["full"] }
virt = "0.4.1"
# virt = "0.4.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
tracing = "0.1"
Expand All @@ -14,16 +17,20 @@ anyhow = "1.0"
async-trait = "0.1"
config = "0.15.6"
axum = { version = "0.8.0", features = ["macros"] }
hyper = { version = "1.0", features = ["full"] }
hyper = { version = "0.14.32", features = ["full"] }
tower = { version = "0.5.2", features = ["limit", "util"] }
tower-http = { version = "0.6.2", features = ["trace", "limit", "add-extension"] }
clap = { version = "4.4", features = ["derive"] }
colored = "3.0"
thiserror = "2.0.11"
chrono = "0.4"
chrono = { version = "0.4", features = ["serde"] }
uuid = { version = "1.8.0", features = ["v4"] }
governor = { version = "0.8", features = ["dashmap"] }
jsonwebtoken = "8.3.0"
jsonwebtoken = "9.3.0"
bollard = "0.15.0"
futures-util = "0.3"
ratatui = "0.26"
crossterm = "0.27"

[target.'cfg(target_os = "linux")'.dependencies]
nvml-wrapper = { version = "0.10.0", optional = true }
Expand All @@ -40,4 +47,12 @@ windows = { version = "0.48", features = ["Win32_Graphics_Dxgi"] }
[features]
default = ["metal"]
metal = ["dep:core-graphics", "dep:metal"]
windows = ["dep:dxgi", "winapi"]
windows = ["dep:dxgi", "winapi"]

[[bin]]
name = "gpu-share-vm-manager"
path = "src/main.rs"

[dev-dependencies]
tokio = { version = "1.0", features = ["full"] }
rand = "0.8"
24 changes: 24 additions & 0 deletions src/api/error.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#[derive(Debug)]
pub struct ErrorResponse {
pub code: ErrorNumber,
pub message: String,
}

impl ErrorResponse {
pub fn new(code: ErrorNumber, message: String) -> Self {
Self { code, message }
}
}

impl IntoResponse for ErrorResponse {
fn into_response(self) -> Response {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({
"error_code": self.code as u32,
"message": self.message
})),
)
.into_response()
}
}
6 changes: 5 additions & 1 deletion src/api/middleware/mod.rs
Original file line number Diff line number Diff line change
@@ -1 +1,5 @@
pub mod rate_limit;
// src/api/middleware/mod.rs
//! This module groups middleware for the API.
//! Currently it only re-exports the rate_limit middleware.

pub mod rate_limit;
45 changes: 45 additions & 0 deletions src/api/middleware/rate_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,52 @@ impl fmt::Display for RateLimitExceeded {
write!(f, "Rate limit exceeded")
}
}
#[derive(Clone)]
pub struct RateLimit<T> {
inner: T,
}

// wrapper for RateLimitLayer
#[derive(Clone)]
pub struct CustomRateLimitLayer {
rate: u64,
per: Duration,
inner: RateLimitLayer,
}

impl CustomRateLimitLayer {
pub fn new(rate: u64, per: Duration) -> Self {
Self {
rate,
per,
inner: RateLimitLayer::new(rate, per),
}
}

pub fn get_rate(&self) -> u64 {
self.rate
}

pub fn get_per(&self) -> Duration {
self.per
}

pub fn into_inner(self) -> RateLimitLayer {
self.inner
}
}

impl From<RateLimitLayer> for CustomRateLimitLayer {
fn from(_layer: RateLimitLayer) -> Self {
Self::new(100, Duration::from_secs(1))
}
}

impl From<CustomRateLimitLayer> for RateLimitLayer {
fn from(custom: CustomRateLimitLayer) -> Self {
custom.into_inner()
}
}

#[cfg(test)]
mod tests {
Expand Down
2 changes: 1 addition & 1 deletion src/api/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
pub mod middleware;
pub mod routes;

pub use routes::{create_router, AppState};
pub use routes::{AppState, create_router};
Loading