1
0
forked from boos/BoosAPI

feat: merge opus comfyui and login updates

This commit is contained in:
boos
2026-05-21 16:32:04 +02:00
parent c16479da7c
commit 9c84759498
105 changed files with 20160 additions and 1034 deletions
+654
View File
@@ -0,0 +1,654 @@
//! 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.
if let Some(entry) = self.containers.read().await.get(&user_id).cloned() {
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.
if let Some(entry) = self.containers.read().await.get(&user_id).cloned() {
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(), self.config.boosapi_base.clone());
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(),
}
}
}
/// 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);
}
}
+120
View File
@@ -0,0 +1,120 @@
//! 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>>;
}
+288
View File
@@ -0,0 +1,288 @@
//! 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
))
// 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);
}
}
+4
View File
@@ -1,9 +1,13 @@
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 route_rebuild;
pub mod shared;
pub mod update_center;
+174 -4
View File
@@ -13,6 +13,14 @@
//! marked `manual_override = 1` are NEVER touched (created, updated, deleted).
//! Other channels are reconciled to match the current model availability:
//! missing combos get inserted, stale combos get hard-deleted.
//! - After channels are reconciled, a route orphan sweep deletes
//! `token_routes` rows that (a) no active account advertises any more,
//! (b) have no surviving channels (so no manual_override channel pins
//! them), and (c) have empty `display_name` / `display_icon` /
//! `model_mapping` (so no operator has touched the route's display
//! metadata). Without this sweep, `/api/routes/summary` keeps surfacing
//! phantom routes for models nothing can serve, and the user-visible
//! "rebuild" button appears to do nothing.
//! - Disabled site models (`site_disabled_models`) are filtered out before the
//! write. So is anything from accounts whose status isn't `active`, or models
//! where `model_availability.available != 1`.
@@ -30,6 +38,7 @@ use std::collections::{HashMap, HashSet};
pub struct RebuildSummary {
pub created_routes: i64,
pub reused_routes: i64,
pub deleted_routes: i64,
pub created_channels: i64,
pub deleted_channels: i64,
pub skipped_manual_channels: i64,
@@ -255,6 +264,50 @@ pub(crate) async fn rebuild_routes_inner(
}
}
// 6. Route orphan sweep. After channels for "no longer routable" models
// have been deleted above, the parent `token_routes` row can also
// safely go away — but only if it's a pure auto-managed leftover.
// We preserve a route when ANY of these is true:
// - the model is still in `by_model` (an active account advertises it);
// - the route still has at least one channel left, which after the
// channel sweep above can only be a `manual_override = 1` channel;
// - an operator has touched the route's display metadata
// (`display_name`, `display_icon`, `model_mapping`).
// Without this, every route ever created sticks around forever — the
// `/api/routes/summary` endpoint then surfaces phantom routes for
// models nothing can serve, and rebuild "does nothing" because the
// main loop only visits models that are in `by_model`.
let route_orphan_candidates: Vec<(i64, Option<String>, Option<String>, Option<String>)> =
sqlx::query_as(
"SELECT tr.id, tr.display_name, tr.display_icon, tr.model_mapping
FROM token_routes tr
WHERE NOT EXISTS (
SELECT 1 FROM route_channels rc WHERE rc.route_id = tr.id
)",
)
.fetch_all(&mut *tx)
.await?;
for (route_id, display_name, display_icon, model_mapping) in route_orphan_candidates {
if processed_route_ids.contains(&route_id) {
// Still backed by an active model — leave it, even if it currently
// has no channels (channels will be re-inserted on next refresh).
continue;
}
let touched_by_operator = display_name.as_deref().map(str::trim).unwrap_or("").is_empty()
== false
|| display_icon.as_deref().map(str::trim).unwrap_or("").is_empty() == false
|| model_mapping.as_deref().map(str::trim).unwrap_or("").is_empty() == false;
if touched_by_operator {
continue;
}
sqlx::query("DELETE FROM token_routes WHERE id = ?")
.bind(route_id)
.execute(&mut *tx)
.await?;
summary.deleted_routes += 1;
}
tx.commit().await?;
Ok(summary)
}
@@ -296,6 +349,9 @@ mod tests {
"CREATE TABLE token_routes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
model_pattern TEXT NOT NULL UNIQUE,
model_mapping TEXT,
display_name TEXT,
display_icon TEXT,
routing_strategy TEXT NOT NULL DEFAULT 'weighted',
route_mode TEXT NOT NULL DEFAULT 'pattern',
enabled INTEGER NOT NULL DEFAULT 1
@@ -387,14 +443,23 @@ mod tests {
summary.deleted_channels, 1,
"summary should report exactly one deletion"
);
// The route row itself stays — operators may have edited display
// metadata on it, and the next time an active account advertises
// `gpt-x` we want to reuse it. The rebuild never deletes routes.
// The route row also goes — it's a pure auto-managed leftover (no
// surviving channels, no operator-edited display metadata, no active
// account advertising the model). Leaving it behind is what causes
// `/api/routes/summary` to surface phantom routes for models nothing
// can serve.
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(route_count, 1, "token_routes row should be reused, not dropped");
assert_eq!(
route_count, 0,
"pure auto-managed route must be swept once its last channel is gone"
);
assert_eq!(
summary.deleted_routes, 1,
"summary should report the route deletion"
);
}
/// Sanity check: manual_override channels survive an orphan sweep. Even
@@ -423,6 +488,10 @@ mod tests {
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
assert_eq!(summary.deleted_channels, 0, "manual channel must survive");
assert_eq!(
summary.deleted_routes, 0,
"route is pinned by a manual_override channel and must survive"
);
let manual_count: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM route_channels WHERE manual_override = 1",
@@ -431,5 +500,106 @@ mod tests {
.await
.unwrap();
assert_eq!(manual_count, 1);
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(route_count, 1);
}
/// User-reported bug: empty DB still shows hundreds of routes on the
/// `/routes` page, and clicking "rebuild" does nothing. Reproduces by
/// pre-seeding `token_routes` rows that no account / availability row
/// references — exactly the state we found in the live `boosapi.db`.
/// The rebuild must hard-delete those phantom routes.
#[tokio::test]
async fn phantom_routes_with_no_accounts_are_swept_on_rebuild() {
let pool = fixture_pool().await;
// No accounts, no availability — but three left-behind auto-managed
// routes. Mirror the production layout: empty display_name /
// display_icon / model_mapping.
sqlx::query(
"INSERT INTO token_routes (id, model_pattern) VALUES
(1, 'gpt-4'),
(2, 'claude-opus-4.5'),
(3, 'glm-5')",
)
.execute(&pool)
.await
.unwrap();
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
route_count, 0,
"all phantom routes must be swept when no account advertises any model"
);
assert_eq!(summary.deleted_routes, 3, "summary reports the three sweeps");
assert_eq!(summary.created_routes, 0);
assert_eq!(summary.created_channels, 0);
}
/// Operator-edited routes survive the sweep. We probe each of the three
/// "touched" columns (`display_name`, `display_icon`, `model_mapping`)
/// independently so an accidental refactor that drops a check in the
/// disjunction gets caught.
#[tokio::test]
async fn route_with_operator_metadata_is_preserved() {
for (label, sql) in [
(
"display_name",
"INSERT INTO token_routes (id, model_pattern, display_name) VALUES (1, 'gpt-4', 'Renamed by ops')",
),
(
"display_icon",
"INSERT INTO token_routes (id, model_pattern, display_icon) VALUES (1, 'gpt-4', 'openai-color')",
),
(
"model_mapping",
"INSERT INTO token_routes (id, model_pattern, model_mapping) VALUES (1, 'gpt-4', '{\"alias\":\"gpt-4o\"}')",
),
] {
let pool = fixture_pool().await;
sqlx::query(sql).execute(&pool).await.unwrap();
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
assert_eq!(
summary.deleted_routes, 0,
"operator-edited {label} must pin the route through a sweep"
);
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(route_count, 1, "route must survive ({label} edited)");
}
}
/// Whitespace-only metadata is not an "edit" — `' '` should still get
/// swept. Otherwise a stray space saved by the admin form would
/// silently keep dead routes around forever.
#[tokio::test]
async fn whitespace_only_metadata_does_not_pin_a_route() {
let pool = fixture_pool().await;
sqlx::query(
"INSERT INTO token_routes (id, model_pattern, display_name, display_icon, model_mapping)
VALUES (1, 'gpt-4', ' ', '', '\t')",
)
.execute(&pool)
.await
.unwrap();
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
assert_eq!(summary.deleted_routes, 1);
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(route_count, 0);
}
}