Skip to content

Commit 831a4d0

Browse files
authored
fix(server): detect bare uvicorn --workers N for per-worker rate-limit warning (#680) (#689)
* fix(server): detect bare `uvicorn --workers N` for per-worker rate-limit warning (#680) The per-worker in-memory rate-limit WARNING (#678) only fired when worker count came from WEB_CONCURRENCY/UVICORN_WORKERS. A bare `uvicorn ... --workers N` on the CLI (no env var) went undetected, so a genuinely multi-worker, in-memory deployment stayed silent. Add a best-effort, dependency-free signal: walk this process's ancestors via /proc and read a `--workers N` / `-w N` flag off a recognized uvicorn/gunicorn command line. `_detect_worker_count()` now takes the max of the env-based and process-tree results. Detection is gated to ASGI-server cmdlines so an unrelated wrapper's `--workers` flag can't false-positive, and degrades silently to the env result on non-Linux or any error (never raises). Linux-only (via /proc); macOS/Windows fall back to the env-var result. * docs(server): note loose substring bound + pin ascii encoding (#680 review nits)
1 parent b0fb421 commit 831a4d0

2 files changed

Lines changed: 318 additions & 9 deletions

File tree

codeframe/ui/server.py

Lines changed: 142 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -171,16 +171,12 @@ def _validate_security_config():
171171
)
172172

173173

174-
def _detect_worker_count() -> int:
175-
"""Best-effort detection of the configured uvicorn/gunicorn worker count.
176-
177-
Reads ``WEB_CONCURRENCY`` (honored by both uvicorn and gunicorn) and
178-
``UVICORN_WORKERS`` (uvicorn's env-var form of ``--workers``) and returns the
179-
largest valid value found, or ``1`` when neither is set or parseable.
174+
def _worker_count_from_env() -> int:
175+
"""Worker count from ``WEB_CONCURRENCY`` / ``UVICORN_WORKERS`` env vars.
180176
181-
Worker count cannot be observed directly from inside a worker process, so
182-
this is a best-effort signal for the startup warning below. A bare
183-
``uvicorn --workers N`` on the command line (no env var) is not detected.
177+
Returns the largest valid value found, or ``1`` when neither is set or
178+
parseable. ``WEB_CONCURRENCY`` is honored by both uvicorn and gunicorn;
179+
``UVICORN_WORKERS`` is uvicorn's env-var form of ``--workers``.
184180
"""
185181
count = 1
186182
for var in ("WEB_CONCURRENCY", "UVICORN_WORKERS"):
@@ -196,6 +192,143 @@ def _detect_worker_count() -> int:
196192
return count
197193

198194

195+
def _parse_worker_count_from_argv(argv: list[str]) -> int | None:
196+
"""Parse a uvicorn/gunicorn worker count from a command-line argv list.
197+
198+
Recognizes ``--workers N``, ``--workers=N``, ``-w N`` and ``-w=N``. Returns
199+
the positive integer worker count, or ``None`` when no parseable, positive
200+
worker flag is present.
201+
"""
202+
for i, token in enumerate(argv):
203+
value: str | None = None
204+
if token in ("--workers", "-w"):
205+
value = argv[i + 1] if i + 1 < len(argv) else None
206+
elif token.startswith("--workers="):
207+
value = token[len("--workers="):]
208+
elif token.startswith("-w="):
209+
value = token[len("-w="):]
210+
if value is None:
211+
continue
212+
try:
213+
count = int(value.strip())
214+
except (ValueError, AttributeError):
215+
continue
216+
if count > 0:
217+
return count
218+
return None
219+
220+
221+
def _is_asgi_server_cmdline(argv: list[str]) -> bool:
222+
"""True when ``argv`` looks like a uvicorn/gunicorn server invocation.
223+
224+
Matches ``uvicorn``, ``python -m uvicorn``, an absolute ``uvicorn`` path, and
225+
``gunicorn`` (incl. ``-k uvicorn.workers.UvicornWorker``). Used to avoid
226+
reading a ``--workers`` flag off an unrelated ancestor (a wrapper, supervisor,
227+
or test runner) and mistaking it for the server's worker count.
228+
229+
This is a loose substring match on any token, so a path that merely contains
230+
``uvicorn``/``gunicorn`` would also match. That's acceptable for an
231+
advisory-only warning: a false positive only emits an unnecessary WARNING, it
232+
never breaks startup.
233+
"""
234+
return any(("uvicorn" in token or "gunicorn" in token) for token in argv)
235+
236+
237+
def _read_proc_cmdline(pid: int) -> list[str] | None:
238+
"""Read ``/proc/<pid>/cmdline`` as an argv list, or ``None`` if unavailable.
239+
240+
Linux-only; returns ``None`` on any error (missing /proc, permission, race).
241+
"""
242+
try:
243+
with open(f"/proc/{pid}/cmdline", "rb") as f:
244+
raw = f.read()
245+
except OSError:
246+
return None
247+
if not raw:
248+
return None
249+
# /proc cmdline is NUL-separated and usually NUL-terminated.
250+
return [part.decode("utf-8", "replace") for part in raw.split(b"\x00") if part]
251+
252+
253+
def _read_proc_ppid(pid: int) -> int | None:
254+
"""Read the parent PID from ``/proc/<pid>/status``, or ``None`` on error."""
255+
try:
256+
with open(f"/proc/{pid}/status", "r", encoding="ascii") as f:
257+
for line in f:
258+
if line.startswith("PPid:"):
259+
return int(line.split()[1])
260+
except (OSError, ValueError, IndexError):
261+
return None
262+
return None
263+
264+
265+
def _iter_ancestor_cmdlines(max_depth: int = 8) -> list[list[str]]:
266+
"""Best-effort: this process's cmdline plus those of its ancestors.
267+
268+
Walks up the process tree via ``/proc`` (Linux only). Returns ``[]`` on
269+
non-Linux platforms or when nothing is readable. Never raises.
270+
"""
271+
cmdlines: list[list[str]] = []
272+
try:
273+
pid: int | None = os.getpid()
274+
except OSError:
275+
return cmdlines
276+
seen: set[int] = set()
277+
for _ in range(max_depth):
278+
if pid is None or pid <= 0 or pid in seen:
279+
break
280+
seen.add(pid)
281+
cmdline = _read_proc_cmdline(pid)
282+
if cmdline:
283+
cmdlines.append(cmdline)
284+
pid = _read_proc_ppid(pid)
285+
return cmdlines
286+
287+
288+
def _detect_workers_from_process_tree() -> int | None:
289+
"""Best-effort worker count from a ``--workers N`` flag on an ancestor.
290+
291+
Worker subprocesses are spawned, so their own ``sys.argv`` does not carry
292+
``--workers`` — the original flag lives on the supervisor (parent) process.
293+
Returns ``None`` when nothing is found or detection isn't possible. Never
294+
raises.
295+
"""
296+
try:
297+
for argv in _iter_ancestor_cmdlines():
298+
if not _is_asgi_server_cmdline(argv):
299+
continue
300+
count = _parse_worker_count_from_argv(argv)
301+
if count is not None:
302+
return count
303+
except Exception: # best-effort signal; never break startup
304+
return None
305+
return None
306+
307+
308+
def _detect_worker_count() -> int:
309+
"""Best-effort detection of the configured uvicorn/gunicorn worker count.
310+
311+
Combines two signals and returns the largest:
312+
313+
- ``WEB_CONCURRENCY`` / ``UVICORN_WORKERS`` env vars (set by most production
314+
process managers), and
315+
- a ``--workers N`` flag on an ancestor process command line, which covers a
316+
bare ``uvicorn ... --workers N`` invocation that sets no env var (Linux
317+
only, via ``/proc``).
318+
319+
This is a best-effort signal for the startup warning below; detection
320+
failure degrades silently to the env-var result (never raises).
321+
"""
322+
count = _worker_count_from_env()
323+
try:
324+
from_tree = _detect_workers_from_process_tree()
325+
except Exception:
326+
from_tree = None
327+
if from_tree is not None and from_tree > count:
328+
count = from_tree
329+
return count
330+
331+
199332
def _per_worker_rate_limit_warning(
200333
*, enabled: bool, storage: str, worker_count: int
201334
) -> str | None:

tests/ui/test_rate_limit_worker_warning.py

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,179 @@ def test_no_warning_when_rate_limiting_disabled():
9898
)
9999
is None
100100
)
101+
102+
103+
# --- _parse_worker_count_from_argv (issue #680) --------------------------
104+
105+
106+
@pytest.mark.parametrize(
107+
"argv, expected",
108+
[
109+
(["uvicorn", "codeframe.ui.server:app", "--workers", "2"], 2),
110+
(["uvicorn", "app", "--workers=4"], 4),
111+
(["gunicorn", "-w", "3", "app"], 3),
112+
(["gunicorn", "-w=5", "app"], 5),
113+
(["python", "-m", "uvicorn", "app", "--workers", " 6 "], 6),
114+
# Not present / not a worker flag
115+
(["uvicorn", "codeframe.ui.server:app"], None),
116+
(["uvicorn", "app", "--port", "8000"], None),
117+
# Trailing flag with no value
118+
(["uvicorn", "app", "--workers"], None),
119+
# Unparseable / non-positive values are ignored
120+
(["uvicorn", "app", "--workers", "notanumber"], None),
121+
(["uvicorn", "app", "--workers", "0"], None),
122+
(["uvicorn", "app", "--workers", "-2"], None),
123+
([], None),
124+
],
125+
)
126+
def test_parse_worker_count_from_argv(argv, expected):
127+
assert server._parse_worker_count_from_argv(argv) == expected
128+
129+
130+
# --- _detect_workers_from_process_tree (issue #680) ----------------------
131+
132+
133+
def test_process_tree_detection_finds_workers_flag(monkeypatch):
134+
"""A `--workers N` on an ancestor cmdline is detected."""
135+
monkeypatch.setattr(
136+
server,
137+
"_iter_ancestor_cmdlines",
138+
lambda max_depth=8: [
139+
["python", "spawn_main"], # the worker's own bootstrap argv
140+
["uvicorn", "codeframe.ui.server:app", "--workers", "3"], # supervisor
141+
],
142+
)
143+
assert server._detect_workers_from_process_tree() == 3
144+
145+
146+
def test_process_tree_detection_ignores_non_server_workers_flag(monkeypatch):
147+
"""A `--workers` on a non-uvicorn/gunicorn ancestor must not false-positive.
148+
149+
A wrapper/test-runner with its own `--workers` flag is skipped; the real
150+
server ancestor (further up) is the one that counts.
151+
"""
152+
monkeypatch.setattr(
153+
server,
154+
"_iter_ancestor_cmdlines",
155+
lambda max_depth=8: [
156+
["some-wrapper", "--workers", "9"], # NOT a server → ignored
157+
["uvicorn", "codeframe.ui.server:app", "--workers", "2"], # real one
158+
],
159+
)
160+
assert server._detect_workers_from_process_tree() == 2
161+
162+
163+
def test_process_tree_detection_ignores_lone_non_server_workers_flag(monkeypatch):
164+
monkeypatch.setattr(
165+
server,
166+
"_iter_ancestor_cmdlines",
167+
lambda max_depth=8: [["pytest", "-p", "xdist", "--workers", "4"]],
168+
)
169+
assert server._detect_workers_from_process_tree() is None
170+
171+
172+
@pytest.mark.parametrize(
173+
"argv, is_server",
174+
[
175+
(["uvicorn", "app"], True),
176+
(["/usr/bin/uvicorn", "app"], True),
177+
(["python", "-m", "uvicorn", "app"], True),
178+
(["gunicorn", "-k", "uvicorn.workers.UvicornWorker", "app"], True),
179+
(["pytest", "--workers", "4"], False),
180+
(["some-wrapper", "--workers", "9"], False),
181+
([], False),
182+
],
183+
)
184+
def test_is_asgi_server_cmdline(argv, is_server):
185+
assert server._is_asgi_server_cmdline(argv) is is_server
186+
187+
188+
def test_process_tree_detection_returns_none_when_absent(monkeypatch):
189+
monkeypatch.setattr(
190+
server, "_iter_ancestor_cmdlines", lambda max_depth=8: [["uvicorn", "app"]]
191+
)
192+
assert server._detect_workers_from_process_tree() is None
193+
194+
195+
def test_process_tree_detection_never_raises(monkeypatch):
196+
"""Detection degrades gracefully if the walker blows up."""
197+
198+
def boom(max_depth=8):
199+
raise OSError("no /proc here")
200+
201+
monkeypatch.setattr(server, "_iter_ancestor_cmdlines", boom)
202+
assert server._detect_workers_from_process_tree() is None
203+
204+
205+
# --- _detect_worker_count integrates env + process tree (issue #680) ------
206+
207+
208+
def test_detect_worker_count_uses_process_tree(monkeypatch):
209+
"""Bare `uvicorn --workers N` (no env vars) is detected via the tree."""
210+
monkeypatch.setattr(
211+
server, "_detect_workers_from_process_tree", lambda: 2
212+
)
213+
assert server._detect_worker_count() == 2
214+
215+
216+
def test_detect_worker_count_env_wins_when_larger(monkeypatch):
217+
monkeypatch.setenv("WEB_CONCURRENCY", "8")
218+
monkeypatch.setattr(
219+
server, "_detect_workers_from_process_tree", lambda: 2
220+
)
221+
assert server._detect_worker_count() == 8
222+
223+
224+
def test_detect_worker_count_tree_wins_when_larger(monkeypatch):
225+
monkeypatch.setenv("WEB_CONCURRENCY", "2")
226+
monkeypatch.setattr(
227+
server, "_detect_workers_from_process_tree", lambda: 5
228+
)
229+
assert server._detect_worker_count() == 5
230+
231+
232+
def test_detect_worker_count_falls_back_to_env_on_tree_failure(monkeypatch):
233+
"""If the process-tree scan fails, the env-based result still stands."""
234+
monkeypatch.setenv("UVICORN_WORKERS", "4")
235+
236+
def boom():
237+
raise RuntimeError("unexpected")
238+
239+
monkeypatch.setattr(server, "_detect_workers_from_process_tree", boom)
240+
assert server._detect_worker_count() == 4
241+
242+
243+
def test_detect_worker_count_defaults_to_one_with_no_signals(monkeypatch):
244+
monkeypatch.setattr(
245+
server, "_detect_workers_from_process_tree", lambda: None
246+
)
247+
assert server._detect_worker_count() == 1
248+
249+
250+
# --- _iter_ancestor_cmdlines via a fake /proc (issue #680) ---------------
251+
252+
253+
def test_iter_ancestor_cmdlines_walks_proc(monkeypatch):
254+
"""The walker reads the current process and its ancestors from /proc."""
255+
# Fake a tiny process tree: 100 (self) -> 50 (supervisor) -> 1 (init)
256+
cmdlines = {
257+
100: ["python", "spawn_main"],
258+
50: ["uvicorn", "app", "--workers", "2"],
259+
1: ["/sbin/init"],
260+
}
261+
ppids = {100: 50, 50: 1, 1: 0}
262+
263+
monkeypatch.setattr(server.os, "getpid", lambda: 100)
264+
monkeypatch.setattr(server, "_read_proc_cmdline", lambda pid: cmdlines.get(pid))
265+
monkeypatch.setattr(server, "_read_proc_ppid", lambda pid: ppids.get(pid))
266+
267+
result = server._iter_ancestor_cmdlines(max_depth=8)
268+
assert ["uvicorn", "app", "--workers", "2"] in result
269+
# And the parser finds the count through it
270+
assert server._detect_workers_from_process_tree() == 2
271+
272+
273+
def test_iter_ancestor_cmdlines_returns_empty_on_missing_proc(monkeypatch):
274+
monkeypatch.setattr(server, "_read_proc_cmdline", lambda pid: None)
275+
monkeypatch.setattr(server, "_read_proc_ppid", lambda pid: None)
276+
assert server._iter_ancestor_cmdlines(max_depth=8) == []

0 commit comments

Comments
 (0)