Приручаем многопоточность в Node.js (часть 5: автомасштабирование под нагрузку)

class Coordinator extends Worker {
  #messages = new Map(); // хранилище всех обрабатываемых сообщений

  #p32 = new Uint32Array(2);               // префикс данных = [messageID, dataSize]
  #p8  = new Uint8Array(this.#p32.buffer); // ... и его uint8-проекция
  #data = [this.#p8]; // массив для склейки частей в двоичный блок
  #message2data;      // функция получения двоичных данных из сообщения и помещения в массив

  #data2buffer() {
    // объединяем префикс и данные
    const buf = Buffer.concat(this.#data);
    // записываем в префикс результирующую длину
    buf.writeUInt32LE(buf.length - 8, 4); // Uint32Array.BYTES_PER_ELEMENT
    // восстанавливаем состояние массива
    this.#data.length = 1;
    // возвращаем итоговый двоичный контент
    return buf.buffer;
  };

  constructor(filename, options) {
    // сигнальный канал
    const {port1, port2} = new MessageChannel();

    /* доинициализируем необходимые опции
      {
        workerData : {
          workerType : 'coordinator'
        , port       : port2
        }
      , transferList : [port2]
      }
    */
    ((options ??= {}).workerData ??= {}).workerType = 'coordinator';
    ((options ??= {}).workerData ??= {}).port = port2;
    ((options ??= {}).transferList ??= []).push(port2);
    super(filename, options);

    const {dataField, dataArray} = options;
    this.#p32[0] = 0; // ID

    this.#message2data = dataField
      ? message => this.#data.push(dataField(message))
      : message => this.#data.push(...dataArray(message));

    const messages = this.#messages;
    port1.on('message', ({threadId, port}) => {
      // ассоциируем открывшийся порт с конкретным потоком
      port.threadId = threadId;
      this.emit('port.open', port);

      port
        .on('message', result => {
          // из результата обработки по ID получаем исходное сообщение ...
          const message = messages.get(result.id);
          if (message) {
            messages.delete(result.id);
            // ... и передаем вместе с результатом
            this.emit('message', result, message);
          }
        })
        .on('close', () => {
          // при закрытии порта - отписываемся от него
          this.emit('port.close', port);
          port.removeAllListeners();
        });
    });
  }

  // передача двоичных данных в поток
  postMessage(message) {
    // сохраняем объект сообщения в хранилище
    this.#messages.set(this.#p32[0], message);
    // добавляем в общий массив для склейки одно или несколько полей
    this.#message2data(message);
    // формируем целевой контейнер данных
    const buffer = this.#data2buffer();
    // передаем двоичный контент с префиксом в поток по ссылке
    super.postMessage(buffer, [buffer]);
    // id = (id + 1) % 0x10000000
    this.#p32[0]++;
    this.#p32[0] &= 0x0FFFFFFF;
  }
}

class WorkersPool {
  #options; // кэш аргументов конструктора

  #mainWorker;              // основной рабочий поток
  #workersPool = [];        // пул свободных дополнительных потоков
  #workersSet  = new Set(); // полный набор всех активных потоков
  #workersRemain;           // потоков еще доступно к созданию

  #queue;      // очередь на кольцевом буфере

  #checking;   // признак активности проверки состояния

  #queueMore1; // "длина" очереди превышает возможности 1 потока
  #queueMoreW; //                                   ... всех существующих потоков

  #checkQueue(workers) {
    // если в очереди задач уже больше, чем потоков,
    // ... и самая старая задача висит дольше, чем обработали бы все активные потоки за время старта нового
    return this.#queue.length > workers && this.#queue[0].ts + this.#options['timeoutSpawn'] * workers < Date.now();
  }

  constructor(workerFile, options) {
    options['workerFile'] ??= workerFile;
    const {poolSize, queuePowMin, queuePowMax, timeoutIdle, intervalCheck} = options;
    this.#options = options;

    this.#workersRemain = poolSize;

    // очередь на кольцевом буфере
    this.#queue = new (require('./Pow2Buffer'))(queuePowMin, queuePowMax);

    // главный вспомогательный поток существует сразу и всегда
    this.#mainWorker = this.#createWorker();
    
    // ...
    
    // периодическая проверка очереди и пула дополнительных потоков
    const pool = this.#workersPool;
    setInterval(() => {
      // пора ли порождать еще один поток?
      this.#queueMoreW = this.#checkQueue(this.#workersSet.size);
      // пора ли отдавать задачи дополнительным потокам?
      this.#queueMore1 = this.#queueMoreW || this.#checkQueue(1);
      // закрытие простаивающих дополнительных потоков
      // ...
    }, intervalCheck);
  }

  #createWorkerIfPossible() {
    // проверяем возможность и необходимость (по состоянмю очереди) запуска потока
    if (this.#workersRemain > 0 && this.#checkQueue(this.#workersSet.size)) {
      this.#createWorker();
    }
  }

  #createWorker() {
    this.#workersRemain--;
    this.#workersRemain = -this.#workersRemain; // wrap flag
    // ...

    worker.on('online', () => {
      // продолжим пробовать стартовать следующий после паузы на timeoutSpawn
      setTimeout(() => {
        this.#workersRemain = -this.#workersRemain; // wrap flag
        this.#createWorkerIfPossible();
      }, this.#options['timeoutSpawn']);
    });

    return worker;
  }

  #destroyWorker(worker) {
    // убираем поток из общего набора
    this.#workersSet.delete(worker);
    // завершаем сам поток
    worker.terminate();
    // исключаем воркер из пула, если он там был
    const idx = this.#workersPool.indexOf(worker);
    if (idx >= 0) {
      this.#workersPool.splice(idx, 1);
    }
    // увеличиваем количество доступных к запуску
    this.#workersRemain++;
  }
}

Тут мы использовали #workersRemain одновременно и как счетчик доступных к запуску, и как признак (отрицательное значение) наличия запускающегося в текущий момент потока.

Второе замечание касается использования для пула свободных именно массива, а не Set, поскольку нам важен приоритет потоков при раздаче им заданий на обработку.

  constructor(workerFile, options) {
    // ...
    
    // основная точка приема сообщений из main-потока
    parentPort.on('message', message => {
      // флаг возможности и необходимости продолжать отправку
      let processed;
      const {lock : [lockState], _pulse} = this.#mainWorker;
      if (lockState == THREAD_FREE) { // если основной рабочий поток свободен
        processed = _pulse(message);  // ... отдаем сообщение основному потоку
        message = undefined;          // ... и обнуляем
      }
      // продолжаем "дергать" потоки, пока хоть какая-то очередь (порта или своя) непуста
      // ... и еще кто-то остался кто-то свободный
      while (!processed) {
        const worker = this.#queueMore1 && this.#workersPool.pop();
        if (worker) {
          processed = worker._pulse(message);
          message &&= undefined;
        }
        else {
          // перекладываем всю очередь порта в свою очередь с меткой времени
          const now = Date.now();
          message ??= receiveMessageOnPort(parentPort)?.message;
          while (message) {
            this.#queue.push(message);
            message.ts = now;
            message = receiveMessageOnPort(parentPort)?.message;
          }
          // отложенная проверка необходимости запуска нового потока
          this.#queueMoreW && this.#workersRemain > 0 && (this.#checking ??= setTimeout(() => {
            this.#createWorkerIfPossible();
            this.#checking = null;
          }, this.#options['intervalCheck']));
          return;
        }
      }
    });
  }
  // ...

  #pulseWorker(worker, message) {
    // если нам передали сообщение - пытаемся отдать его потоку
    let processed = !message || worker._send(message);
    // пока он говорит "уже все сразу обработал"...
    while (processed) {
      // извлекаем следующее сообщение из очереди порта основного потока
      const recv = receiveMessageOnPort(parentPort);
      // ... или своей локальной очереди
      message = recv?.message ?? this.#queue.shift();
      if (!message) {
        if (worker !== this.#mainWorker) {
          // когда все очереди закончились - возвращаем дополнительный поток в пул свободных
          const pool = this.#workersPool;
          !pool.includes(worker) && pool.push(worker);
        }
        return true;
      }
      // передаем данные в поток
      processed = worker._send(message);
    }
  };

const {
  Worker
, isMainThread
, parentPort
, workerData
, MessageChannel
, receiveMessageOnPort
, threadId
} = require('node:worker_threads');

const THREAD_FREE = -1;

class WorkersPool {
  #options; // кэш аргументов конструктора

  #mainWorker;              // основной рабочий поток
  #workersPool = [];        // пул свободных дополнительных потоков
  #workersSet  = new Set(); // полный набор всех активных потоков
  #workersRemain;           // потоков еще доступно к созданию

  #queue;      // очередь на кольцевом буфере

  #checking;   // признак активности проверки состояния

  #queueMore1; // "длина" очереди превышает возможности 1 потока
  #queueMoreW; //                                   ... всех существующих потоков

  #checkQueue(workers) {
    // если в очереди задач уже больше, чем потоков,
    // ... и самая старая задача висит дольше, чем обработали бы все активные потоки за время старта нового
    return this.#queue.length > workers && this.#queue[0].ts + this.#options['timeoutSpawn'] * workers < Date.now();
  }

  constructor(workerFile, options) {
    options['workerFile'] ??= workerFile;
    const {poolSize, queuePowMin, queuePowMax, timeoutIdle, intervalCheck} = options;
    this.#options = options;

    this.#workersRemain = poolSize;

    // очередь на кольцевом буфере
    this.#queue = new (require('./Pow2Buffer'))(queuePowMin, queuePowMax);

    // главный вспомогательный поток существует сразу и всегда
    this.#mainWorker = this.#createWorker();

    // основная точка приема сообщений из main-потока
    parentPort.on('message', message => {
      // флаг возможности и необходимости продолжать отправку
      let processed;
      const {lock : [lockState], _pulse} = this.#mainWorker;
      if (lockState == THREAD_FREE) { // если основной рабочий поток свободен
        processed = _pulse(message);  // ... отдаем сообщение основному потоку
        message = undefined;          // ... и обнуляем
      }
      // продолжаем "дергать" потоки, пока хоть какая-то очередь (порта или своя) непуста
      // ... и еще кто-то остался кто-то свободный
      while (!processed) {
        const worker = this.#queueMore1 && this.#workersPool.pop();
        if (worker) {
          processed = worker._pulse(message);
          message &&= undefined;
        }
        else {
          // перекладываем всю очередь порта в свою очередь с меткой времени
          const now = Date.now();
          message ??= receiveMessageOnPort(parentPort)?.message;
          while (message) {
            this.#queue.push(message);
            message.ts = now;
            message = receiveMessageOnPort(parentPort)?.message;
          }
          // отложенная проверка необходимости запуска нового потока
          this.#queueMoreW && this.#workersRemain > 0 && (this.#checking ??= setTimeout(() => {
            this.#createWorkerIfPossible();
            this.#checking = null;
          }, this.#options['intervalCheck']));
          return;
        }
      }
    });

    // периодическая проверка очереди и пула дополнительных потоков
    const pool = this.#workersPool;
    setInterval(() => {
      // пора ли порождать еще один поток?
      this.#queueMoreW = this.#checkQueue(this.#workersSet.size);
      // пора ли отдавать задачи дополнительным потокам?
      this.#queueMore1 = this.#queueMoreW || this.#checkQueue(1);
      // закрытие простаивающих дополнительных потоков
      if (pool.length) {
        const deadline = Date.now() - timeoutIdle;
        for (const {activity, lock : [lockState], _destroy} of pool) {
          activity < deadline && lockState == THREAD_FREE && _destroy();
        }
      }
    }, intervalCheck);
  }

  #createWorkerIfPossible() {
    // проверяем возможность и необходимость (по состоянмю очереди) запуска потока
    if (this.#workersRemain > 0 && this.#checkQueue(this.#workersSet.size)) {
      this.#createWorker();
    }
  }

  #createWorker() {
    this.#workersRemain--;
    this.#workersRemain = -this.#workersRemain; // wrap flag

    const shared = new SharedArrayBuffer(Int32Array.BYTES_PER_ELEMENT + Uint32Array.BYTES_PER_ELEMENT + this.#options['dataSize']);
    const data = new Uint8Array(shared);
    const lock = new Int32Array(shared, 0, 1); // lock/id - первые 4 байта data
    lock[0] = 0;

    // ждем, пока поток не проинициализируется и не сменит состояние, чтобы его "подергать"
    let worker;
    Atomics.waitAsync(lock, 0, 0)
      .value
      .then(() => lock[0] == THREAD_FREE && worker._pulse());
    
    const {port1, port2} = new MessageChannel();
    worker = new Worker(
      this.#options['workerFile']
    , {
        workerData : {
          workerType : 'worker'
        , shared
        , port       : port2 // передаем порт в рабочий поток
        }
      , transferList : [port2]
      }
    );

    // отправляем порт в основной поток, в Coordinator
    this.#options['port'].postMessage(
      {
        threadId : worker.threadId
      , port     : port1
      }
    , [port1]
    );

    this.#workersSet.add(worker);

    worker.id = worker.threadId;
    worker.port = port2;
    worker.lock = lock;
    worker.data = data;

    worker._pulse   = this.#pulseWorker.bind(this, worker);
    worker._destroy = this.#destroyWorker.bind(this, worker);
    worker._send    = this.#sendMessage.bind(this, worker);

    worker.on('online', () => {
      // продолжим пробовать стартовать следующий после паузы на timeoutSpawn
      setTimeout(() => {
        this.#workersRemain = -this.#workersRemain; // wrap flag
        this.#createWorkerIfPossible();
      }, this.#options['timeoutSpawn']);
    });

    return worker;
  }

  #pulseWorker(worker, message) {
    // если нам передали сообщение - пытаемся отдать его потоку
    let processed = !message || worker._send(message);
    // пока он говорит "уже все сразу обработал"...
    while (processed) {
      // извлекаем следующее сообщение из очереди порта основного потока
      const recv = receiveMessageOnPort(parentPort);
      // ... или своей локальной очереди
      message = recv?.message ?? this.#queue.shift();
      if (!message) {
        if (worker !== this.#mainWorker) {
          // когда все очереди закончились - возвращаем дополнительный поток в пул свободных
          const pool = this.#workersPool;
          !pool.includes(worker) && pool.push(worker);
        }
        return true;
      }
      // передаем данные в поток
      processed = worker._send(message);
    }
  };

  #sendMessage(worker, message) {
    const {lock, data, _pulse} = worker;
    // записываем из ui8-проекции в разделяемую память вместе с id
    data.set(new Uint8Array(message));
    // id - это первые 4 байта data
    const id = lock[0];
    // уведомляем поток
    Atomics.notify(lock, 0, 1);
    // фиксируем момент последней активности потока
    worker.activity = Date.now();
    // ждем, пока поток не обработает и не позовет нас
    const {value} = Atomics.waitAsync(lock, 0, id);
    if (value === 'not-equal') {
      // если он сразу успел обработать, то в него можно и дальше писать
      return true;
    }
    else {
      // если не сразу - "подергаем", когда освободится
      value.then(() => lock[0] == THREAD_FREE && _pulse());
    }
  }

  #destroyWorker(worker) {
    // убираем поток из общего набора
    this.#workersSet.delete(worker);
    // закрываем обе стороны ассоциированного MessageChannel
    worker.port.close();
    // завершаем сам поток
    worker.terminate();
    // исключаем воркер из пула, если он там был
    const idx = this.#workersPool.indexOf(worker);
    if (idx >= 0) {
      this.#workersPool.splice(idx, 1);
    }
    // увеличиваем количество доступных к запуску
    this.#workersRemain++;
  }
}

class Coordinator extends Worker {
  #messages = new Map(); // хранилище всех обрабатываемых сообщений

  #p32 = new Uint32Array(2);               // префикс данных = [messageID, dataSize]
  #p8  = new Uint8Array(this.#p32.buffer); // ... и его uint8-проекция
  #data = [this.#p8]; // массив для склейки частей в двоичный блок
  #message2data;      // функция получения двоичных данных из сообщения и помещения в массив

  #data2buffer() {
    // объединяем префикс и данные
    const buf = Buffer.concat(this.#data);
    // записываем в префикс результирующую длину
    buf.writeUInt32LE(buf.length - 8, 4); // Uint32Array.BYTES_PER_ELEMENT
    // восстанавливаем состояние массива
    this.#data.length = 1;
    // возвращаем итоговый двоичный контент
    return buf.buffer;
  };

  constructor(filename, options) {
    // сигнальный канал
    const {port1, port2} = new MessageChannel();

    /* доинициализируем необходимые опции
      {
        workerData : {
          workerType : 'coordinator'
        , port       : port2
        }
      , transferList : [port2]
      }
    */
    ((options ??= {}).workerData ??= {}).workerType = 'coordinator';
    ((options ??= {}).workerData ??= {}).port = port2;
    ((options ??= {}).transferList ??= []).push(port2);
    super(filename, options);

    const {dataField, dataArray} = options;
    this.#p32[0] = 0; // ID

    this.#message2data = dataField
      ? message => this.#data.push(dataField(message))
      : message => this.#data.push(...dataArray(message));

    const messages = this.#messages;
    port1.on('message', ({threadId, port}) => {
      // ассоциируем открывшийся порт с конкретным потоком
      port.threadId = threadId;
      this.emit('port.open', port);

      port
        .on('message', result => {
          // из результата обработки по ID получаем исходное сообщение ...
          const message = messages.get(result.id);
          if (message) {
            messages.delete(result.id);
            // ... и передаем вместе с результатом
            this.emit('message', result, message);
          }
        })
        .on('close', () => {
          // при закрытии порта - отписываемся от него
          this.emit('port.close', port);
          port.removeAllListeners();
        });
    });
  }

  // передача двоичных данных в поток
  postMessage(message) {
    // сохраняем объект сообщения в хранилище
    this.#messages.set(this.#p32[0], message);
    // добавляем в общий массив для склейки одно или несколько полей
    this.#message2data(message);
    // формируем целевой контейнер данных
    const buffer = this.#data2buffer();
    // передаем двоичный контент с префиксом в поток по ссылке
    super.postMessage(buffer, [buffer]);
    // id = (id + 1) % 0x10000000
    this.#p32[0]++;
    this.#p32[0] &= 0x0FFFFFFF;
  }
}

const taskSize = 1 << 16;

if (isMainThread) {
  const {randomBytes} = require('node:crypto');
  const fs = require('node:fs');
  const {tmpdir} = require('node:os');
  const {sep} = require('node:path');

  console.log(Date.now(), 'Main : online');
  const messages = Array(1 << 12).fill().map(_ => randomBytes(taskSize));

  // создаем временную папку и в ней файлы со всеми "сообщениями"
  const dir = fs.mkdtempSync(tmpdir() + sep);
  messages.forEach((data, i) => {
    const fn = i.toString(16).padStart(3, '0');
    fs.writeFileSync(dir + sep + fn, data);
  });
  console.log(Date.now(), 'Main : generated');

  const hashes = messages.map(() => undefined);
  let remain;

  const coordinator = new Coordinator(
    __filename
  , {
      dataField : message => message.data // двоичные данные лежат в одном поле
    }
  );
  coordinator
    .on('online', () => {
      console.log(Date.now(), 'Coordinator : online');
      // получаем список всех сообщений
      const fns = fs.readdirSync(dir)
        .sort()
        .map(fn => dir + sep + fn);
      remain = fns.length;

      fns.forEach((fn, id) => {
        const data = fs.readFileSync(fn); // тяжелый синхронный код
        coordinator.postMessage({id, fn, data});
      });
      console.log(Date.now(), 'Main : all send');
    })
    .on('port.open', port => {
      console.log(Date.now(), `Coordinator : dataPort open  = +1 worker [${port.threadId}]`);
    })
    .on('port.close', port => {
      console.log(Date.now(), `Coordinator : dataPort close = -1 worker [${port.threadId}]`);
    })
    .on('message', (result, message) => {
      hashes[message.id] = result.hash;
      if (!--remain) {
        console.log(Date.now(), 'Main : all recv');
        // подождем, пока все завершатся все желающие потоки
        setTimeout(() => {
          process.exit();
        }, 1000);
      }
    });
}
else {
  const {workerType, port} = workerData;
  switch (workerType) {
    case 'coordinator':
      // в потоке-координаторе нет активности, кроме управления пулом рабочих потоков
      const pool = new WorkersPool(
        __filename
      , {
          poolSize      : require('node:os').cpus().length // по количеству CPU-ядер
        , dataSize      : taskSize // предельный размер данных
        , port
        , queuePowMin   : 8  //   256
        , queuePowMax   : 16 // 65536
        , intervalCheck : 10
        , timeoutSpawn  : 100
        , timeoutIdle   : 10
        }
      );
      break;
    case 'worker':
      const {createHash} = require('node:crypto');

      const {shared} = workerData;

      // [shared] = {lock:int32} + {size:uint32} + {data:uint8[]}
      const lock = new Int32Array(shared, 0, 1);
      const size = new Uint32Array(shared, Int32Array.BYTES_PER_ELEMENT, 1);
      const data = new Uint8Array(shared, Int32Array.BYTES_PER_ELEMENT + Uint32Array.BYTES_PER_ELEMENT);

      const processMessage = () => {
        // обрабатываем сообщение
        port.postMessage({
          id   : lock[0]              // связующий ID
        , hash : createHash('sha256').update(
            data.subarray(0, size[0]) // входящий контент
          ).digest('hex')
        });
        // уведомляем координатор о своей доступности
        lock[0] = THREAD_FREE;
        Atomics.notify(lock, 0, 1);
        // ... и возвращаемся к ожиданию блокировки
        wait();
      };

      const wait = () => {
        const {value} = Atomics.waitAsync(lock, 0, THREAD_FREE);
        if (value === 'not-equal') {
          // если значение изменилось, то поток уже обработал задачу, и реагируем сразу
          processMessage();
        }
        else {
          // иначе ждем разрешения Promise блокировки
          value.then(processMessage);
        }
      };

      // "освобождаем" блокировку и уведомляем координатора
      lock[0] = THREAD_FREE;
      Atomics.notify(lock, 0, 1);
      wait();

      // подвешиваем поток в бесконечное ожидание
      port.on('message', () => {});
      break;
  }
}