wangxuming/makelore
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: settle Pi child capacity reservations

Browse Source
This commit is contained in:
brother7
2026-08-23 13:34:41 +08:00
parent 3de85d61d7
commit a10b98e484
5 changed files with 275 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
electron/coding-runtime/pi/subagent.ts
View File

@@ -58,7 +58,7 @@ export interface PiSubagentDispatchOptions {
export interface PiSubagentSchedulerOptions {
openChild(input: PiSubagentChildOpenInput): Promise<PiSubagentChild>;
processBudget: PiProcessBudget;
reclaimProcessCapacity?(): Promise<boolean>;
reclaimProcessCapacity?(signal?: AbortSignal): Promise<boolean>;
createId?: (kind: 'dispatch' | 'task') => string;
}
@@ -216,7 +216,7 @@ export class PiSubagentScheduler {
private readonly openChild: PiSubagentSchedulerOptions['openChild'];
private readonly processBudget: PiProcessBudget;
private readonly childPermits: FifoSemaphore;
private readonly reclaimProcessCapacity: (() => Promise<boolean>) | undefined;
private readonly reclaimProcessCapacity: ((signal?: AbortSignal) => Promise<boolean>) | undefined;
private readonly createId: NonNullable<PiSubagentSchedulerOptions['createId']>;
private readonly dispatches = new Map<string, DispatchRecord>();
private readonly parentDispatches = new Map<string, Set<string>>();
@@ -353,12 +353,7 @@ export class PiSubagentScheduler {
let child: PiSubagentChild | undefined;
try {
releaseChild = await this.childPermits.acquire(record.controller.signal);
const budgetWasFull = this.processBudget.activeCount >= this.processBudget.maxProcesses;
const processLeaseFlight = this.processBudget.acquire(record.controller.signal, 'child');
if (budgetWasFull && this.reclaimProcessCapacity) {
await this.reclaimProcessCapacity();
}
processLease = await processLeaseFlight;
processLease = await this.acquireProcessLease(record.controller.signal);
if (record.controller.signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED');
projected.status = 'running';
this.emit(details, onUpdate);
@@ -395,6 +390,33 @@ export class PiSubagentScheduler {
}
}
private async acquireProcessLease(signal: AbortSignal): Promise<PiProcessLease> {
const reservation = new AbortController();
const abortReservation = () => reservation.abort();
signal.addEventListener('abort', abortReservation, { once: true });
if (signal.aborted) reservation.abort();
const budgetWasFull = this.processBudget.activeCount >= this.processBudget.maxProcesses;
const processLeaseFlight = this.processBudget.acquire(reservation.signal, 'child');
try {
if (budgetWasFull && this.reclaimProcessCapacity) {
const outcome = await Promise.race([
processLeaseFlight.then((lease) => ({ lease })),
this.reclaimProcessCapacity(reservation.signal).then(() => null),
]);
if (outcome) return outcome.lease;
}
return await processLeaseFlight;
} catch (error) {
reservation.abort();
const orphanedLease = await processLeaseFlight.catch(() => undefined);
orphanedLease?.release();
throw error;
} finally {
signal.removeEventListener('abort', abortReservation);
reservation.abort();
}
}
private markRemaining(
details: SubagentDetailsV1,
from: number,

68
electron/coding-runtime/pi/worker-pool.ts
View File

@@ -230,6 +230,7 @@ export class PiWorkerPool {
private readonly activeRuns = new Map<string, { runId: string; generation: number }>();
private readonly generations = new Map<string, number>();
private readonly listeners = new Set<(event: PiWorkerPoolEvent) => void>();
private readonly reclaimWaiters = new Set<() => void>();
private commandSequence = 0;
private runningCount = 0;
private useSequence = 0;
@@ -274,6 +275,7 @@ export class PiWorkerPool {
.then(async ({ opened: { worker, session }, lease }) => {
const record = this.createRecord(conversation, worker, session, generation, revision, lease);
this.workers.set(conversation.conversationId, record);
this.notifyReclaimableWorker();
const state = this.publicState(record);
await this.trimIdleWorkers();
return state;
@@ -328,6 +330,7 @@ export class PiWorkerPool {
lease,
);
this.workers.set(conversation.conversationId, record);
this.notifyReclaimableWorker();
await this.trimIdleWorkers();
return this.publicState(record);
})().finally(() => {
@@ -344,13 +347,19 @@ export class PiWorkerPool {
return record ? this.publicState(record) : null;
}
async reclaimIdleWorker(): Promise<boolean> {
const record = [...this.workers.values()]
.filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
if (!record) return false;
await this.evict(record);
return true;
async reclaimIdleWorker(signal?: AbortSignal): Promise<boolean> {
while (true) {
if (signal?.aborted) throw new Error('Pi idle worker reclaim cancelled');
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
const record = [...this.workers.values()]
.filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
if (record) {
if (await this.evict(record)) return true;
continue;
}
await this.waitForReclaimableWorker(signal);
}
}
subscribe(listener: (event: PiWorkerPoolEvent) => void): () => void {
@@ -406,6 +415,7 @@ export class PiWorkerPool {
const record = this.workers.get(conversationId);
if (record && record.state !== 'crashed') {
record.state = 'idle';
this.notifyReclaimableWorker();
try {
this.revisions.settleRun(record.revisionWorkerId);
} catch {
@@ -605,6 +615,7 @@ export class PiWorkerPool {
if (current && this.workers.get(run.conversationId) === current
&& current.state !== 'crashed') {
current.state = 'idle';
this.notifyReclaimableWorker();
if (beganRun) this.revisions.settleRun(current.revisionWorkerId);
void this.trimIdleWorkers();
}
@@ -619,6 +630,7 @@ export class PiWorkerPool {
this.activeRuns.delete(conversationId);
this.runningCount -= 1;
record.state = 'idle';
this.notifyReclaimableWorker();
const revisionAction = this.revisions.settleRun(record.revisionWorkerId);
if (!this.shuttingDown
&& (record.reconfigureAfterSettled || revisionAction.action === 'rebuild-after-settled')) {
@@ -716,6 +728,7 @@ export class PiWorkerPool {
? 'running'
: 'ready';
this.workers.set(conversationId, replacement);
if (replacement.state === 'ready') this.notifyReclaimableWorker();
this.emit({
type: 'worker.replaced',
conversationId,
@@ -796,15 +809,16 @@ export class PiWorkerPool {
await Promise.all(idle.slice(0, excess).map((record) => this.evict(record)));
}
private async evict(record: WorkerRecord): Promise<void> {
private async evict(record: WorkerRecord): Promise<boolean> {
const conversationId = record.conversation.conversationId;
if (this.workers.get(conversationId) !== record) return;
if (this.workers.get(conversationId) !== record) return false;
record.unsubscribeEvent();
record.unsubscribeInvalidation();
this.cancelGenerationResources(record);
this.revisions.removeWorker(record.revisionWorkerId);
this.workers.delete(conversationId);
await this.stopAndRelease(record);
return true;
}
private launchWaitingRuns(): void {
@@ -902,9 +916,39 @@ export class PiWorkerPool {
}
private async stopAndRelease(record: WorkerRecord): Promise<void> {
await record.worker.stop();
record.processLease?.release();
record.processLease = null;
try {
await record.worker.stop();
} finally {
record.processLease?.release();
record.processLease = null;
}
}
private waitForReclaimableWorker(signal?: AbortSignal): Promise<void> {
return new Promise((resolve, reject) => {
const cleanup = () => {
this.reclaimWaiters.delete(wake);
signal?.removeEventListener('abort', cancel);
this.shutdownController.signal.removeEventListener('abort', cancel);
};
const wake = () => {
cleanup();
resolve();
};
const cancel = () => {
cleanup();
reject(new Error(
this.shuttingDown ? 'Pi worker pool is shutting down' : 'Pi idle worker reclaim cancelled',
));
};
this.reclaimWaiters.add(wake);
signal?.addEventListener('abort', cancel, { once: true });
this.shutdownController.signal.addEventListener('abort', cancel, { once: true });
});
}
private notifyReclaimableWorker(): void {
for (const wake of [...this.reclaimWaiters]) wake();
}
private emit(event: PiWorkerPoolEvent): void {