Compare commits

...
Author SHA1 Message Date
Daniel Salazar f3542a6aa3 fix(backend): shutdown hooks don't skip each other, and final flushes wait for in-flight work (#4499)
A throwing prepare/shutdown hook stopped the loop, so shutdown() skipped
closeAllConnections and every later hook, including the final metering
flush; each hook now runs and failures are logged. Broadcast, events
forwarding and share digests joined nothing already in flight and could
drop what was queued behind it; their final flushes now wait for it, ship
the tail, and drain once more at shutdown. Local worker disposal and the
OpenAI image metering call are awaited. settleWithin and
FINAL_FLUSH_TIMEOUT_MS bound each final flush.
2026-11-04 19:41:17 -07:00
Daniel Salazar 5a59db7b6b feat(backend): drain-start hooks run concurrently and can hold the listener open
The listener closes once every onServerDrainStart hook settles, and the
drain window now starts at the signal, so time a hook spends holding the
listener open counts against drainMs instead of extending shutdown. A
failing hook no longer skips the others.
2026-11-04 18:38:33 -07:00
17 changed files with 1100 additions and 92 deletions

No files matched your search

@@ -235,6 +235,56 @@ describe('OpenAiImageProvider.generate argument validation', () => {
// ── Credit gate ─────────────────────────────────────────────────────
describe('OpenAiImageProvider.generate credit gate', () => {
it('waits for the real usage write before returning an image', async () => {
const actor = makeActor({
user: {
id: 62,
uuid: `image-write-${Date.now()}`,
username: 'metered',
},
});
const redis = server.clients.redis as unknown as {
meterIncr: (...args: string[]) => Promise<unknown>;
};
const realIncr = redis.meterIncr.bind(redis);
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
let blocked = false;
vi.spyOn(redis, 'meterIncr').mockImplementation(async (...args) => {
if (!blocked && args[5].includes(`:actor:${actor.user!.uuid}:`)) {
blocked = true;
await held;
}
return realIncr(...args);
});
generateMock.mockResolvedValueOnce({
data: [{ url: 'https://image.example/result.png' }],
});
const provider = makeProvider();
let finished = false;
const generating = withTestActor(
() =>
provider.generate({
model: 'gpt-image-2',
prompt: 'an image',
ratio: { w: 1024, h: 1024 },
}),
actor,
).then((result) => {
finished = true;
return result;
});
try {
await vi.waitFor(() => expect(blocked).toBe(true));
expect(finished).toBe(false);
} finally {
release();
await generating;
}
});
it('throws 402 BEFORE hitting OpenAI when actor lacks credits', async () => {
const provider = makeProvider();
hasCreditsSpy.mockResolvedValueOnce(false);
@@ -293,7 +293,10 @@ export class OpenAiImageProvider implements IImageProvider {
});
}
if (usageEntries.length) {
this.#meteringService.batchIncrementUsages(actor, usageEntries);
await this.#meteringService.batchIncrementUsages(
actor,
usageEntries,
);
}
const url =
+165
View File
@@ -1492,6 +1492,104 @@ describe('PuterServer keep-alive timeout', () => {
});
describe('PuterServer lifecycle hooks', () => {
it('fails startup before running later hooks when a start hook rejects', async () => {
const failure = new Error('startup failed');
const later = vi.fn();
class FailingStart extends JobScheduler {
override async onServerStart() {
throw failure;
}
}
class LaterStart extends JobScheduler {
override async onServerStart() {
later();
}
}
extension.registerService('failingStart', FailingStart as never);
extension.registerService('laterStart', LaterStart as never);
try {
await expect(setupTestServer()).rejects.toBe(failure);
expect(later).not.toHaveBeenCalled();
} finally {
delete (extensionStore.services as Record<string, unknown>)
.failingStart;
delete (extensionStore.services as Record<string, unknown>)
.laterStart;
}
});
it.each([false, true])(
'continues teardown after failing hooks with listener callback %s',
async (prepareFirst) => {
const port = await allocateEphemeralPort();
const server = await setupTestServer(
{ port } as unknown as IConfig,
{ listen: true },
);
const logged = vi
.spyOn(console, 'error')
.mockImplementation(() => {});
const ran: string[] = [];
class FailingHooks {
onServerPrepareShutdown() {
ran.push('prepare:failed');
throw new Error('prepare failed');
}
async onServerShutdown() {
ran.push('shutdown:failed');
throw new Error('shutdown failed');
}
}
(server.services as Record<string, unknown>).fails =
new FailingHooks();
(server.services as Record<string, unknown>).after = {
onServerPrepareShutdown: () => ran.push('prepare:after'),
onServerShutdown: () => ran.push('shutdown:after'),
};
(server.clients as Record<string, unknown>).last = {
onServerShutdown: () => ran.push('shutdown:client'),
};
try {
if (prepareFirst) {
await server.prepareShutdown();
await vi.waitFor(() =>
expect(ran).toContain('prepare:after'),
);
}
await server.shutdown();
expect(ran).toEqual([
'prepare:failed',
'prepare:after',
'shutdown:failed',
'shutdown:after',
'shutdown:client',
]);
expect(logged).toHaveBeenCalledWith(
'[server] onServerPrepareShutdown failed in FailingHooks',
expect.any(Error),
);
expect(logged).toHaveBeenCalledWith(
'[server] onServerShutdown failed in FailingHooks',
expect.any(Error),
);
await expect(
new Promise<void>((resolve, reject) => {
const socket = net.connect(port, '127.0.0.1');
socket.on('connect', () => {
socket.destroy();
resolve();
});
socket.on('error', reject);
}),
).rejects.toMatchObject({ code: 'ECONNREFUSED' });
} finally {
delete (server.services as Record<string, unknown>).fails;
await server.shutdown();
logged.mockRestore();
}
},
);
it('runs drain-start and prepare-shutdown bottom-up, once, before tearing down', async () => {
const port = await allocateEphemeralPort();
const server = await setupTestServer({ port } as unknown as IConfig, {
@@ -1541,6 +1639,73 @@ describe('PuterServer lifecycle hooks', () => {
'shutdown:client',
]);
});
it('keeps accepting connections until every drain-start hook settles, without one holding up another', async () => {
const port = await allocateEphemeralPort();
const server = await setupTestServer({ port } as unknown as IConfig, {
listen: true,
});
let release!: () => void;
const held = new Promise<void>((resolve) => (release = resolve));
let otherRan = false;
(server.services as unknown as Record<string, unknown>).holdsListener =
{ onServerDrainStart: () => held };
(server.drivers as unknown as Record<string, unknown>).other = {
onServerDrainStart: () => {
otherRan = true;
},
};
const connect = () =>
new Promise<void>((resolve, reject) => {
const socket = net.connect(port, '127.0.0.1');
socket.on('connect', () => {
socket.destroy();
resolve();
});
socket.on('error', reject);
});
const prepared = server.prepareShutdown({ drainMs: 1_000 });
await new Promise((resolve) => setTimeout(resolve, 20));
expect(otherRan).toBe(true);
await expect(connect()).resolves.toBeUndefined();
release();
await prepared;
await expect(connect()).rejects.toMatchObject({
code: 'ECONNREFUSED',
});
await server.shutdown();
});
it('runs every drain-start hook when one of them throws', async () => {
const port = await allocateEphemeralPort();
const server = await setupTestServer({ port } as unknown as IConfig, {
listen: true,
});
vi.spyOn(console, 'error').mockImplementation(() => {});
const ran: string[] = [];
(server.stores as unknown as Record<string, unknown>).fails = {
onServerDrainStart: () => {
throw new Error('drain hook failed');
},
};
(server.services as unknown as Record<string, unknown>).after = {
onServerDrainStart: () => {
ran.push('after');
},
};
await server.prepareShutdown({ drainMs: 1_000 });
await server.shutdown();
expect(ran).toEqual(['after']);
expect(console.error).toHaveBeenCalledWith(
expect.stringContaining('onServerDrainStart failed'),
expect.any(Error),
);
});
});
describe('PuterServer shutdown during prepare', () => {
+47 -13
View File
@@ -1743,12 +1743,7 @@ export class PuterServer {
*/
async #runHook(
order: 'bottomUp' | 'topDown',
hook:
| 'onServerStart'
| 'onServerDrainStart'
| 'onServerPrepareShutdown'
| 'onServerShutdown',
info?: ServerDrainInfo,
hook: 'onServerStart' | 'onServerPrepareShutdown' | 'onServerShutdown',
) {
const layers = [
this.clients,
@@ -1760,11 +1755,47 @@ export class PuterServer {
if (order === 'topDown') layers.reverse();
for (const layer of layers) {
for (const instance of Object.values(layer) as WithLifecycle[]) {
await instance[hook]?.(info);
try {
await instance[hook]?.();
} catch (err) {
if (hook === 'onServerStart') throw err;
console.error(
`[server] ${hook} failed in ${instance.constructor.name}`,
err,
);
}
}
}
}
/**
* Run every `onServerDrainStart` concurrently, so one hook holding the
* listener open doesn't delay the others, and one failing doesn't skip the
* rest.
*/
async #runDrainStartHooks(info: ServerDrainInfo) {
const instances = [
this.clients,
this.stores,
this.services,
this.controllers,
this.drivers,
].flatMap((layer) => Object.values(layer) as WithLifecycle[]);
const results = await Promise.allSettled(
instances.map(async (instance) =>
instance.onServerDrainStart?.(info),
),
);
results.forEach((result, i) => {
if (result.status === 'rejected') {
console.error(
`[server] onServerDrainStart failed in ${instances[i].constructor.name}`,
result.reason,
);
}
});
}
#prepareShutdownHooks: Promise<void> | null = null;
/**
@@ -1782,16 +1813,19 @@ export class PuterServer {
async prepareShutdown(info: ServerDrainInfo = {}) {
if (this.#server) {
try {
await this.#runHook('bottomUp', 'onServerDrainStart', info);
} catch (err) {
console.error('[server] onServerDrainStart failed', err);
}
await this.#runDrainStartHooks(info);
this.#server.close(async () => {
console.log(
'PuterServer has stopped accepting new connections',
);
await this.#runPrepareShutdownHooks();
try {
await this.#runPrepareShutdownHooks();
} catch (err) {
console.error(
'[server] prepare-shutdown hooks failed',
err,
);
}
});
}
}
@@ -1149,3 +1149,154 @@ describe('BroadcastService — same-cluster pub/sub fan-out', () => {
}
});
});
describe('BroadcastService shutdown drain', () => {
let server: PuterServer;
let broadcast: BroadcastService;
beforeEach(async () => {
axiosRequestMock.mockReset();
axiosRequestMock.mockResolvedValue({ status: 200, data: 'ok' });
server = await setupTestServer({
broadcast: {
webhook: { peerId: SELF_PEER_ID, secret: SELF_SECRET },
peers: [
{
peerId: PEER_ID,
webhook: true,
webhook_url: PEER_URL,
webhook_secret: PEER_SECRET,
},
],
outbound_flush_ms: 10,
},
} as never);
broadcast = server.services.broadcast as unknown as BroadcastService;
});
afterEach(async () => {
vi.useRealTimers();
await broadcast.onServerPrepareShutdown();
await broadcast.onServerShutdown();
await server.shutdown();
});
const emit = (hash: string) =>
server.clients.event.emit(
'outer.fs.write-hash' as never,
{ hash, uuid: hash } as never,
{},
);
it.each(['onServerPrepareShutdown', 'onServerShutdown'] as const)(
'waits for the active flush and sends its tail in %s',
async (hook) => {
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
let releaseTail!: () => void;
const tailHeld = new Promise<void>((resolve) => {
releaseTail = resolve;
});
axiosRequestMock
.mockImplementationOnce(async () => {
await held;
return { status: 200, data: 'ok' };
})
.mockImplementationOnce(async () => {
await tailHeld;
return { status: 200, data: 'ok' };
});
emit('first');
await vi.waitFor(() =>
expect(axiosRequestMock).toHaveBeenCalledTimes(1),
);
emit('tail');
let finished = false;
const draining = broadcast[hook]().then(() => {
finished = true;
});
try {
await new Promise((resolve) => setTimeout(resolve, 20));
expect(finished).toBe(false);
release();
await vi.waitFor(() =>
expect(axiosRequestMock).toHaveBeenCalledTimes(2),
);
expect(finished).toBe(false);
} finally {
release();
releaseTail();
await draining;
}
expect(axiosRequestMock).toHaveBeenCalledTimes(2);
expect(axiosRequestMock.mock.calls[1][0].data).toContain('tail');
},
);
it('keeps normal timer flushing after prepare on an idle server', async () => {
await broadcast.onServerPrepareShutdown();
vi.useFakeTimers();
axiosRequestMock.mockClear();
emit('during-drain');
await vi.advanceTimersByTimeAsync(10);
expect(axiosRequestMock).toHaveBeenCalledTimes(1);
expect(axiosRequestMock.mock.calls[0][0].data).toContain(
'during-drain',
);
});
it('sends events queued after prepare during shutdown and clears the timer', async () => {
await broadcast.onServerPrepareShutdown();
vi.useFakeTimers();
axiosRequestMock.mockClear();
emit('after-prepare');
await broadcast.onServerShutdown();
expect(axiosRequestMock).toHaveBeenCalledTimes(1);
expect(axiosRequestMock.mock.calls[0][0].data).toContain(
'after-prepare',
);
axiosRequestMock.mockClear();
await vi.advanceTimersByTimeAsync(100);
expect(axiosRequestMock).not.toHaveBeenCalled();
});
it.each(['onServerPrepareShutdown', 'onServerShutdown'] as const)(
'bounds a hung flush in %s',
async (hook) => {
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
axiosRequestMock.mockImplementationOnce(async () => {
await held;
return { status: 200, data: 'ok' };
});
emit('hung');
await vi.waitFor(() =>
expect(axiosRequestMock).toHaveBeenCalledTimes(1),
);
vi.useFakeTimers();
emit('tail');
let finished = false;
const draining = broadcast[hook]().then(() => {
finished = true;
});
try {
await vi.advanceTimersByTimeAsync(4_999);
expect(finished).toBe(false);
await vi.advanceTimersByTimeAsync(1);
await draining;
release();
await vi.advanceTimersByTimeAsync(100);
expect(axiosRequestMock).toHaveBeenCalledTimes(
hook === 'onServerPrepareShutdown' ? 2 : 1,
);
} finally {
release();
vi.useRealTimers();
}
},
);
});
@@ -21,6 +21,8 @@ import axios from 'axios';
import { createHmac, randomUUID, timingSafeEqual } from 'node:crypto';
import { Agent as HttpsAgent } from 'node:https';
import { IBroadcastPeerConfig } from '../../types.js';
import { settleWithin } from '../../util/concurrency.js';
import { FINAL_FLUSH_TIMEOUT_MS } from '../../util/gracefulShutdown.js';
import { PuterService } from '../types.js';
// -- Wire types ------------------------------------------------------
@@ -121,7 +123,8 @@ export class BroadcastService extends PuterService {
/** Per peer, serialized events it has not accepted yet, oldest first. */
#outboundBacklogs = new Map<IBroadcastPeerConfig, string[]>();
#outboundFlushTimer: ReturnType<typeof setTimeout> | null = null;
#outboundIsFlushing = false;
#outboundFlush: Promise<void> | null = null;
#outboundDraining = false;
#webhookReplayWindowSeconds = 300;
#outboundFlushMs = 2000;
@@ -165,21 +168,40 @@ export class BroadcastService extends PuterService {
}
override async onServerPrepareShutdown(): Promise<void> {
await this.#drainOutboundEvents();
if (this.#redisSub) {
try {
await this.#redisSub.unsubscribe('pubsub');
this.#redisSub.quit();
} catch (err) {
console.warn('[broadcast] unsubscribe failed', err);
} finally {
this.#redisSub = null;
}
}
}
override async onServerShutdown(): Promise<void> {
this.#outboundDraining = true;
await this.#drainOutboundEvents();
}
async #drainOutboundEvents(): Promise<void> {
if (this.#outboundFlushTimer) {
clearTimeout(this.#outboundFlushTimer);
this.#outboundFlushTimer = null;
}
// Best-effort drain — try one final flush so events queued near
// shutdown make it out.
try {
await this.#flushOutboundEvents();
} catch (err) {
console.warn('[broadcast] final flush failed', err);
}
if (this.#redisSub) {
await this.#redisSub.unsubscribe('pubsub');
this.#redisSub.quit();
this.#redisSub = null;
const deadline = Date.now() + FINAL_FLUSH_TIMEOUT_MS;
const drain = async () => {
try {
await this.#outboundFlush;
if (Date.now() < deadline) await this.#flushOutboundEvents();
} catch (err) {
console.warn('[broadcast] final flush failed', err);
}
};
if (!(await settleWithin(drain(), FINAL_FLUSH_TIMEOUT_MS))) {
console.warn('[broadcast] final flush timed out');
}
}
@@ -453,24 +475,27 @@ export class BroadcastService extends PuterService {
}
#scheduleOutboundFlush(): void {
if (this.#outboundFlushTimer) return;
this.#outboundFlushTimer = setTimeout(() => {
if (this.#outboundDraining || this.#outboundFlushTimer) return;
this.#outboundFlushTimer = setTimeout(async () => {
this.#outboundFlushTimer = null;
void this.#flushOutboundEvents().catch((err) => {
try {
await this.#flushOutboundEvents();
} catch (err) {
console.warn('[broadcast] outbound flush failed', err);
});
}
}, this.#outboundFlushMs);
}
async #flushOutboundEvents(): Promise<void> {
if (
this.#outboundIsFlushing ||
(this.#outboundQueue.size === 0 &&
this.#outboundBacklogs.size === 0)
)
return;
#flushOutboundEvents(): Promise<void> {
if (this.#outboundFlush) return this.#outboundFlush;
if (this.#outboundQueue.size === 0 && this.#outboundBacklogs.size === 0)
return Promise.resolve();
this.#outboundIsFlushing = true;
this.#outboundFlush = this.#sendOutboundEvents();
return this.#outboundFlush;
}
async #sendOutboundEvents(): Promise<void> {
try {
const fresh = this.#outboundQueue;
this.#outboundQueue = new Set();
@@ -504,7 +529,7 @@ export class BroadcastService extends PuterService {
for (const result of flushed)
if (result.status === 'rejected') throw result.reason;
} finally {
this.#outboundIsFlushing = false;
this.#outboundFlush = null;
// Anything that arrived during flush, or is still owed to a peer,
// gets the next tick.
if (
@@ -217,6 +217,97 @@ describe('the addressed peer queue', () => {
).toEqual(['b']);
});
it.each([false, true])(
'waits for an active send and ships its tail after failure %s',
async (failFirst) => {
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
let releaseTail!: () => void;
const tailHeld = new Promise<void>((resolve) => {
releaseTail = resolve;
});
const sent: ForwardItem[][] = [];
const { queue } = makeQueue({
maxItems: 1,
flushMs: 60_000,
send: async (_peerId, items) => {
sent.push(items);
if (sent.length === 1) {
await held;
if (failFirst) throw new Error('peer unavailable');
} else {
await tailHeld;
}
},
});
queue.push('east', delivery('first'));
await vi.waitFor(() => expect(sent).toHaveLength(1));
queue.push('east', delivery('tail'));
let finished = false;
const draining = queue.flushAll().then(() => {
finished = true;
});
try {
await new Promise((resolve) => setTimeout(resolve, 20));
expect(finished).toBe(false);
release();
await vi.waitFor(() => expect(sent).toHaveLength(2));
expect(finished).toBe(false);
} finally {
release();
releaseTail();
await draining;
queue.stop();
}
expect(sent).toEqual([[delivery('first')], [delivery('tail')]]);
},
);
it('releases the peer after a send throws synchronously', async () => {
const send = vi.fn(() => {
throw new Error('send failed');
});
const { queue } = makeQueue({ maxItems: 1, send });
queue.push('east', delivery('first'));
await queue.flushAll();
queue.push('east', delivery('second'));
await queue.flushAll();
expect(send).toHaveBeenCalledTimes(2);
});
it('bounds a final drain while a send is hung', async () => {
vi.useFakeTimers();
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
const send = vi.fn(async () => {
await held;
});
const { queue } = makeQueue({ maxItems: 1, send });
try {
queue.push('east', delivery('first'));
queue.push('east', delivery('tail'));
let finished = false;
const draining = queue.flushAll().then(() => {
finished = true;
});
await vi.advanceTimersByTimeAsync(4_999);
expect(finished).toBe(false);
await vi.advanceTimersByTimeAsync(1);
await draining;
queue.stop();
release();
await vi.advanceTimersByTimeAsync(100);
expect(send).toHaveBeenCalledTimes(1);
expect(vi.getTimerCount()).toBe(0);
} finally {
release();
vi.useRealTimers();
}
});
it('sends nothing more once it is stopped', async () => {
const { queue, sent } = makeQueue();
queue.stop();
+28 -9
View File
@@ -18,6 +18,8 @@
*/
import type { EventKey, KvOp } from '../../clients/event/types.js';
import { settleWithin } from '../../util/concurrency.js';
import { FINAL_FLUSH_TIMEOUT_MS } from '../../util/gracefulShutdown.js';
import type { DeliverableEvent } from './registry.js';
/**
@@ -204,7 +206,7 @@ interface PeerQueue {
items: ForwardItem[];
bytes: number;
timer: ReturnType<typeof setTimeout> | null;
flushing: boolean;
flushing: Promise<void> | null;
}
export class PeerForwardQueue {
@@ -272,19 +274,27 @@ export class PeerForwardQueue {
}
/** Ship whatever is held for one peer. */
async flush(peerId: string): Promise<void> {
flush(peerId: string): Promise<void> {
const queue = this.#peers.get(peerId);
if (!queue || queue.flushing || queue.items.length === 0) return;
if (!queue) return Promise.resolve();
if (queue.flushing) return queue.flushing;
if (queue.items.length === 0) return Promise.resolve();
if (queue.timer) {
clearTimeout(queue.timer);
queue.timer = null;
}
queue.flushing = true;
queue.flushing = this.#send(peerId, queue);
return queue.flushing;
}
async #send(peerId: string, queue: PeerQueue): Promise<void> {
const batch = queue.items;
queue.items = [];
queue.bytes = 0;
try {
// Register the in-flight promise before a send can throw synchronously.
await Promise.resolve();
await this.#options.send(peerId, batch);
} catch (err) {
// A lost batch is a lost at-most-once delivery, and a `single`
@@ -292,7 +302,7 @@ export class PeerForwardQueue {
// hand the client an event whose moment has passed.
console.warn(`[events] forward to ${peerId} failed`, err);
} finally {
queue.flushing = false;
queue.flushing = null;
// Anything that arrived mid-flush gets the next window.
if (queue.items.length > 0 && !queue.timer && !this.#stopped) {
queue.timer = setTimeout(() => {
@@ -305,7 +315,16 @@ export class PeerForwardQueue {
}
async flushAll(): Promise<void> {
await Promise.all([...this.#peers.keys()].map((id) => this.flush(id)));
const deadline = Date.now() + FINAL_FLUSH_TIMEOUT_MS;
const work = Promise.allSettled(
[...this.#peers.entries()].map(async ([id, queue]) => {
await queue.flushing;
if (Date.now() < deadline) await this.flush(id);
}),
);
if (!(await settleWithin(work, FINAL_FLUSH_TIMEOUT_MS))) {
console.warn('[events] final flush timed out');
}
}
stop(): void {
@@ -323,7 +342,7 @@ export class PeerForwardQueue {
items: [],
bytes: 0,
timer: null,
flushing: false,
flushing: null,
};
this.#peers.set(peerId, queue);
return queue;
@@ -349,7 +368,7 @@ const isGapMarker = (item: ForwardItem): boolean =>
*/
const shed = (queue: PeerQueue, count: number): ForwardItem[] => {
const dropped: ForwardItem[] = [];
for (let i = 0; i < queue.items.length && dropped.length < count;) {
for (let i = 0; i < queue.items.length && dropped.length < count; ) {
if (isGapMarker(queue.items[i])) {
i++;
continue;
@@ -372,7 +391,7 @@ const shed = (queue: PeerQueue, count: number): ForwardItem[] => {
const shedBytes = (queue: PeerQueue, maxBytesHeld: number): ForwardItem[] => {
const dropped: ForwardItem[] = [];
let remaining = queue.bytes;
for (let i = 0; i < queue.items.length && remaining > maxBytesHeld;) {
for (let i = 0; i < queue.items.length && remaining > maxBytesHeld; ) {
if (isGapMarker(queue.items[i])) {
i++;
continue;
@@ -0,0 +1,99 @@
import {
afterAll,
afterEach,
beforeAll,
describe,
expect,
it,
vi,
} from 'vitest';
import type { PuterServer } from '../../server.js';
import { setupTestServer } from '../../testUtil.js';
const { disposals } = vi.hoisted(() => ({
disposals: [] as Array<ReturnType<typeof vi.fn>>,
}));
vi.mock('miniflare', () => ({
Miniflare: class {
dispose = vi.fn(async () => {});
constructor() {
disposals.push(this.dispose);
}
},
}));
describe('LocalWorkerService shutdown disposal', () => {
let server: PuterServer;
beforeAll(async () => {
server = await setupTestServer({
api_base_url: 'http://api.puter.localhost',
} as never);
});
afterAll(async () => {
await server.shutdown();
});
afterEach(() => {
vi.useRealTimers();
disposals.length = 0;
});
it('awaits every disposal even if one fails', async () => {
const service = server.services.localworkerservice;
await service.cfDeployLocal('first', 'auth', 'export default {}');
await service.cfDeployLocal('second', 'auth', 'export default {}');
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
disposals[0].mockRejectedValueOnce(new Error('dispose failed'));
disposals[1].mockImplementationOnce(async () => {
await held;
});
let finished = false;
const draining = Promise.resolve(service.onServerShutdown()).then(
() => {
finished = true;
},
);
try {
await new Promise((resolve) => setTimeout(resolve, 20));
expect(
disposals.every((dispose) => dispose.mock.calls.length === 1),
).toBe(true);
expect(finished).toBe(false);
} finally {
release();
await draining;
}
});
it('bounds disposal of a hung worker to five seconds', async () => {
const service = server.services.localworkerservice;
await service.cfDeployLocal('hung', 'auth', 'export default {}');
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
disposals[0].mockImplementationOnce(async () => {
await held;
});
vi.useFakeTimers();
try {
let finished = false;
const draining = Promise.resolve(service.onServerShutdown()).then(
() => {
finished = true;
},
);
await vi.advanceTimersByTimeAsync(4_999);
expect(finished).toBe(false);
await vi.advanceTimersByTimeAsync(1);
await draining;
expect(vi.getTimerCount()).toBe(0);
} finally {
release();
vi.useRealTimers();
}
});
});
@@ -10,6 +10,8 @@ import {
import { puterStores } from '../../stores';
import type { SubdomainRow } from '../../stores/subdomain/SubdomainStore';
import { LayerInstances } from '../../types';
import { settleWithin } from '../../util/concurrency.js';
import { FINAL_FLUSH_TIMEOUT_MS } from '../../util/gracefulShutdown.js';
import { PuterService } from '../types';
const MAX_SOURCE_SIZE = 10 * 1024 * 1024; // 10 MB
@@ -240,13 +242,16 @@ export class LocalWorkerService extends PuterService {
idleSweepTimer.unref?.();
}
override onServerShutdown(): void {
override async onServerShutdown(): Promise<void> {
if (idleSweepTimer) {
clearInterval(idleSweepTimer);
idleSweepTimer = null;
}
for (const name of [...activeWorkers.keys()]) {
void this.#disposeWorker(name);
const work = Promise.allSettled(
[...activeWorkers.keys()].map((name) => this.#disposeWorker(name)),
);
if (!(await settleWithin(work, FINAL_FLUSH_TIMEOUT_MS))) {
console.warn('[localworker] final flush timed out');
}
}
async reconstructDeployArgs(workerName: string, row: SubdomainRow) {
@@ -0,0 +1,175 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { Actor } from '../../core/actor.js';
import type { PuterServer } from '../../server.js';
import { setupTestServer } from '../../testUtil.js';
import type { ResolvedShare } from './ShareService.js';
describe('ShareNotificationService shutdown drain', () => {
let server: PuterServer;
beforeEach(async () => {
server = await setupTestServer({
email: { from: '[email protected]', host: '127.0.0.1', port: 1 },
share_notify_limits: { emailBatchSeconds: 0.01 },
} as never);
});
afterEach(async () => {
vi.useRealTimers();
await server.services.shareNotification.onServerPrepareShutdown();
await server.services.shareNotification.onServerShutdown();
await server.shutdown();
vi.restoreAllMocks();
});
const queueInvite = async (email: string, issuerId = 7) => {
await server.services.shareNotification.notifyShared(
{
user: { id: issuerId, username: 'sender' },
effectiveApp: null,
} as Actor,
[
{
pending: true,
isNew: true,
recipientEmail: email,
name: 'shared.txt',
} as ResolvedShare,
],
);
};
it.each([false, true])(
'awaits an active timer send and its failure requeue %s',
async (failSend) => {
const email = `active-${failSend}@puter.local`;
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
const send = vi
.spyOn(server.clients.email, 'sendRaw')
.mockImplementation(async () => {
await held;
if (failSend) throw new Error('transport unavailable');
return null;
});
await queueInvite(email);
await vi.waitFor(() => expect(send).toHaveBeenCalledTimes(1));
const list = async () =>
server.stores.kv.list({
as: 'keys',
pattern: `share:digest:invite:${email}:`,
});
expect((await list()).res).toEqual([]);
let finished = false;
const draining = server.services.shareNotification
.onServerPrepareShutdown()
.then(() => {
finished = true;
});
try {
await new Promise((resolve) => setTimeout(resolve, 20));
expect(finished).toBe(false);
} finally {
release();
await draining;
}
expect((await list()).res).toHaveLength(failSend ? 1 : 0);
},
);
it('waits for an active send before draining entries queued meanwhile', async () => {
let releaseFirst!: () => void;
let releaseTail!: () => void;
const first = new Promise<void>((resolve) => {
releaseFirst = resolve;
});
const tail = new Promise<void>((resolve) => {
releaseTail = resolve;
});
const send = vi
.spyOn(server.clients.email, 'sendRaw')
.mockImplementationOnce(async () => {
await first;
return null;
})
.mockImplementationOnce(async () => {
await tail;
return null;
});
await queueInvite('[email protected]');
await vi.waitFor(() => expect(send).toHaveBeenCalledTimes(1));
(
server.services.shareNotification.config.share_notify_limits as {
emailBatchSeconds: number;
}
).emailBatchSeconds = 600;
await queueInvite('[email protected]', 8);
let finished = false;
const draining = server.services.shareNotification
.onServerPrepareShutdown()
.then(() => {
finished = true;
});
try {
releaseFirst();
await vi.waitFor(() => expect(send).toHaveBeenCalledTimes(2));
expect(finished).toBe(false);
} finally {
releaseFirst();
releaseTail();
await draining;
}
expect(send).toHaveBeenCalledTimes(2);
});
it.each(['onServerPrepareShutdown', 'onServerShutdown'] as const)(
'bounds %s while an active send is hung',
async (hook) => {
let release!: () => void;
const held = new Promise<void>((resolve) => {
release = resolve;
});
const send = vi
.spyOn(server.clients.email, 'sendRaw')
.mockImplementationOnce(async () => {
await held;
return null;
});
await queueInvite(`hung-${hook}@puter.local`);
await vi.waitFor(() => expect(send).toHaveBeenCalledTimes(1));
vi.useFakeTimers();
try {
let finished = false;
const draining = server.services.shareNotification[hook]().then(
() => {
finished = true;
},
);
await vi.advanceTimersByTimeAsync(4_999);
expect(finished).toBe(false);
await vi.advanceTimersByTimeAsync(1);
await draining;
} finally {
release();
vi.useRealTimers();
}
},
);
it('drains a digest queued after prepare at shutdown', async () => {
await server.services.shareNotification.onServerPrepareShutdown();
const send = vi
.spyOn(server.clients.email, 'sendRaw')
.mockResolvedValue(null);
(
server.services.shareNotification.config.share_notify_limits as {
emailBatchSeconds: number;
}
).emailBatchSeconds = 600;
await queueInvite('[email protected]');
await server.services.shareNotification.onServerShutdown();
expect(send).toHaveBeenCalledTimes(1);
expect(send.mock.calls[0][0].to).toBe('[email protected]');
await new Promise((resolve) => setTimeout(resolve, 30));
expect(send).toHaveBeenCalledTimes(1);
});
});
@@ -20,7 +20,11 @@
import { createHash, randomUUID } from 'node:crypto';
import type { Actor } from '../../core/actor';
import { checkRateLimit } from '../../core/http/middleware/rateLimit.js';
import { runWithConcurrencyLimitSettled } from '../../util/concurrency.js';
import {
runWithConcurrencyLimitSettled,
settleWithin,
} from '../../util/concurrency.js';
import { FINAL_FLUSH_TIMEOUT_MS } from '../../util/gracefulShutdown.js';
import { PuterService } from '../types';
import {
digestItemPaths,
@@ -178,6 +182,7 @@ export class ShareNotificationService extends PuterService {
* queued sends live in KV, where any node's flush can pick them up.
*/
#digestTimers = new Map<string, ReturnType<typeof setTimeout>>();
#digestFlushes = new Map<string, Promise<void>>();
/** The recovery sweep; see `onServerStart`. */
#digestSweep: ReturnType<typeof setInterval> | null = null;
@@ -199,12 +204,46 @@ export class ShareNotificationService extends PuterService {
/** Send what's still waiting while the transport is alive to send it. */
override async onServerPrepareShutdown(): Promise<void> {
await this.#drainDigests();
}
override async onServerShutdown(): Promise<void> {
await this.#drainDigests();
}
async #drainDigests(): Promise<void> {
if (this.#digestSweep) clearInterval(this.#digestSweep);
this.#digestSweep = null;
const keys = [...this.#digestTimers.keys()];
const keys = new Set([
...this.#digestTimers.keys(),
...this.#digestFlushes.keys(),
]);
for (const [, timer] of this.#digestTimers) clearTimeout(timer);
this.#digestTimers.clear();
await Promise.all(keys.map((key) => this.#flushDigest(key)));
const active = [...this.#digestFlushes.values()];
const deadline = Date.now() + FINAL_FLUSH_TIMEOUT_MS;
const drain = async () => {
try {
await Promise.allSettled(active);
if (Date.now() >= deadline) return;
for (const [key, timer] of this.#digestTimers) {
keys.add(key);
clearTimeout(timer);
}
this.#digestTimers.clear();
await Promise.allSettled(
[...keys].map((key) => this.#flushDigest(key)),
);
} catch (err) {
console.warn(
'[share-notify] shutdown digest flush failed:',
err,
);
}
};
if (!(await settleWithin(drain(), FINAL_FLUSH_TIMEOUT_MS))) {
console.warn('[share-notify] final flush timed out');
}
}
/** The recovery sweep, for tests that can't wait out an interval. */
@@ -878,25 +917,37 @@ export class ShareNotificationService extends PuterService {
* everything queued), and across regions the KV `take` is the real claim,
* handing each entry to exactly one flusher.
*/
async #flushDigest(key: string): Promise<void> {
await withRedisLock(
this.clients.redis,
`share:digest:lock:${key}`,
() => this.#flushDigestLocked(key),
{
ttlMs: DIGEST_LOCK_SECONDS * 1000,
// No lock beats no mail; `take` still prevents double sends.
onUnavailable: 'run',
onBusy: () => {
console.log(
'[share-notify] another flush holds the lock:',
{
key,
},
);
#flushDigest(key: string): Promise<void> {
const active = this.#digestFlushes.get(key);
if (active) return active;
const flushing = this.#sendDigest(key);
this.#digestFlushes.set(key, flushing);
return flushing;
}
async #sendDigest(key: string): Promise<void> {
try {
await withRedisLock(
this.clients.redis,
`share:digest:lock:${key}`,
() => this.#flushDigestLocked(key),
{
ttlMs: DIGEST_LOCK_SECONDS * 1000,
// No lock beats no mail; `take` still prevents double sends.
onUnavailable: 'run',
onBusy: () => {
console.log(
'[share-notify] another flush holds the lock:',
{ key },
);
},
},
},
);
);
} catch (err) {
console.warn('[share-notify] digest flush failed:', err);
} finally {
this.#digestFlushes.delete(key);
}
}
async #flushDigestLocked(key: string): Promise<void> {
+10 -4
View File
@@ -256,7 +256,12 @@ export interface IPreludeConfig {
* an RCS agent provisioned in the Prelude account to actually use RCS.
*/
preferredChannel?:
'sms' | 'rcs' | 'whatsapp' | 'viber' | 'zalo' | 'telegram';
| 'sms'
| 'rcs'
| 'whatsapp'
| 'viber'
| 'zalo'
| 'telegram';
}
/**
@@ -1375,7 +1380,9 @@ export interface WithLifecycle extends Object {
/**
* When the drain starts, before `onServerPrepareShutdown`, which waits for
* every connection (upgraded sockets included) to close or the drain to
* end.
* end. Hooks run concurrently, and the server keeps accepting connections
* until all of them settle, so a hook can hold the listener open (until
* traffic stops arriving, say) by returning a promise.
*/
onServerDrainStart?: (info?: ServerDrainInfo) => Promise<void> | void;
}
@@ -1386,8 +1393,7 @@ export interface ServerDrainInfo {
}
export interface WithCostsReporting extends WithLifecycle {
getReportedCosts?: () =>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
getReportedCosts?: () => // eslint-disable-next-line @typescript-eslint/no-explicit-any
| Promise<Record<string, any>[]>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| Record<string, any>[];
+63 -10
View File
@@ -3,23 +3,25 @@
*
* This file is part of Puter.
*
* Puter is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published
* by the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
* Puter is free software: you can redistribute it and/or modify it under the
* terms of the GNU Affero General Public License as published by the Free
* Software Foundation, either version 3 of the License, or (at your option) any
* later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
* This program is distributed in the hope that it will be useful, but WITHOUT
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
* FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
* details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <https://www.gnu.org/licenses/>.
* along with this program. If not, see
* [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/).
*/
import { describe, expect, it } from 'vitest';
import { afterEach, describe, expect, it, vi } from 'vitest';
import {
awaitLater,
settleWithin,
runWithConcurrencyLimit,
runWithConcurrencyLimitSettled,
} from './concurrency.ts';
@@ -221,3 +223,54 @@ describe('awaitLater', () => {
expect(await awaitLater(Promise.resolve(42))).toBe(42);
});
});
describe('settleWithin', () => {
afterEach(() => vi.useRealTimers());
it('returns true and clears the timer when work settles in time', async () => {
vi.useFakeTimers();
const work = Promise.withResolvers<void>();
const waiting = settleWithin(work.promise, 5000);
expect(vi.getTimerCount()).toBe(1);
work.resolve();
await expect(waiting).resolves.toBe(true);
expect(vi.getTimerCount()).toBe(0);
});
it('returns false and clears the timer when time runs out', async () => {
vi.useFakeTimers();
const cleared = vi.spyOn(globalThis, 'clearTimeout');
try {
const waiting = settleWithin(new Promise(() => {}), 5000);
await vi.advanceTimersByTimeAsync(5000);
await expect(waiting).resolves.toBe(false);
expect(cleared).toHaveBeenCalledTimes(1);
expect(vi.getTimerCount()).toBe(0);
} finally {
cleared.mockRestore();
}
});
it.each([false, true])(
'handles a rejection after timeout %s',
async (afterTimeout) => {
const unhandled: unknown[] = [];
const onUnhandled = (reason: unknown) => unhandled.push(reason);
process.on('unhandledRejection', onUnhandled);
try {
const work = Promise.withResolvers<void>();
const waiting = settleWithin(
work.promise,
afterTimeout ? 0 : 5000,
);
if (afterTimeout) await expect(waiting).resolves.toBe(false);
work.reject(new Error('work failed'));
await expect(waiting).resolves.toBe(!afterTimeout);
await new Promise((resolve) => setTimeout(resolve, 10));
expect(unhandled).toEqual([]);
} finally {
process.off('unhandledRejection', onUnhandled);
}
},
);
});
+29
View File
@@ -110,3 +110,32 @@ export function awaitLater<T>(promise: Promise<T>): Promise<T> {
})();
return promise;
}
/**
* Waits for `promise` for at most `ms`, without rejecting. Resolves `true` if
* it settled in time, `false` if the time ran out first.
*/
export async function settleWithin(
promise: Promise<unknown>,
ms: number,
): Promise<boolean> {
let timer: ReturnType<typeof setTimeout> | undefined;
const settled = async () => {
try {
await promise;
} catch {
// Callers report work failures.
}
return true;
};
try {
return await Promise.race([
settled(),
new Promise<boolean>((resolve) => {
timer = setTimeout(() => resolve(false), ms);
}),
]);
} finally {
clearTimeout(timer);
}
}
+46
View File
@@ -83,6 +83,52 @@ describe('createGracefulShutdown', () => {
expect(drainInfo).toEqual([{ drainMs: 90_000 }]);
});
it('counts time spent in prepareShutdown against the drain window', async () => {
const exit = vi.fn((code: number) => events.push(`exit:${code}`));
const target = makeTarget();
const trigger = createGracefulShutdown(
{
...target,
prepareShutdown: async (info) => {
await target.prepareShutdown(info);
await new Promise((resolve) => setTimeout(resolve, 15_000));
},
},
{ drainMs: 90_000, exit },
);
const done = trigger('SIGTERM');
await vi.advanceTimersByTimeAsync(89_999);
expect(events).toEqual(['prepare']);
await vi.advanceTimersByTimeAsync(1);
await done;
expect(events).toEqual(['prepare', 'shutdown', 'telemetry', 'exit:0']);
});
it('waits for prepareShutdown when it outlasts the drain window', async () => {
const exit = vi.fn((code: number) => events.push(`exit:${code}`));
const target = makeTarget();
const trigger = createGracefulShutdown(
{
...target,
prepareShutdown: async (info) => {
await target.prepareShutdown(info);
await new Promise((resolve) => setTimeout(resolve, 2_000));
},
},
{ drainMs: 1_000, exit },
);
const done = trigger('SIGTERM');
await vi.advanceTimersByTimeAsync(1_999);
expect(events).toEqual(['prepare']);
await vi.advanceTimersByTimeAsync(1);
await done;
expect(events).toEqual(['prepare', 'shutdown', 'telemetry', 'exit:0']);
});
it('SIGINT with a drain function returning 0 for SIGINT exits without waiting', async () => {
const exit = vi.fn((code: number) => events.push(`exit:${code}`));
const drainMs = (signal: string) => (signal === 'SIGTERM' ? 90_000 : 0);
+12 -6
View File
@@ -27,9 +27,9 @@ export interface ShutdownTarget {
export interface GracefulShutdownOptions {
/**
* How long to keep serving after new connections stop, before teardown. A
* function is called with the signal that triggered shutdown, so a
* deployment can drain on `SIGTERM` but exit `SIGINT` (Ctrl-C) instantly.
* How long to keep serving after the signal, before teardown. A function is
* called with the signal that triggered shutdown, so a deployment can drain
* on `SIGTERM` but exit `SIGINT` (Ctrl-C) instantly.
*/
drainMs: number | ((signal: string) => number);
/** Cap on the final telemetry flush. */
@@ -45,6 +45,8 @@ export interface GracefulShutdownOptions {
export const DEFAULT_TELEMETRY_SHUTDOWN_TIMEOUT_MS = 5_000;
export const DEFAULT_SHUTDOWN_HARD_LIMIT_MS = 25_000;
/** Cap on one component's final flush, so teardown fits the hard limit. */
export const FINAL_FLUSH_TIMEOUT_MS = 5_000;
/**
* The process's one shutdown sequence: stop accepting connections, drain, tear
@@ -79,14 +81,18 @@ export const createGracefulShutdown = (
exitOnce(1);
}, drain + hardLimitMs);
let exitCode = 0;
// The drain window starts at the signal, so time `prepareShutdown`
// spends holding the listener open counts against it.
const drained =
drain > 0
? new Promise<void>((resolve) => setTimeout(resolve, drain))
: Promise.resolve();
try {
await target.prepareShutdown({ drainMs: drain });
} catch (err) {
console.error('[shutdown] prepareShutdown failed', err);
}
if (drain > 0) {
await new Promise<void>((resolve) => setTimeout(resolve, drain));
}
await drained;
try {
await target.shutdown();
} catch (err) {