fast-worker-channel
Version:
High-performance communication channels for Node.js worker threads
251 lines • 8.68 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.Channel = void 0;
exports.handshake = handshake;
const msgpack_1 = require("@msgpack/msgpack");
const CLOSED_FLAG = 1;
class RingBuffer {
constructor(shared, bufferSize) {
this.closed = false;
this.header = new Int32Array(shared, 0, 4);
if (bufferSize != null) {
this.header[RingBuffer.SLOTS.SIZE] = bufferSize;
this.bufferSize = bufferSize;
}
else {
this.bufferSize = this.header[RingBuffer.SLOTS.SIZE];
}
this.lengthView = new Uint32Array(shared, RingBuffer.HEADER_BYTES, 1);
this.payload = new Uint8Array(shared, RingBuffer.HEADER_BYTES);
}
advance(pointer, by) {
const next = pointer + by;
if (next >= this.bufferSize) {
return next - this.bufferSize;
}
return next;
}
writeLength(writePos, length) {
this.lengthView[0] = length;
new DataView(this.payload.buffer, this.payload.byteOffset + writePos, 4).setUint32(0, length, true);
}
readLength(tailPos) {
return new DataView(this.payload.buffer, this.payload.byteOffset + tailPos, 4).getUint32(0, true);
}
close() {
if (this.closed) {
return;
}
this.closed = true;
Atomics.store(this.header, RingBuffer.SLOTS.FLAGS, CLOSED_FLAG);
Atomics.notify(this.header, RingBuffer.SLOTS.FLAGS, Infinity);
Atomics.notify(this.header, RingBuffer.SLOTS.HEAD, Infinity);
}
}
RingBuffer.SLOTS = {
SIZE: 0,
HEAD: 1,
TAIL: 2,
FLAGS: 3
};
RingBuffer.HEADER_BYTES = Object.keys(RingBuffer.SLOTS).length * 4;
class WritePort {
constructor(data) {
if (typeof data === 'number') {
this.shared = new SharedArrayBuffer(RingBuffer.HEADER_BYTES + data);
this.buffer = new RingBuffer(this.shared, data);
}
else {
this.shared = data;
this.buffer = new RingBuffer(this.shared);
}
Atomics.store(this.buffer.header, RingBuffer.SLOTS.HEAD, 0);
Atomics.store(this.buffer.header, RingBuffer.SLOTS.TAIL, 0);
}
write(message) {
if (message === null) {
return;
}
const encoded = (0, msgpack_1.encode)(message);
const totalLen = encoded.length + 4;
const head = Atomics.load(this.buffer.header, RingBuffer.SLOTS.HEAD);
const tail = Atomics.load(this.buffer.header, RingBuffer.SLOTS.TAIL);
const used = head >= tail
? head - tail
: this.buffer.bufferSize - (tail - head);
const free = this.buffer.bufferSize - used;
if (totalLen > free) {
throw new Error('Buffer full');
}
let writePos = head;
if (writePos + totalLen > this.buffer.bufferSize) {
writePos = 0;
}
this.buffer.writeLength(writePos, encoded.length);
const afterLen = writePos + 4;
const firstPart = Math.min(encoded.length, this.buffer.bufferSize - afterLen);
this.buffer.payload.set(encoded.subarray(0, firstPart), afterLen);
if (firstPart < encoded.length) {
this.buffer.payload.set(encoded.subarray(firstPart), 0);
}
const newHead = this.buffer.advance(writePos, totalLen);
Atomics.store(this.buffer.header, RingBuffer.SLOTS.HEAD, newHead);
Atomics.notify(this.buffer.header, RingBuffer.SLOTS.HEAD);
}
}
class ReadPort {
constructor(shared) {
this.shared = shared;
this.buffer = new RingBuffer(shared);
}
async *read() {
try {
while (true) {
const flags = Atomics.load(this.buffer.header, RingBuffer.SLOTS.FLAGS);
if (flags === CLOSED_FLAG) {
break;
}
const head = Atomics.load(this.buffer.header, RingBuffer.SLOTS.HEAD);
let tail = Atomics.load(this.buffer.header, RingBuffer.SLOTS.TAIL);
if (head === tail) {
Atomics.wait(this.buffer.header, RingBuffer.SLOTS.HEAD, head);
continue;
}
if (tail + 4 > this.buffer.bufferSize) {
Atomics.store(this.buffer.header, RingBuffer.SLOTS.TAIL, 0);
continue;
}
const len = this.buffer.readLength(tail);
if (len === 0) {
Atomics.store(this.buffer.header, RingBuffer.SLOTS.TAIL, 0);
continue;
}
let value;
const afterLen = tail + 4;
if (afterLen + len > this.buffer.bufferSize) {
const firstPart = this.buffer.bufferSize - afterLen;
value = new Uint8Array(len);
value.set(this.buffer.payload.subarray(afterLen, this.buffer.bufferSize), 0);
value.set(this.buffer.payload.subarray(0, len - firstPart), firstPart);
}
else {
value = this.buffer.payload.subarray(afterLen, afterLen + len);
}
const newTail = this.buffer.advance(tail, len + 4);
Atomics.store(this.buffer.header, RingBuffer.SLOTS.TAIL, newTail);
yield (0, msgpack_1.decode)(value);
}
}
finally {
// @ts-ignore
this.buffer = null;
// @ts-ignore
this.shared = null;
}
}
}
class SABProducer {
constructor(writer, inbound) {
this.writer = writer;
this.outboundBuffer = writer.shared;
this.inboundBuffer = inbound ?? writer.shared;
}
send(msg) {
this.writer.write(msg);
}
close() {
this.writer.write(CLOSED_FLAG);
}
}
class SABConsumer {
constructor(reader, outbound) {
this.reader = reader;
this.inboundBuffer = reader.shared;
this.outboundBuffer = outbound ?? reader.shared;
}
async *receive() {
for await (const x of this.reader.read()) {
if (x === CLOSED_FLAG) {
break;
}
yield x;
}
}
}
class SABProducerConsumer {
constructor(writer, reader) {
this.writer = writer;
this.reader = reader;
this.inboundBuffer = reader.shared;
this.outboundBuffer = writer.shared;
}
send(msg) {
this.writer.write(msg);
}
close() {
this.writer.write(CLOSED_FLAG);
}
async *receive() {
for await (const x of this.reader.read()) {
if (x === CLOSED_FLAG) {
break;
}
yield x;
}
}
}
class Channel {
static createUnidirectional(bufferSize = 1024 * 1024) {
const writer = new WritePort(bufferSize);
const reader = new ReadPort(writer.shared);
return {
producer: new SABProducer(writer),
consumer: new SABConsumer(reader)
};
}
static createBidirectional(bufferSize = 1024 * 1024) {
const w1 = new WritePort(bufferSize);
const r1 = new ReadPort(w1.shared);
const w2 = new WritePort(bufferSize);
const r2 = new ReadPort(w2.shared);
return {
producer: new SABProducerConsumer(w1, r2),
consumer: new SABProducerConsumer(w2, r1),
};
}
static adoptConsumer(inbound) {
const reader = new ReadPort(inbound);
return new SABConsumer(reader);
}
static adoptProducerConsumer(inbound, outbound) {
const reader = new ReadPort(inbound);
const writer = new WritePort(outbound);
return new SABProducerConsumer(writer, reader);
}
}
exports.Channel = Channel;
async function handshake(type, worker, arg2) {
const hello = new Promise((resolve, reject) => {
const timeout = setTimeout(() => {
reject(new Error('Handshake timeout'));
}, 10000);
worker.once('message', (e) => {
if (e.type === 'hello') {
resolve();
}
else {
reject(new Error('Handshake failed'));
}
clearTimeout(timeout);
});
});
if (type === 'unidirectional') {
worker.postMessage({ inboundBuffer: arg2.inboundBuffer, type: 'hello' });
}
else if (type === 'bidirectional') {
const { inboundBuffer, outboundBuffer } = arg2;
worker.postMessage({ inboundBuffer, outboundBuffer, type: 'hello' });
}
return hello;
}
//# sourceMappingURL=port.js.map