mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(rust): leave the fork gate untouched when a late reservation is refused
reserve() stored fork_only_pid before noticing the runtime already ran under that pid, so a refused reservation still reserved the process: the next enter() cleared the runtime claim and children forked afterwards inherited a dead runtime and hung. Undo the reservation on the error path so the gate is exactly as it was. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
18a1491bd2
commit
38fa8a7f55
1 changed files with 20 additions and 0 deletions
|
|
@ -51,9 +51,18 @@ impl ForkGate {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
/// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
|
||||
/// the mirror image, so when the two race at least one of them sees the other. A refused
|
||||
/// reservation leaves the gate exactly as it was, so a process already running the runtime
|
||||
/// keeps refusing the children it forks.
|
||||
pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
|
||||
self.fork_only_pid.store(pid, Ordering::SeqCst);
|
||||
if self.runtime_pid.load(Ordering::SeqCst) == pid {
|
||||
// Nothing may change for a process that already runs the runtime: its children
|
||||
// must still be refused.
|
||||
let _ =
|
||||
self.fork_only_pid
|
||||
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
|
||||
return Err(RuntimeAlreadyStarted);
|
||||
}
|
||||
Ok(())
|
||||
|
|
@ -109,6 +118,17 @@ mod tests {
|
|||
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
|
||||
let gate = ForkGate::new();
|
||||
gate.enter(MASTER).unwrap();
|
||||
|
||||
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
|
||||
assert_eq!(gate.enter(MASTER), Ok(()));
|
||||
assert!(gate.started(MASTER));
|
||||
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
|
||||
let gate = ForkGate::new();
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue