Running many WebSocket servers

#websocket #koa #backend

At Woovi, we have 7 frontends, and each frontend has their own server.
To provide a better DX, we run all the servers inside only one server in development.
Our solution to handle many WebSocket servers in the same server is to have 1 WebSocket server per endpoint.

WebSocket Server per endpoint

import http from 'http';

import WebSocket, { WebSocketServer as WSWebSocketServer } from 'ws';

// work with commonjs and esm
const WebSocketServer = WebSocket.Server || WSWebSocketServer;

export const createWebsocketMiddleware = (
  propertyName = 'ws',
  options = {},
) => {
  if (options instanceof http.Server) options = { server: options };

  const wsServers = {};

  const getOrCreateWebsocketServer = (url: string) => {
    const server = wsServers[url];

    if (server) {
      return server;
    }

    const newServer = new WebSocketServer({
      ...(options.wsOptions || {}),
      noServer: true,
    });

    wsServers[url] = newServer;

    return newServer;
  };

  const websocketMiddleware = async (ctx, next) => {
    const upgradeHeader = (ctx.request.headers.upgrade || '')
      .split(',')
      .map((s) => s.trim());

    if (~upgradeHeader.indexOf('websocket')) {
      const wss = getOrCreateWebsocketServer(ctx.url);

      ctx[propertyName] = () =>
        new Promise((resolve) => {
          wss.handleUpgrade(
            ctx.req,
            ctx.request.socket,
            Buffer.alloc(0),
            (ws) => {
              wss.emit('connection', ws, ctx.req);
              resolve(ws);
            },
          );
          ctx.respond = false;
        });
      ctx.wss = wss;
    }

    await next();
  };

  return websocketMiddleware;
};

Usage

const app = new Koa();

app.use(createWebsocketMiddleware());

const routes = new Router();

routes.all('/ws', async (ctx) => {
   if (ctx.wss) { // check if this is a websocket connection
      // connect to websocket - upgrade
      const client = await ctx.ws(); 
   }
});

Let's deep dive:

getOrCreateWebsocketServer get an existing WebSocket Server or create a new one based on URL
upgradeHeader checks if the request wants to use WebSocket
wss.handleUpgrade upgrades the requests to use WebSocket

To sum up

This is a generic approach, that let you have many WebSocket Server in a single server per endpoint.
It can be used in localhost to improve DX or even in production if you don't have much traffic.

Woovi
Woovi is a Startup that enables shoppers to pay as they like. To make this possible, Woovi provides instant payment solutions for merchants to accept orders.

If you want to work with us, we are hiring!

Photo by Tanbir Mahmud on Unsplash