You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by jl...@apache.org on 2026/06/09 17:30:15 UTC
(superset) branch master updated: feat(websocket): add configurable per-channel and total connection limits (#40856)
This is an automated email from the ASF dual-hosted git repository.
sadpandajoe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 8117488fd82 feat(websocket): add configurable per-channel and total connection limits (#40856)
8117488fd82 is described below
commit 8117488fd82ffbad9c99995d4f752746799a51fd
Author: Evan Rusackas <ev...@preset.io>
AuthorDate: Tue Jun 9 10:30:07 2026 -0700
feat(websocket): add configurable per-channel and total connection limits (#40856)
Co-authored-by: Claude Code <no...@anthropic.com>
---
superset-websocket/spec/index.test.ts | 166 ++++++++++++++++++++++++++++++++++
superset-websocket/src/config.ts | 8 ++
superset-websocket/src/index.ts | 72 +++++++++++++++
3 files changed, 246 insertions(+)
diff --git a/superset-websocket/spec/index.test.ts b/superset-websocket/spec/index.test.ts
index 6a535b24cf7..8a2ded35526 100644
--- a/superset-websocket/spec/index.test.ts
+++ b/superset-websocket/spec/index.test.ts
@@ -632,6 +632,172 @@ describe('server', () => {
});
});
+ describe('connection limits', () => {
+ const getRequest = (token: string, url: string): http.IncomingMessage => {
+ const request = new http.IncomingMessage(new net.Socket());
+ request.method = 'GET';
+ request.headers = { cookie: `${config.jwtCookieName}=${token}` };
+ request.url = url;
+ return request;
+ };
+
+ afterEach(() => {
+ // restore opt-in limits to their disabled default
+ server.opts.maxTotalConnections = 0;
+ server.opts.maxConnectionsPerChannel = 0;
+ });
+
+ test('no limit when disabled (0)', () => {
+ server.opts.maxTotalConnections = 0;
+ server.opts.maxConnectionsPerChannel = 0;
+ const socketInstance = {
+ ws: new wsMock('localhost'),
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, socketInstance);
+ expect(server.connectionLimitReason(channelId)).toBeNull();
+ });
+
+ test('total connection limit reached', () => {
+ server.opts.maxTotalConnections = 1;
+ const ws = new wsMock('localhost');
+ setReadyState(ws, WebSocket.OPEN);
+ const socketInstance = {
+ ws,
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, socketInstance);
+ expect(server.connectionLimitReason('some-other-channel')).toMatch(
+ /total connection limit/,
+ );
+ });
+
+ test('per-channel connection limit reached', () => {
+ server.opts.maxConnectionsPerChannel = 1;
+ const ws = new wsMock('localhost');
+ setReadyState(ws, WebSocket.OPEN);
+ const socketInstance = {
+ ws,
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, socketInstance);
+ expect(server.connectionLimitReason(channelId)).toMatch(
+ /per-channel connection limit/,
+ );
+ });
+
+ test('stale closed socket does not count toward total limit', () => {
+ server.opts.maxTotalConnections = 1;
+ const ws = new wsMock('localhost');
+ const socketInstance = {
+ ws,
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, socketInstance);
+ // simulate the socket having closed but not yet been GC'd
+ setReadyState(ws, WebSocket.CLOSED);
+ expect(server.connectionLimitReason('some-other-channel')).toBeNull();
+ });
+
+ test('stale closed socket does not count toward per-channel limit', () => {
+ server.opts.maxConnectionsPerChannel = 1;
+ const ws = new wsMock('localhost');
+ const socketInstance = {
+ ws,
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, socketInstance);
+ // simulate the socket having closed but not yet been GC'd
+ setReadyState(ws, WebSocket.CLOSED);
+ expect(server.connectionLimitReason(channelId)).toBeNull();
+ });
+
+ test('isSocketActive reflects the socket readyState', () => {
+ const ws = new wsMock('localhost');
+ setReadyState(ws, WebSocket.OPEN);
+ const socketId = server.trackClient(channelId, {
+ ws,
+ channel: channelId,
+ pongTs: Date.now(),
+ });
+ expect(server.isSocketActive(socketId)).toBe(true);
+ // CONNECTING is also considered active (see SOCKET_ACTIVE_STATES)
+ setReadyState(ws, WebSocket.CONNECTING);
+ expect(server.isSocketActive(socketId)).toBe(true);
+ setReadyState(ws, WebSocket.CLOSED);
+ expect(server.isSocketActive(socketId)).toBe(false);
+ // unknown socket ids are never active
+ expect(server.isSocketActive('does-not-exist')).toBe(false);
+ });
+
+ test('activeSocketCount counts only active sockets', () => {
+ const openWs = new wsMock('localhost');
+ setReadyState(openWs, WebSocket.OPEN);
+ server.trackClient(channelId, {
+ ws: openWs,
+ channel: channelId,
+ pongTs: Date.now(),
+ });
+ const closedWs = new wsMock('localhost');
+ setReadyState(closedWs, WebSocket.CLOSED);
+ server.trackClient(channelId, {
+ ws: closedWs,
+ channel: channelId,
+ pongTs: Date.now(),
+ });
+ expect(server.activeSocketCount()).toBe(1);
+ });
+
+ test('activeChannelSocketCount counts only active sockets on the channel', () => {
+ const openWs = new wsMock('localhost');
+ setReadyState(openWs, WebSocket.OPEN);
+ server.trackClient(channelId, {
+ ws: openWs,
+ channel: channelId,
+ pongTs: Date.now(),
+ });
+ const closedWs = new wsMock('localhost');
+ setReadyState(closedWs, WebSocket.CLOSED);
+ server.trackClient(channelId, {
+ ws: closedWs,
+ channel: channelId,
+ pongTs: Date.now(),
+ });
+ expect(server.activeChannelSocketCount(channelId)).toBe(1);
+ // unknown channels report zero active sockets
+ expect(server.activeChannelSocketCount('no-such-channel')).toBe(0);
+ });
+
+ test('wsConnection refuses over-limit connection without tracking', () => {
+ server.opts.maxConnectionsPerChannel = 1;
+ const existingWs = new wsMock('localhost');
+ setReadyState(existingWs, WebSocket.OPEN);
+ const existing = {
+ ws: existingWs,
+ channel: channelId,
+ pongTs: Date.now(),
+ };
+ server.trackClient(channelId, existing);
+
+ const trackClientSpy = jest.spyOn(server, 'trackClient');
+ const ws = new wsMock('localhost');
+ const validToken = jwt.sign({ channel: channelId }, config.jwtSecret);
+ server.wsConnection(ws, getRequest(validToken, 'http://localhost'));
+
+ expect(ws.close).toHaveBeenCalledWith(
+ 1013,
+ expect.stringMatching(/limit/),
+ );
+ expect(trackClientSpy).not.toHaveBeenCalled();
+ trackClientSpy.mockRestore();
+ });
+ });
+
describe('httpUpgrade', () => {
let socket: net.Socket;
let socketDestroySpy: jest.SpiedFunction<typeof socket.destroy>;
diff --git a/superset-websocket/src/config.ts b/superset-websocket/src/config.ts
index 3fd9d85a773..72f7a40fcb8 100644
--- a/superset-websocket/src/config.ts
+++ b/superset-websocket/src/config.ts
@@ -53,6 +53,8 @@ type ConfigType = {
gcChannelsIntervalMs: number;
maxSocketBufferBytes: number;
eventYieldBatchSize: number;
+ maxConnectionsPerChannel: number;
+ maxTotalConnections: number;
};
function defaultConfig(): ConfigType {
@@ -78,6 +80,9 @@ function defaultConfig(): ConfigType {
// Number of stream events to process before yielding to the event loop.
// 0 disables yielding (process the whole batch synchronously).
eventYieldBatchSize: 100,
+ // 0 disables the limit (unlimited); set a positive value to opt in.
+ maxConnectionsPerChannel: 0,
+ maxTotalConnections: 0,
statsd: {
host: '127.0.0.1',
port: 8125,
@@ -160,6 +165,9 @@ function applyEnvOverrides(config: ConfigType): ConfigType {
val,
config.eventYieldBatchSize,
)),
+ MAX_CONNECTIONS_PER_CHANNEL: val =>
+ (config.maxConnectionsPerChannel = toNumber(val)),
+ MAX_TOTAL_CONNECTIONS: val => (config.maxTotalConnections = toNumber(val)),
REDIS_HOST: val => (config.redis.host = val),
REDIS_PORT: val => (config.redis.port = toNumber(val)),
REDIS_PASSWORD: val => (config.redis.password = val),
diff --git a/superset-websocket/src/index.ts b/superset-websocket/src/index.ts
index 275aaeae5dc..5bafe37950b 100644
--- a/superset-websocket/src/index.ts
+++ b/superset-websocket/src/index.ts
@@ -163,6 +163,67 @@ export const setLastFirehoseId = (id: string): void => {
lastFirehoseId = id;
};
+// WebSocket close code used when a connection is refused because a configured
+// connection limit has been reached (1013 = "Try Again Later").
+const CONNECTION_LIMIT_CLOSE_CODE = 1013;
+
+/**
+ * Returns whether the socket with the given id is currently active, i.e. it is
+ * still registered and its underlying connection is in an active readyState.
+ *
+ * Closed sockets are only removed from the registries asynchronously (via the
+ * `checkSockets`/`cleanChannel` GC routines), so connection-limit checks must
+ * filter on live socket state rather than trusting the raw registry sizes.
+ */
+export const isSocketActive = (socketId: string): boolean => {
+ const socketInstance = sockets[socketId];
+ return (
+ !!socketInstance &&
+ SOCKET_ACTIVE_STATES.includes(socketInstance.ws.readyState)
+ );
+};
+
+/**
+ * Counts the sockets in the global registry that are still active.
+ */
+export const activeSocketCount = (): number =>
+ Object.keys(sockets).filter(isSocketActive).length;
+
+/**
+ * Counts the active sockets currently registered on the given channel.
+ */
+export const activeChannelSocketCount = (channel: string): number =>
+ channels[channel]?.sockets.filter(isSocketActive).length ?? 0;
+
+/**
+ * Determines whether accepting a new connection on the given channel would
+ * exceed a configured connection limit. Returns a human-readable reason when a
+ * limit is reached, or `null` when the connection is within limits.
+ *
+ * Both limits are opt-in: a value of `0` (the default) disables the check.
+ *
+ * Counts are derived from active socket state rather than raw registry sizes:
+ * recently closed sockets linger in the registries until the next GC pass, so
+ * counting them would spuriously reject new connections even when no active
+ * connection is consuming capacity.
+ */
+export const connectionLimitReason = (channel: string): string | null => {
+ const { maxTotalConnections, maxConnectionsPerChannel } = opts;
+
+ if (maxTotalConnections > 0 && activeSocketCount() >= maxTotalConnections) {
+ return `total connection limit (${maxTotalConnections}) reached`;
+ }
+
+ if (
+ maxConnectionsPerChannel > 0 &&
+ activeChannelSocketCount(channel) >= maxConnectionsPerChannel
+ ) {
+ return `per-channel connection limit (${maxConnectionsPerChannel}) reached`;
+ }
+
+ return null;
+};
+
/**
* Adds the passed channel and socket instance to the internal registries.
*/
@@ -379,6 +440,17 @@ export const incrementId = (id: string): string => {
*/
export const wsConnection = (ws: WebSocket, request: http.IncomingMessage) => {
const channel: string = readChannelId(request);
+
+ // Refuse the connection if a configured connection limit has been reached,
+ // before tracking it against the internal registries.
+ const limitReason = connectionLimitReason(channel);
+ if (limitReason) {
+ statsd.increment('ws_connection_rejected');
+ logger.warn(`Refusing connection on channel ${channel}: ${limitReason}`);
+ ws.close(CONNECTION_LIMIT_CLOSE_CODE, limitReason);
+ return;
+ }
+
const socketInstance: SocketInstance = { ws, channel, pongTs: Date.now() };
// add this ws instance to the internal registry