Spaces:
Sleeping
Sleeping
| # Copyright (C) 2026 Hengzhe Zhao. All rights reserved. | |
| # Licensed under dual license: AGPL-3.0 (open-source) or commercial. See LICENSE. | |
| """Concurrent user queue for Prefero on shared hosting (e.g. HF Spaces). | |
| Uses a module-level dict (shared across all Streamlit sessions within the | |
| same process) to track active users. Thread-safe via a lock. | |
| Toggle with PREFERO_QUEUE_ENABLED env var ("true" to enable). | |
| Max concurrent users controlled by PREFERO_MAX_CONCURRENT (default 2). | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import threading | |
| import time | |
| import uuid | |
| import streamlit as st | |
| # --------------------------------------------------------------------------- | |
| # Configuration | |
| # --------------------------------------------------------------------------- | |
| _MAX_CONCURRENT = int(os.environ.get("PREFERO_MAX_CONCURRENT", "2")) | |
| _SESSION_TIMEOUT = 1800 # 30 minutes of inactivity → evicted | |
| def _queue_enabled() -> bool: | |
| return os.environ.get("PREFERO_QUEUE_ENABLED", "").lower() == "true" | |
| # --------------------------------------------------------------------------- | |
| # Shared state (module-level, shared across all sessions in one process) | |
| # --------------------------------------------------------------------------- | |
| _lock = threading.Lock() | |
| _active_sessions: dict[str, float] = {} # session_id → last_heartbeat | |
| _session_usernames: dict[str, str] = {} # session_id → username | |
| def _cleanup_stale() -> None: | |
| """Remove sessions that haven't sent a heartbeat recently.""" | |
| now = time.time() | |
| stale = [sid for sid, ts in _active_sessions.items() if now - ts > _SESSION_TIMEOUT] | |
| for sid in stale: | |
| del _active_sessions[sid] | |
| _session_usernames.pop(sid, None) | |
| def _ensure_session_id() -> str: | |
| """Get or create a unique session identifier.""" | |
| if "_queue_session_id" not in st.session_state: | |
| st.session_state["_queue_session_id"] = str(uuid.uuid4()) | |
| return st.session_state["_queue_session_id"] | |
| # --------------------------------------------------------------------------- | |
| # Public API | |
| # --------------------------------------------------------------------------- | |
| def try_enter() -> bool: | |
| """Try to claim a slot. Returns True if the user is admitted.""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _cleanup_stale() | |
| if sid in _active_sessions: | |
| _active_sessions[sid] = time.time() | |
| return True | |
| if len(_active_sessions) < _MAX_CONCURRENT: | |
| _active_sessions[sid] = time.time() | |
| return True | |
| return False | |
| def heartbeat() -> None: | |
| """Refresh the current session's timestamp (call on every page load).""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| if sid in _active_sessions: | |
| _active_sessions[sid] = time.time() | |
| def leave() -> None: | |
| """Release the current session's slot.""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _active_sessions.pop(sid, None) | |
| _session_usernames.pop(sid, None) | |
| def register_username(username: str) -> None: | |
| """Associate the current session with a username.""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _session_usernames[sid] = username | |
| def is_username_active(username: str) -> bool: | |
| """Check if a username is logged in on another active session.""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _cleanup_stale() | |
| for other_sid, uname in _session_usernames.items(): | |
| if uname == username and other_sid != sid: | |
| return True | |
| return False | |
| def force_evict_username(username: str) -> None: | |
| """Evict all other sessions using this username so the caller can log in.""" | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _cleanup_stale() | |
| to_remove = [ | |
| other_sid | |
| for other_sid, uname in _session_usernames.items() | |
| if uname == username and other_sid != sid | |
| ] | |
| for other_sid in to_remove: | |
| _active_sessions.pop(other_sid, None) | |
| _session_usernames.pop(other_sid, None) | |
| def active_count() -> int: | |
| """How many sessions are currently active.""" | |
| with _lock: | |
| _cleanup_stale() | |
| return len(_active_sessions) | |
| def spots_available() -> int: | |
| """How many open slots remain.""" | |
| with _lock: | |
| _cleanup_stale() | |
| return max(0, _MAX_CONCURRENT - len(_active_sessions)) | |
| def is_session_active() -> bool: | |
| """Check whether the current session still holds a slot. | |
| Returns False if the session was evicted due to inactivity (stale | |
| heartbeat). Callers should clear auth state and redirect to login. | |
| """ | |
| sid = _ensure_session_id() | |
| with _lock: | |
| _cleanup_stale() | |
| return sid in _active_sessions | |
| # --------------------------------------------------------------------------- | |
| # Waiting-room gate (Streamlit UI) | |
| # --------------------------------------------------------------------------- | |
| _SLOWBRO_IMG = ( | |
| "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites" | |
| "/pokemon/other/official-artwork/80.png" | |
| ) | |
| def queue_gate() -> bool: | |
| """Show waiting room if the server is full. Returns True if admitted. | |
| When queue is disabled, always returns True. | |
| """ | |
| if not _queue_enabled(): | |
| return True | |
| # Always send a heartbeat so active sessions stay fresh | |
| heartbeat() | |
| if try_enter(): | |
| return True | |
| # ── Waiting room UI ───────────────────────────────────────── | |
| from waiting_facts import WAITING_FACTS | |
| from utils import language_banner | |
| import random | |
| # Scrolling multilingual banner | |
| language_banner() | |
| _SLOWPOKE_IMG = ( | |
| "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites" | |
| "/pokemon/other/official-artwork/79.png" | |
| ) | |
| n_active = active_count() | |
| spots_left = spots_available() | |
| # ── Slowpoke waiting illustration ── | |
| st.markdown( | |
| "<div style='text-align:center; margin-top:20px;'>" | |
| f"<img src='{_SLOWPOKE_IMG}' width='120' />" | |
| "</div>", | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| "<h3 style='text-align:center;'>Slowbro is busy crunching numbers...</h3>" | |
| "<p style='text-align:center; color:gray;'>" | |
| "All seats are taken! But don't worry — Slowpoke is keeping " | |
| "your spot warm. You'll get in as soon as someone finishes.</p>", | |
| unsafe_allow_html=True, | |
| ) | |
| # ── Queue status ── | |
| q1, q2 = st.columns(2) | |
| with q1: | |
| st.metric("Active users", f"{n_active} / {_MAX_CONCURRENT}") | |
| with q2: | |
| st.metric("Seats available", str(spots_left)) | |
| # ── Session policy note ── | |
| st.warning( | |
| "**How the queue works:** Each user gets a seat for as long as " | |
| "they're active. Sessions expire after **30 minutes** of inactivity " | |
| "to keep things moving — but if you're running a model, your seat " | |
| "is safe until estimation completes." | |
| ) | |
| # ── Rolling cultural facts ── | |
| st.markdown("---") | |
| st.markdown( | |
| "<p style='text-align:center; font-weight:600; margin-bottom:4px;'>" | |
| "While you wait — queuing around the world</p>", | |
| unsafe_allow_html=True, | |
| ) | |
| rng = random.Random(int(time.time()) // 8) # change every 8 seconds | |
| indices = list(range(len(WAITING_FACTS))) | |
| rng.shuffle(indices) | |
| display_facts = [WAITING_FACTS[i] for i in indices[:3]] | |
| for fact in display_facts: | |
| st.info(fact) | |
| # auto-refresh every 5 seconds | |
| time.sleep(5) | |
| st.rerun() | |
| return False # unreachable, but for type-checker | |