boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add workflow studio and postgres support

Browse Source
This commit is contained in:
boos
2026-05-25 16:03:07 +08:00
parent 7cdb02ef1d
commit cb35d5fc6e
96 changed files with 12414 additions and 15548 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/services/comfyui_pool.rs
-672
View File
@@ -1,672 +0,0 @@
//! Per-user ComfyUI container pool.
//!
//! Owns the in-memory map `user_id -> PoolEntry`, allocates ports from the
//! 8200..=8299 range, enforces the 30-container ceiling, runs the readiness
//! probe, and reaps idle containers.
//!
//! All Docker/K3S specifics are encapsulated by the [`ContainerRuntime`]
//! trait; this file MUST stay backend-agnostic.
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use parking_lot::Mutex;
use tokio::sync::{Mutex as AsyncMutex, RwLock};
use tracing::{debug, error, info, warn};
use super::container_runtime::{
BindMount, ContainerRuntime, ContainerSpec, ResourceLimits, RuntimeError, TmpfsMount,
};
/// Maximum concurrent containers (§design Decision 2: 30 × 2 GB ≈ 60 GB).
pub const DEFAULT_MAX_CONCURRENT: usize = 30;
/// Idle timeout before the reaper stops a container.
pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(60 * 60);
/// Host port range for per-user containers.
pub const PORT_RANGE_START: u16 = 8200;
pub const PORT_RANGE_END: u16 = 8299;
/// Time we wait for a freshly started container to answer `/system_stats`.
pub const READINESS_TIMEOUT: Duration = Duration::from_secs(30);
/// How often the reaper task wakes up.
pub const REAPER_TICK: Duration = Duration::from_secs(60);
/// Image tag we run for per-user tenants.
pub const DEFAULT_IMAGE: &str = "boosapi-comfyui:multi-tenant";
/// Host directory under which per-user volumes live.
pub const DEFAULT_DATA_ROOT: &str = "/data/comfyui-users";
/// Lifecycle state visible to callers (frontend renders loading vs ready).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContainerState {
Starting,
Ready,
Stopped,
Failed,
}
impl ContainerState {
pub fn as_str(self) -> &'static str {
match self {
ContainerState::Starting => "starting",
ContainerState::Ready => "ready",
ContainerState::Stopped => "stopped",
ContainerState::Failed => "failed",
}
}
}
/// One running (or starting) per-user container.
#[derive(Debug, Clone)]
pub struct PoolEntry {
pub user_id: i64,
pub container_name: String,
pub host_port: u16,
pub state: ContainerState,
pub last_used_at: Instant,
pub active_ws: u32,
pub last_error: Option<String>,
}
#[derive(thiserror::Error, Debug)]
pub enum PoolError {
#[error("pool at capacity ({0} containers running)")]
AtCapacity(usize),
#[error("no free port in pool range")]
OutOfPorts,
#[error("container failed to become ready in {0:?}")]
StartTimeout(Duration),
#[error("runtime error: {0}")]
Runtime(#[from] RuntimeError),
#[error("internal: {0}")]
Internal(String),
}
pub type PoolResult<T> = Result<T, PoolError>;
/// Tunable knobs the binary wires up at boot.
#[derive(Debug, Clone)]
pub struct PoolConfig {
pub image: String,
pub max_concurrent: usize,
pub idle_timeout: Duration,
pub data_root: PathBuf,
pub boosapi_base: String,
pub builtin_key: String,
}
impl PoolConfig {
pub fn from_env(boosapi_base: String, builtin_key: String) -> Self {
let image = std::env::var("BOOSAPI_COMFY_IMAGE").unwrap_or_else(|_| DEFAULT_IMAGE.into());
let max_concurrent: usize = std::env::var("BOOSAPI_COMFY_MAX")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(DEFAULT_MAX_CONCURRENT);
let idle_secs: u64 = std::env::var("BOOSAPI_COMFY_IDLE_SECS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(DEFAULT_IDLE_TIMEOUT.as_secs());
let data_root = std::env::var("BOOSAPI_COMFY_DATA_ROOT")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from(DEFAULT_DATA_ROOT));
Self {
image,
max_concurrent,
idle_timeout: Duration::from_secs(idle_secs),
data_root,
boosapi_base,
builtin_key,
}
}
}
/// Thread-safe allocator over `PORT_RANGE_START..=PORT_RANGE_END`.
struct PortAllocator {
in_use: Mutex<HashSet<u16>>,
}
impl PortAllocator {
fn new() -> Self {
Self {
in_use: Mutex::new(HashSet::new()),
}
}
fn claim(&self, port: u16) {
self.in_use.lock().insert(port);
}
fn allocate(&self) -> Option<u16> {
let mut g = self.in_use.lock();
for p in PORT_RANGE_START..=PORT_RANGE_END {
if !g.contains(&p) {
g.insert(p);
return Some(p);
}
}
None
}
fn release(&self, port: u16) {
self.in_use.lock().remove(&port);
}
}
/// Per-user start lock. Prevents two concurrent calls to
/// `get_or_start(7)` from spawning two containers for user 7.
type StartLock = Arc<AsyncMutex<()>>;
/// The pool itself.
pub struct ComfyUIPool {
runtime: Arc<dyn ContainerRuntime>,
config: PoolConfig,
containers: RwLock<HashMap<i64, PoolEntry>>,
start_locks: Mutex<HashMap<i64, StartLock>>,
ports: PortAllocator,
http: reqwest::Client,
}
impl ComfyUIPool {
pub fn new(runtime: Arc<dyn ContainerRuntime>, config: PoolConfig) -> Arc<Self> {
Arc::new(Self {
runtime,
config,
containers: RwLock::new(HashMap::new()),
start_locks: Mutex::new(HashMap::new()),
ports: PortAllocator::new(),
http: reqwest::Client::builder()
.timeout(Duration::from_secs(5))
.build()
.expect("reqwest client"),
})
}
pub fn config(&self) -> &PoolConfig {
&self.config
}
/// Start the user's container if it isn't already running and return a
/// handle. Idempotent.
pub async fn get_or_start(&self, user_id: i64) -> PoolResult<PoolEntry> {
// Fast path: already ready.
let existing = {
let map = self.containers.read().await;
map.get(&user_id).cloned()
};
if let Some(entry) = existing {
if entry.state == ContainerState::Ready {
self.touch_inner(user_id).await;
return Ok(entry);
}
}
// Per-user start lock — only one task gets to spawn per user.
let lock = self
.start_locks
.lock()
.entry(user_id)
.or_insert_with(|| Arc::new(AsyncMutex::new(())))
.clone();
let _g = lock.lock().await;
// Re-check after acquiring the lock; another task may have done the work.
let existing = {
let map = self.containers.read().await;
map.get(&user_id).cloned()
};
if let Some(entry) = existing {
if entry.state == ContainerState::Ready {
self.touch_inner(user_id).await;
return Ok(entry);
}
if entry.state == ContainerState::Starting {
// Probe ourselves — the previous winner may have died mid-start.
return self.probe_until_ready(user_id, entry).await;
}
}
// Capacity check.
{
let map = self.containers.read().await;
let live = map
.values()
.filter(|e| matches!(e.state, ContainerState::Starting | ContainerState::Ready))
.count();
if live >= self.config.max_concurrent {
return Err(PoolError::AtCapacity(live));
}
}
// Allocate port.
let port = self.ports.allocate().ok_or(PoolError::OutOfPorts)?;
let name = format!("comfyui-user-{}", user_id);
// Insert a Starting marker so concurrent callers see it.
{
let mut map = self.containers.write().await;
map.insert(
user_id,
PoolEntry {
user_id,
container_name: name.clone(),
host_port: port,
state: ContainerState::Starting,
last_used_at: Instant::now(),
active_ws: 0,
last_error: None,
},
);
}
let spec = self.build_spec(user_id, &name, port);
match self.runtime.start(spec).await {
Ok(handle) => {
debug!(user_id, container = %handle.name, port = handle.host_port, "container started");
}
Err(e) => {
self.ports.release(port);
let mut map = self.containers.write().await;
if let Some(entry) = map.get_mut(&user_id) {
entry.state = ContainerState::Failed;
entry.last_error = Some(e.to_string());
}
return Err(PoolError::Runtime(e));
}
}
let entry_now = self
.containers
.read()
.await
.get(&user_id)
.cloned()
.ok_or_else(|| PoolError::Internal("entry vanished mid-start".into()))?;
self.probe_until_ready(user_id, entry_now).await
}
/// Wait for the container to answer `/system_stats`, then mark Ready.
async fn probe_until_ready(&self, user_id: i64, entry: PoolEntry) -> PoolResult<PoolEntry> {
let url = format!("http://127.0.0.1:{}/system_stats", entry.host_port);
let deadline = Instant::now() + READINESS_TIMEOUT;
while Instant::now() < deadline {
match self.http.get(&url).send().await {
Ok(r) if r.status().is_success() => {
let mut map = self.containers.write().await;
if let Some(e) = map.get_mut(&user_id) {
e.state = ContainerState::Ready;
e.last_used_at = Instant::now();
e.last_error = None;
return Ok(e.clone());
}
return Err(PoolError::Internal("entry removed during probe".into()));
}
_ => {
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
}
// Timeout. Mark Failed and let the next call retry.
let mut map = self.containers.write().await;
if let Some(e) = map.get_mut(&user_id) {
e.state = ContainerState::Failed;
e.last_error = Some("readiness probe timed out".into());
}
// Best-effort stop so the slot frees on the next reaper tick.
let name = entry.container_name.clone();
let runtime = self.runtime.clone();
tokio::spawn(async move {
let _ = runtime.stop(&name).await;
});
Err(PoolError::StartTimeout(READINESS_TIMEOUT))
}
/// Stop the user's container, release its port.
pub async fn stop(&self, user_id: i64) -> PoolResult<()> {
let entry = {
let mut map = self.containers.write().await;
map.remove(&user_id)
};
let Some(entry) = entry else {
return Ok(());
};
self.ports.release(entry.host_port);
self.runtime.stop(&entry.container_name).await?;
Ok(())
}
/// Current state for a user (or `Stopped` if no entry).
pub async fn state(&self, user_id: i64) -> ContainerState {
self.containers
.read()
.await
.get(&user_id)
.map(|e| e.state)
.unwrap_or(ContainerState::Stopped)
}
/// Snapshot of the user's pool entry, if any.
pub async fn entry(&self, user_id: i64) -> Option<PoolEntry> {
self.containers.read().await.get(&user_id).cloned()
}
/// Bump `last_used_at`. Cheap; called on every proxied frame.
pub async fn touch(&self, user_id: i64) {
self.touch_inner(user_id).await
}
async fn touch_inner(&self, user_id: i64) {
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
e.last_used_at = Instant::now();
}
}
/// Increment / decrement the active-WS counter for a user.
pub async fn ws_open(&self, user_id: i64) {
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
e.active_ws = e.active_ws.saturating_add(1);
e.last_used_at = Instant::now();
}
}
pub async fn ws_close(&self, user_id: i64) {
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
e.active_ws = e.active_ws.saturating_sub(1);
e.last_used_at = Instant::now();
}
}
/// Reconcile against the runtime at startup. Adopts existing
/// `comfyui-user-*` containers so a Rust restart doesn't orphan them.
pub async fn reconcile_on_startup(&self) {
let listings = match self.runtime.list_user_containers().await {
Ok(v) => v,
Err(e) => {
warn!(error = %e, "could not reconcile existing containers; assuming none");
return;
}
};
let mut map = self.containers.write().await;
for listing in listings {
// Names look like `comfyui-user-7`.
let Some(suffix) = listing.name.strip_prefix("comfyui-user-") else {
continue;
};
let Ok(user_id) = suffix.parse::<i64>() else {
continue;
};
self.ports.claim(listing.host_port);
map.insert(
user_id,
PoolEntry {
user_id,
container_name: listing.name.clone(),
host_port: listing.host_port,
state: ContainerState::Ready, // assume ready; first request re-probes
last_used_at: Instant::now(),
active_ws: 0,
last_error: None,
},
);
info!(user_id, container = %listing.name, port = listing.host_port, "adopted existing container");
}
}
/// Stop every container we manage (graceful-shutdown hook).
pub async fn shutdown_all(&self) {
let names: Vec<String> = {
let map = self.containers.read().await;
map.values().map(|e| e.container_name.clone()).collect()
};
for name in names {
if let Err(e) = self.runtime.stop(&name).await {
warn!(container = %name, error = %e, "shutdown stop failed");
}
}
self.containers.write().await.clear();
}
/// Build the [`ContainerSpec`] for a given user.
fn build_spec(&self, user_id: i64, name: &str, host_port: u16) -> ContainerSpec {
let user_root = self.config.data_root.join(user_id.to_string());
// Subdirs we mount in. Created on demand by the host operator (or
// the runtime can `mkdir -p` them — Docker `bind` will fail if the
// path is missing).
let mk = |sub: &str| BindMount {
host_path: user_root.join(sub),
container_path: format!("/workspace/ComfyUI/{}", sub),
read_only: false,
};
// Best-effort mkdir so a fresh user just works. Failures here surface
// through the Docker error message later.
for sub in ["input", "output", "temp", "user", "custom_nodes_user"] {
let _ = std::fs::create_dir_all(user_root.join(sub));
}
let mut env = std::collections::HashMap::new();
env.insert(
"BOOSAPI_BUILTIN_KEY".into(),
self.config.builtin_key.clone(),
);
env.insert(
"BOOSAPI_BASE".into(),
container_reachable_boosapi_base(&self.config.boosapi_base),
);
env.insert("BOOSAPI_KEY".into(), self.config.builtin_key.clone());
ContainerSpec {
image: self.config.image.clone(),
name: name.to_string(),
host_port,
env,
binds: vec![
mk("input"),
mk("output"),
mk("temp"),
mk("user"),
mk("custom_nodes_user"),
],
tmpfs: vec![TmpfsMount {
container_path: "/tmp".into(),
size_bytes: 512 * 1024 * 1024,
}],
limits: ResourceLimits::default(),
}
}
}
fn container_reachable_boosapi_base(base: &str) -> String {
if let Ok(override_base) = std::env::var("BOOSAPI_COMFY_CONTAINER_BASE") {
let trimmed = override_base.trim();
if !trimmed.is_empty() {
return trimmed.trim_end_matches('/').to_string();
}
}
let trimmed = base.trim().trim_end_matches('/');
for local_host in ["127.0.0.1", "localhost"] {
let needle = format!("//{}:", local_host);
if let Some((prefix, rest)) = trimmed.split_once(&needle) {
return format!("{}//host.docker.internal:{}", prefix, rest);
}
}
trimmed.to_string()
}
/// Spawn the idle-reaper task. Lives for the life of the process.
pub fn spawn_reaper(pool: Arc<ComfyUIPool>) {
tokio::spawn(async move {
let mut tick = tokio::time::interval(REAPER_TICK);
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tick.tick().await;
reap_once(&pool).await;
}
});
}
async fn reap_once(pool: &Arc<ComfyUIPool>) {
let now = Instant::now();
let idle = pool.config.idle_timeout;
let candidates: Vec<i64> = {
let map = pool.containers.read().await;
map.values()
.filter(|e| {
e.active_ws == 0
&& now.saturating_duration_since(e.last_used_at) > idle
&& matches!(e.state, ContainerState::Ready | ContainerState::Failed)
})
.map(|e| e.user_id)
.collect()
};
for user_id in candidates {
info!(user_id, "reaping idle comfyui container");
if let Err(e) = pool.stop(user_id).await {
error!(user_id, error = %e, "reap failed");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::services::container_runtime::{ContainerHandle, ContainerListing, RuntimeResult};
/// In-memory runtime that succeeds and pretends `/system_stats` is up.
/// Real readiness probe is tested separately via the live Docker setup.
struct MockRuntime {
starts: AtomicUsize,
stops: AtomicUsize,
}
impl MockRuntime {
fn new() -> Self {
Self {
starts: AtomicUsize::new(0),
stops: AtomicUsize::new(0),
}
}
}
#[async_trait]
impl ContainerRuntime for MockRuntime {
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle> {
self.starts.fetch_add(1, Ordering::SeqCst);
Ok(ContainerHandle {
id: format!("mock-{}", spec.name),
name: spec.name,
host_port: spec.host_port,
})
}
async fn stop(&self, _name: &str) -> RuntimeResult<()> {
self.stops.fetch_add(1, Ordering::SeqCst);
Ok(())
}
async fn is_alive(&self, _name: &str) -> RuntimeResult<bool> {
Ok(true)
}
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>> {
Ok(vec![])
}
}
fn cfg() -> PoolConfig {
PoolConfig {
image: "test".into(),
max_concurrent: 2,
idle_timeout: Duration::from_secs(60),
data_root: std::env::temp_dir().join("boosapi-pool-test"),
boosapi_base: "http://127.0.0.1:9999".into(),
builtin_key: "sk-test".into(),
}
}
#[tokio::test]
async fn capacity_cap_enforced() {
let runtime = Arc::new(MockRuntime::new());
let pool = ComfyUIPool::new(runtime.clone(), cfg());
// Inject 2 containers manually so we don't have to satisfy the
// /system_stats probe (this test isolates the cap check).
{
let mut map = pool.containers.write().await;
map.insert(
1,
PoolEntry {
user_id: 1,
container_name: "comfyui-user-1".into(),
host_port: 8200,
state: ContainerState::Ready,
last_used_at: Instant::now(),
active_ws: 0,
last_error: None,
},
);
map.insert(
2,
PoolEntry {
user_id: 2,
container_name: "comfyui-user-2".into(),
host_port: 8201,
state: ContainerState::Ready,
last_used_at: Instant::now(),
active_ws: 0,
last_error: None,
},
);
}
let err = pool.get_or_start(3).await.unwrap_err();
assert!(matches!(err, PoolError::AtCapacity(2)), "got: {:?}", err);
}
#[tokio::test]
async fn stop_releases_port() {
let runtime = Arc::new(MockRuntime::new());
let pool = ComfyUIPool::new(runtime.clone(), cfg());
{
let mut map = pool.containers.write().await;
map.insert(
7,
PoolEntry {
user_id: 7,
container_name: "comfyui-user-7".into(),
host_port: 8205,
state: ContainerState::Ready,
last_used_at: Instant::now(),
active_ws: 0,
last_error: None,
},
);
}
pool.ports.claim(8205);
pool.stop(7).await.unwrap();
assert_eq!(runtime.stops.load(Ordering::SeqCst), 1);
// Port is back in the free pool.
let next = pool.ports.allocate().unwrap();
assert_eq!(next, PORT_RANGE_START);
}
#[test]
fn port_allocator_assigns_in_order_and_releases() {
let p = PortAllocator::new();
let a = p.allocate().unwrap();
let b = p.allocate().unwrap();
assert_eq!(a, PORT_RANGE_START);
assert_eq!(b, PORT_RANGE_START + 1);
p.release(a);
let c = p.allocate().unwrap();
assert_eq!(c, PORT_RANGE_START);
}
}
src/services/container_runtime.rs
-120
View File
@@ -1,120 +0,0 @@
//! Backend-agnostic container lifecycle abstraction.
//!
//! v1 ships [`crate::services::docker_runtime::DockerRuntime`] against the
//! local Docker daemon. v2 will add a K3S implementation behind the same
//! trait so the pool, proxy, and route handlers do NOT have to change.
//!
//! Callers (the pool, in particular) MUST go through this trait — never
//! import `bollard`, the `docker` CLI, or `kube` types directly.
use async_trait::async_trait;
use std::collections::HashMap;
use std::path::PathBuf;
/// Resource and security caps for a container.
///
/// Mirrors §design Decision 5 verbatim: 2 CPU, 2 GB RAM, 2 GB swap,
/// `pids_limit=512`, read-only rootfs, drop all caps, `no-new-privileges`.
#[derive(Debug, Clone)]
pub struct ResourceLimits {
pub cpus: f64,
pub memory_bytes: u64,
pub memory_swap_bytes: u64,
pub pids_limit: u32,
pub read_only_rootfs: bool,
pub cap_drop_all: bool,
pub no_new_privileges: bool,
}
impl Default for ResourceLimits {
fn default() -> Self {
Self {
cpus: 2.0,
memory_bytes: 2 * 1024 * 1024 * 1024,
memory_swap_bytes: 2 * 1024 * 1024 * 1024,
pids_limit: 512,
read_only_rootfs: true,
cap_drop_all: true,
no_new_privileges: true,
}
}
}
/// One bind mount: `host_path -> container_path`.
#[derive(Debug, Clone)]
pub struct BindMount {
pub host_path: PathBuf,
pub container_path: String,
pub read_only: bool,
}
/// Tmpfs mount (used to satisfy a few writable paths under read-only rootfs).
#[derive(Debug, Clone)]
pub struct TmpfsMount {
pub container_path: String,
pub size_bytes: u64,
}
/// A complete container start request.
#[derive(Debug, Clone)]
pub struct ContainerSpec {
pub image: String,
pub name: String,
/// Host port to bind to. The container's internal port is fixed at 8188.
pub host_port: u16,
pub env: HashMap<String, String>,
pub binds: Vec<BindMount>,
pub tmpfs: Vec<TmpfsMount>,
pub limits: ResourceLimits,
}
/// Result of a successful start: identifies the container and tells the
/// caller which loopback port to forward traffic to.
#[derive(Debug, Clone)]
pub struct ContainerHandle {
/// Backend-specific id (Docker container id, k8s pod name, ...).
pub id: String,
/// Container name we asked for, e.g. `comfyui-user-7`.
pub name: String,
/// Loopback host port the runtime mapped for us.
pub host_port: u16,
}
/// Listing entry returned by [`ContainerRuntime::list_user_containers`].
#[derive(Debug, Clone)]
pub struct ContainerListing {
pub id: String,
pub name: String,
pub host_port: u16,
}
#[derive(thiserror::Error, Debug)]
pub enum RuntimeError {
#[error("backend error: {0}")]
Backend(String),
#[error("container not found: {0}")]
NotFound(String),
}
pub type RuntimeResult<T> = Result<T, RuntimeError>;
/// Backend-agnostic container lifecycle.
#[async_trait]
pub trait ContainerRuntime: Send + Sync + 'static {
/// Start a container and return a handle once it is running.
///
/// The runtime is NOT required to wait for the application inside the
/// container to be ready — that probe lives in the pool.
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle>;
/// Stop and remove a container by name. Idempotent: removing a
/// nonexistent container returns `Ok`.
async fn stop(&self, name: &str) -> RuntimeResult<()>;
/// Check whether a container is running.
async fn is_alive(&self, name: &str) -> RuntimeResult<bool>;
/// List every running container whose name starts with `comfyui-user-`.
/// Used at API startup to reconstruct the pool's in-memory map.
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>>;
}
src/services/docker_runtime.rs
-291
View File
@@ -1,291 +0,0 @@
//! Docker implementation of [`ContainerRuntime`].
//!
//! We shell out to the `docker` CLI rather than pull in a Rust Docker SDK,
//! to keep the dependency footprint small. The CLI is universally
//! available on the hosts we ship to, and the surface area we need
//! (`run / stop / rm / inspect / ps`) is tiny.
use async_trait::async_trait;
use std::process::Stdio;
use tokio::process::Command;
use tracing::{debug, warn};
use super::container_runtime::{
ContainerHandle, ContainerListing, ContainerRuntime, ContainerSpec, RuntimeError, RuntimeResult,
};
/// Internal port the ComfyUI process listens on inside the container.
/// Hard-coded to match the image's ENTRYPOINT.
const CONTAINER_INTERNAL_PORT: u16 = 8188;
/// Label we stamp onto every container we manage. Used by
/// `list_user_containers` and as a sanity-check on `stop`.
pub const MANAGED_LABEL: &str = "boosapi.comfy.tenant=1";
/// Default Docker runtime. Stateless — every method shells out fresh.
pub struct DockerRuntime {
docker_bin: String,
}
impl DockerRuntime {
pub fn new() -> Self {
Self {
docker_bin: std::env::var("BOOSAPI_DOCKER_BIN").unwrap_or_else(|_| "docker".into()),
}
}
fn cmd(&self) -> Command {
let mut c = Command::new(&self.docker_bin);
c.stdin(Stdio::null());
c
}
}
impl Default for DockerRuntime {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ContainerRuntime for DockerRuntime {
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle> {
// Best-effort: kill any stale container with the same name.
// `docker rm -f <name>` is a no-op if the name is free.
let _ = self
.cmd()
.args(["rm", "-f", &spec.name])
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await;
let mut cmd = self.cmd();
cmd.arg("run")
.arg("-d")
.arg("--rm")
.arg("--name")
.arg(&spec.name)
.arg("--label")
.arg(MANAGED_LABEL)
.arg("-p")
.arg(format!(
"127.0.0.1:{}:{}",
spec.host_port, CONTAINER_INTERNAL_PORT
))
.arg("--add-host")
.arg("host.docker.internal:host-gateway")
// Resource limits (§design Decision 5)
.arg(format!("--cpus={}", spec.limits.cpus))
.arg(format!("--memory={}", spec.limits.memory_bytes))
.arg(format!("--memory-swap={}", spec.limits.memory_swap_bytes))
.arg(format!("--pids-limit={}", spec.limits.pids_limit));
if spec.limits.read_only_rootfs {
cmd.arg("--read-only");
}
if spec.limits.cap_drop_all {
cmd.args(["--cap-drop", "ALL"]);
}
if spec.limits.no_new_privileges {
cmd.args(["--security-opt", "no-new-privileges"]);
}
for (k, v) in &spec.env {
cmd.arg("-e").arg(format!("{}={}", k, v));
}
for bind in &spec.binds {
// Using `--mount` over `-v` so the syntax is unambiguous and
// we get a clear error if the host path is wrong.
let mut s = format!(
"type=bind,source={},target={}",
bind.host_path.display(),
bind.container_path
);
if bind.read_only {
s.push_str(",readonly");
}
cmd.arg("--mount").arg(s);
}
for tmp in &spec.tmpfs {
cmd.arg("--tmpfs")
.arg(format!("{}:size={}", tmp.container_path, tmp.size_bytes));
}
cmd.arg(&spec.image);
debug!(?cmd, "starting comfyui container");
let output = cmd.output().await.map_err(io_to_runtime)?;
if !output.status.success() {
return Err(RuntimeError::Backend(format!(
"docker run failed (exit {}): {}",
output.status,
String::from_utf8_lossy(&output.stderr).trim()
)));
}
let id = String::from_utf8_lossy(&output.stdout).trim().to_string();
if id.is_empty() {
return Err(RuntimeError::Backend(
"docker run returned empty container id".into(),
));
}
Ok(ContainerHandle {
id,
name: spec.name,
host_port: spec.host_port,
})
}
async fn stop(&self, name: &str) -> RuntimeResult<()> {
// `docker stop` flushes nicely; `--rm` from start removes it after.
// If it's already gone, treat as success.
let output = self
.cmd()
.args(["stop", "-t", "10", name])
.output()
.await
.map_err(io_to_runtime)?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
if stderr.contains("No such container") {
return Ok(());
}
warn!(name, stderr = stderr.trim(), "docker stop failed");
// Fall through to rm -f as a last resort.
}
let _ = self
.cmd()
.args(["rm", "-f", name])
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await;
Ok(())
}
async fn is_alive(&self, name: &str) -> RuntimeResult<bool> {
let output = self
.cmd()
.args(["inspect", "-f", "{{.State.Running}}", name])
.output()
.await
.map_err(io_to_runtime)?;
if !output.status.success() {
// `docker inspect` returns 1 if the container is unknown.
return Ok(false);
}
Ok(String::from_utf8_lossy(&output.stdout).trim() == "true")
}
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>> {
// `--format {{json .}}` gives one JSON object per line. We pull the
// fields we need without depending on a JSON-deserialize crate
// schema for Docker's somewhat-fluid output.
let output = self
.cmd()
.args([
"ps",
"--filter",
"name=comfyui-user-",
"--filter",
&format!("label={}", MANAGED_LABEL),
"--format",
"{{.Names}}\t{{.ID}}\t{{.Ports}}",
])
.output()
.await
.map_err(io_to_runtime)?;
if !output.status.success() {
return Err(RuntimeError::Backend(format!(
"docker ps failed: {}",
String::from_utf8_lossy(&output.stderr).trim()
)));
}
let stdout = String::from_utf8_lossy(&output.stdout);
let mut out = Vec::new();
for line in stdout.lines() {
let mut parts = line.splitn(3, '\t');
let name = parts.next().unwrap_or("").trim().to_string();
let id = parts.next().unwrap_or("").trim().to_string();
let ports = parts.next().unwrap_or("").trim();
if name.is_empty() || id.is_empty() {
continue;
}
let host_port = match parse_host_port(ports) {
Some(p) => p,
None => {
warn!(
name,
ports, "could not parse host port from docker ps; skipping"
);
continue;
}
};
out.push(ContainerListing {
id,
name,
host_port,
});
}
Ok(out)
}
}
/// Pull the host-side port out of a `docker ps` Ports column.
///
/// Examples it must handle:
/// "127.0.0.1:8200->8188/tcp"
/// "127.0.0.1:8200->8188/tcp, 127.0.0.1:8200->8188/udp"
fn parse_host_port(ports: &str) -> Option<u16> {
for spec in ports.split(',') {
let spec = spec.trim();
// Accept `<host>:<port>->...`
let arrow = spec.find("->")?;
let left = &spec[..arrow];
let port_str = left.rsplit(':').next()?;
if let Ok(p) = port_str.parse::<u16>() {
return Some(p);
}
}
None
}
fn io_to_runtime(e: std::io::Error) -> RuntimeError {
RuntimeError::Backend(format!("docker cli io error: {}", e))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_simple_port_mapping() {
assert_eq!(parse_host_port("127.0.0.1:8200->8188/tcp"), Some(8200));
}
#[test]
fn parses_port_mapping_with_extra_entries() {
assert_eq!(
parse_host_port("127.0.0.1:8201->8188/tcp, 127.0.0.1:8201->8188/udp"),
Some(8201)
);
}
#[test]
fn returns_none_for_garbage() {
assert_eq!(parse_host_port(""), None);
assert_eq!(parse_host_port("nope"), None);
}
}
src/services/mod.rs
+1 -3
View File
@@ -1,13 +1,11 @@
pub mod announcement;
pub mod billing;
pub mod builtin_keys;
pub mod comfyui_pool;
pub mod comfyui_process;
pub mod container_runtime;
pub mod cron;
pub mod docker_runtime;
pub mod notification;
pub mod oauth;
pub mod registration_limiter;
pub mod route_rebuild;
pub mod shared;
pub mod update_center;
src/services/registration_limiter.rs
+221
View File
@@ -0,0 +1,221 @@
use std::{
collections::{HashMap, VecDeque},
net::IpAddr,
time::{Duration, Instant},
};
use parking_lot::Mutex;
const IP_WINDOW: Duration = Duration::from_secs(15 * 60);
const EMAIL_WINDOW: Duration = Duration::from_secs(60 * 60);
const GLOBAL_WINDOW: Duration = Duration::from_secs(5 * 60);
const SUCCESS_LOCK_WINDOW: Duration = Duration::from_secs(10 * 60);
const MAX_ATTEMPTS_PER_IP: usize = 8;
const MAX_ATTEMPTS_PER_EMAIL: usize = 3;
const MAX_GLOBAL_ATTEMPTS: usize = 120;
const MAX_SUCCESS_PER_IP: usize = 2;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RegistrationLimitExceeded {
pub retry_after_secs: u64,
pub message: &'static str,
}
#[derive(Default)]
struct LimiterState {
attempts_by_ip: HashMap<String, VecDeque<Instant>>,
attempts_by_email: HashMap<String, VecDeque<Instant>>,
successful_by_ip: HashMap<String, VecDeque<Instant>>,
global_attempts: VecDeque<Instant>,
}
#[derive(Default)]
pub struct RegistrationLimiter {
inner: Mutex<LimiterState>,
}
impl RegistrationLimiter {
pub fn new() -> Self {
Self::default()
}
pub fn check_and_record_attempt(
&self,
ip: Option<IpAddr>,
email: &str,
) -> Result<(), RegistrationLimitExceeded> {
let now = Instant::now();
let ip_key = ip
.map(|addr| addr.to_string())
.unwrap_or_else(|| "unknown".to_string());
let email_key = normalize_email(email);
let mut guard = self.inner.lock();
prune_deque(&mut guard.global_attempts, now, GLOBAL_WINDOW);
if guard.global_attempts.len() >= MAX_GLOBAL_ATTEMPTS {
return Err(limit_error(
&guard.global_attempts,
now,
GLOBAL_WINDOW,
"注册请求过于频繁，请稍后再试。",
));
}
{
let ip_attempts = guard.attempts_by_ip.entry(ip_key.clone()).or_default();
prune_deque(ip_attempts, now, IP_WINDOW);
if ip_attempts.len() >= MAX_ATTEMPTS_PER_IP {
return Err(limit_error(
ip_attempts,
now,
IP_WINDOW,
"当前网络注册过于频繁，请稍后再试。",
));
}
}
{
let ip_successes = guard.successful_by_ip.entry(ip_key.clone()).or_default();
prune_deque(ip_successes, now, SUCCESS_LOCK_WINDOW);
if ip_successes.len() >= MAX_SUCCESS_PER_IP {
return Err(limit_error(
ip_successes,
now,
SUCCESS_LOCK_WINDOW,
"当前网络短时间内已注册多个账号，请稍后再试。",
));
}
}
{
let email_attempts = guard
.attempts_by_email
.entry(email_key.clone())
.or_default();
prune_deque(email_attempts, now, EMAIL_WINDOW);
if email_attempts.len() >= MAX_ATTEMPTS_PER_EMAIL {
return Err(limit_error(
email_attempts,
now,
EMAIL_WINDOW,
"该邮箱注册尝试过于频繁，请稍后再试。",
));
}
}
guard.global_attempts.push_back(now);
guard
.attempts_by_ip
.entry(ip_key)
.or_default()
.push_back(now);
guard
.attempts_by_email
.entry(email_key)
.or_default()
.push_back(now);
Ok(())
}
pub fn record_success(&self, ip: Option<IpAddr>) {
let now = Instant::now();
let ip_key = ip
.map(|addr| addr.to_string())
.unwrap_or_else(|| "unknown".to_string());
let mut guard = self.inner.lock();
let entries = guard.successful_by_ip.entry(ip_key).or_default();
prune_deque(entries, now, SUCCESS_LOCK_WINDOW);
entries.push_back(now);
}
}
fn normalize_email(email: &str) -> String {
email.trim().to_ascii_lowercase()
}
fn prune_deque(entries: &mut VecDeque<Instant>, now: Instant, window: Duration) {
while entries
.front()
.is_some_and(|timestamp| now.duration_since(*timestamp) >= window)
{
entries.pop_front();
}
}
fn limit_error(
entries: &VecDeque<Instant>,
now: Instant,
window: Duration,
message: &'static str,
) -> RegistrationLimitExceeded {
let retry_after_secs = entries
.front()
.map(|first| {
window
.saturating_sub(now.duration_since(*first))
.as_secs()
.max(1)
})
.unwrap_or_else(|| window.as_secs().max(1));
RegistrationLimitExceeded {
retry_after_secs,
message,
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::{IpAddr, Ipv4Addr};
fn ip() -> Option<IpAddr> {
Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
}
#[test]
fn limits_repeated_attempts_from_same_ip() {
let limiter = RegistrationLimiter::new();
for index in 0..MAX_ATTEMPTS_PER_IP {
limiter
.check_and_record_attempt(ip(), &format!("u{index}@example.com"))
.unwrap();
}
let err = limiter
.check_and_record_attempt(ip(), "blocked@example.com")
.unwrap_err();
assert_eq!(err.message, "当前网络注册过于频繁，请稍后再试。");
}
#[test]
fn limits_repeated_attempts_for_same_email_case_insensitive() {
let limiter = RegistrationLimiter::new();
for _ in 0..MAX_ATTEMPTS_PER_EMAIL {
limiter
.check_and_record_attempt(None, "User@Example.com")
.unwrap();
}
let err = limiter
.check_and_record_attempt(None, "user@example.com")
.unwrap_err();
assert_eq!(err.message, "该邮箱注册尝试过于频繁，请稍后再试。");
}
#[test]
fn limits_successful_registrations_from_same_ip() {
let limiter = RegistrationLimiter::new();
for index in 0..MAX_SUCCESS_PER_IP {
limiter
.check_and_record_attempt(ip(), &format!("u{index}@example.com"))
.unwrap();
limiter.record_success(ip());
}
let err = limiter
.check_and_record_attempt(ip(), "extra@example.com")
.unwrap_err();
assert_eq!(err.message, "当前网络短时间内已注册多个账号，请稍后再试。");
}
}
src/services/workspace.rs
+30 -14
View File
@@ -22,6 +22,7 @@ pub struct WorkspaceAssetInput {
#[derive(Debug, Clone)]
pub struct StoredWorkspaceAsset {
pub public_id: String,
pub file_name: Option<String>,
pub file_path: Option<String>,
pub remote_url: Option<String>,
pub status: String,
@@ -56,7 +57,12 @@ fn safe_segment(value: &str) -> String {
}
fn extension_for_mime(mime: &str) -> &'static str {
let mime = mime.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
let mime = mime
.split(';')
.next()
.unwrap_or("")
.trim()
.to_ascii_lowercase();
match mime.as_str() {
"image/jpeg" | "image/jpg" => "jpg",
"image/webp" => "webp",
@@ -156,6 +162,7 @@ async fn insert_asset(
Ok(StoredWorkspaceAsset {
public_id: public_id.to_string(),
file_name,
file_path: file_path_string,
remote_url: input.remote_url.clone(),
status: status.to_string(),
@@ -193,12 +200,13 @@ pub async fn record_remote_asset(
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.unwrap_or_else(|| "application/octet-stream".to_string());
let bytes = response
.bytes()
.await
.map_err(|err| AppError::UpstreamError(format!("download workspace asset: {err}")))?;
let ext = extension_from_url(url).unwrap_or_else(|| extension_for_mime(&mime_type).to_string());
let path = write_asset_file(input.user_id, &input.source, &public_id, &ext, &bytes).await?;
let bytes = response.bytes().await.map_err(|err| {
AppError::UpstreamError(format!("download workspace asset: {err}"))
})?;
let ext = extension_from_url(url)
.unwrap_or_else(|| extension_for_mime(&mime_type).to_string());
let path =
write_asset_file(input.user_id, &input.source, &public_id, &ext, &bytes).await?;
insert_asset(
state,
&input,
@@ -283,17 +291,22 @@ pub async fn record_image_assets_from_response(
prompt: Option<String>,
model: Option<String>,
response: &Value,
) {
) -> Vec<StoredWorkspaceAsset> {
let Some(items) = response.get("data").and_then(Value::as_array) else {
return;
return Vec::new();
};
let mut stored_assets = Vec::new();
for (index, item) in items.iter().enumerate() {
let title = prompt
.as_deref()
.map(|value| value.chars().take(36).collect::<String>())
.filter(|value| !value.trim().is_empty())
.or_else(|| Some(format!("图片产出 #{}", index + 1)));
if let Some(url) = item.get("url").and_then(Value::as_str).filter(|value| !value.trim().is_empty()) {
if let Some(url) = item
.get("url")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
{
let input = WorkspaceAssetInput {
user_id,
source: source.to_string(),
@@ -304,8 +317,9 @@ pub async fn record_image_assets_from_response(
remote_url: Some(url.trim().to_string()),
metadata: json!({ "responseItem": item }),
};
if let Err(err) = record_remote_asset(state, input).await {
tracing::warn!(%err, "failed to archive remote image asset");
match record_remote_asset(state, input).await {
Ok(asset) => stored_assets.push(asset),
Err(err) => tracing::warn!(%err, "failed to archive remote image asset"),
}
} else if let Some(b64) = item
.get("b64_json")
@@ -322,9 +336,11 @@ pub async fn record_image_assets_from_response(
remote_url: None,
metadata: json!({ "responseItem": item }),
};
if let Err(err) = record_base64_asset(state, input, b64).await {
tracing::warn!(%err, "failed to archive base64 image asset");
match record_base64_asset(state, input, b64).await {
Ok(asset) => stored_assets.push(asset),
Err(err) => tracing::warn!(%err, "failed to archive base64 image asset"),
}
}
}
stored_assets
}