diff --git a/src/index.ts b/src/index.ts --- a/src/index.ts +++ b/src/index.ts @@ -4,6 +4,7 @@ export { AsyncIterableTask } from './stream-task.ts'; export { channel } from './channel.ts'; export type { ChannelOptions } from './channel.ts'; +export { pushChannel, PushChannel } from './push-channel.ts'; export { workers } from './worker-pool.ts'; export { transfer } from './transfer.ts'; export { assign } from './assign.ts'; diff --git a/src/push-channel.ts b/src/push-channel.ts new file mode 100644 --- /dev/null +++ b/src/push-channel.ts @@ -0,0 +1,54 @@ +/** + * A simple push-based async queue that implements AsyncIterable. + * Values are pushed via `.send(item)` and the iteration ends when `.close()` is called. + */ +export class PushChannel implements AsyncIterable { + private readonly queue: T[] = []; + private readonly resolvers: Array<(result: IteratorResult) => void> = []; + private closed = false; + + /** Push a value to waiting consumers (or enqueue if none are waiting). */ + send(item: T): void { + if (this.closed) throw new Error('PushChannel is closed'); + if (this.resolvers.length > 0) { + const resolve = this.resolvers.shift()!; + resolve({ value: item, done: false }); + } else { + this.queue.push(item); + } + } + + /** Signal end-of-stream; iteration will finish after queued items. */ + close(): void { + this.closed = true; + for (const resolve of this.resolvers.splice(0)) { + resolve({ value: undefined as unknown as T, done: true }); + } + } + + [Symbol.asyncIterator](): AsyncIterator { + return { + next: (): Promise> => { + if (this.queue.length > 0) { + return Promise.resolve({ value: this.queue.shift()!, done: false }); + } + if (this.closed) { + return Promise.resolve({ value: undefined as unknown as T, done: true }); + } + return new Promise>((resolve) => { + this.resolvers.push(resolve); + }); + }, + return: (): Promise> => { + return Promise.resolve({ value: undefined as unknown as T, done: true }); + }, + }; + } +} + +/** + * Create a push-based channel with `.send(item)` and `.close()` for async iteration. + */ +export function pushChannel(_opts?: { highWaterMark?: number }): PushChannel { + return new PushChannel(); +} diff --git a/src/serve/index.ts b/src/serve/index.ts --- a/src/serve/index.ts +++ b/src/serve/index.ts @@ -1,2 +1,3 @@ export type { ListenArgs, ListenOptions, Balance, ServerThreads, ServerThreadsOptions } from './types.ts'; export { leastConns, roundRobin } from './strategies.ts'; +export { listen } from './listen.ts'; diff --git a/src/serve/listen.ts b/src/serve/listen.ts new file mode 100644 --- /dev/null +++ b/src/serve/listen.ts @@ -0,0 +1,80 @@ +import { Socket } from 'node:net'; +import type { Server } from 'node:net'; +import type { PushChannel } from '../push-channel.ts'; +import type { Tuple } from '../shared/tuple.ts'; +import type { Int32Atomic } from '../shared/int32-atomic.ts'; +import type { ListenOptions } from './types.ts'; + +/** + * Drain fds from `ch` and emit each as a `'connection'` on `server`. Maintains + * the per-worker counter (decrements on each socket's 'close' event). Resolves + * after the channel ends and either: (a) all active sockets have closed, or + * (b) `opts.drainTimeout` elapses — in which case any remaining sockets are + * force-destroyed. + */ +export async function listen( + server: Server, + ch: PushChannel, + counters: Tuple, + slot: number, + opts: Required, +): Promise { + const active = new Set(); + let onDrained: (() => void) | null = null; + + const tryResolve = () => { + if (active.size === 0 && onDrained) { + const cb = onDrained; + onDrained = null; + cb(); + } + }; + + for await (const fd of ch) { + const socket = new Socket({ fd, readable: true, writable: true }); + socket.setNoDelay(true); + active.add(socket); + socket.once('close', () => { + active.delete(socket); + counters.elements[slot].sub(1); + tryResolve(); + }); + server.emit('connection', socket); + } + + // Channel closed; begin drain. + const httpSrv = server as Server & { + closeIdleConnections?: () => void; + closeAllConnections?: () => void; + }; + httpSrv.closeIdleConnections?.(); + + await new Promise((resolve) => { + if (active.size === 0) return resolve(); + + onDrained = resolve; + + const timer = setTimeout(() => { + // Force-close all remaining sockets directly. + for (const socket of active) { + socket.destroy(); + } + // Also try the server-level helpers if available. + httpSrv.closeAllConnections?.(); + // Poll until active drains (socket 'close' events fire async after destroy). + const poll = setInterval(() => { + if (active.size === 0) { + clearInterval(poll); + onDrained = null; + resolve(); + } + }, 10); + }, opts.drainTimeout); + + // Override resolve to also clear the timer. + onDrained = () => { + clearTimeout(timer); + resolve(); + }; + }); +} diff --git a/src/serve/types.ts b/src/serve/types.ts --- a/src/serve/types.ts +++ b/src/serve/types.ts @@ -1,4 +1,4 @@ -import type { Channel } from '../channel.ts'; +import type { PushChannel } from '../push-channel.ts'; import type { Tuple } from '../shared/tuple.ts'; import type { Int32Atomic } from '../shared/int32-atomic.ts'; import type { WorkerHandle } from '../runner.ts'; @@ -14,7 +14,7 @@ * The shape may evolve; users only ever spread it into `listen()`. */ export type ListenArgs = readonly [ - channel: Channel, + channel: PushChannel, counters: Tuple, slot: number, opts: Required, diff --git a/test/serve/listen.test.ts b/test/serve/listen.test.ts new file mode 100644 --- /dev/null +++ b/test/serve/listen.test.ts @@ -0,0 +1,71 @@ +import { describe, it } from 'node:test'; +import assert from 'node:assert/strict'; +import { once } from 'node:events'; +import { openSync } from 'node:fs'; +import { createServer as createNetServer, connect } from 'node:net'; +import { createServer } from 'node:http'; +import { pushChannel as channel, shared, int32atomic } from 'moroutine'; +import { listen } from 'moroutine/serve'; + +// Helper: accept a real TCP connection on a scratch listener and yield its fd. +// We dup the fd via /dev/fd so that the new fd is not registered in libuv's +// event loop — this lets new Socket({ fd }) succeed in the same thread. +async function acquireLocalFd(): Promise<{ fd: number; close: () => void }> { + const srv = createNetServer(); + srv.listen(0); + await once(srv, 'listening'); + const port = (srv.address() as any).port; + const client = connect(port); + const [peer] = (await once(srv, 'connection')) as [any]; + const origFd: number = peer._handle.fd; + const fd = openSync(`/dev/fd/${origFd}`, 'r+'); // dup — new fd, not tracked by libuv + peer.pause(); + peer._handle = null; + peer.destroy(); + return { + fd, + close: () => { + srv.close(); + client.destroy(); + }, + }; +} + +describe('listen()', () => { + it('emits received fds as connections on the user server', { timeout: 5000 }, async () => { + const ch = channel({ highWaterMark: 4 }); + const counters = shared([int32atomic]); + const http = createServer((req, res) => { + res.end('ok'); + }); + const drained = listen(http, ch, counters, 0, { drainTimeout: 5_000 }); + + const { fd, close: closeSrc } = await acquireLocalFd(); + ch.send(fd); + // Wait briefly for the socket to be emitted, then destroy the client so + // the connection closes and drained can resolve. + await new Promise((r) => setTimeout(r, 50)); + closeSrc(); // destroy client so socket closes naturally + ch.close(); + await drained; + }); + + it('decrements counter on socket close', { timeout: 5000 }, async () => { + const ch = channel({ highWaterMark: 4 }); + const counters = shared([int32atomic]); + counters.elements[0].store(1); // simulate main having incremented + const http = createServer((req, res) => { + res.end('ok'); + }); + const drained = listen(http, ch, counters, 0, { drainTimeout: 5_000 }); + + const { fd, close: closeSrc } = await acquireLocalFd(); + ch.send(fd); + await new Promise((r) => setTimeout(r, 50)); + closeSrc(); // destroy client so socket closes → counter decrements + await new Promise((r) => setTimeout(r, 100)); + assert.equal(counters.elements[0].load(), 0); + ch.close(); + await drained; + }); +});