From 694955f864375515d84d463e9e7ee89a156a4a2b Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Thu, 27 Aug 2026 18:54:34 +0200 Subject: [PATCH 1/2] fix(api-gateway): Make pre-agg build job status per-entry and data-source aware (#11666) --- packages/cubejs-api-gateway/src/gateway.ts | 23 ++-- .../cubejs-api-gateway/test/index.test.ts | 103 +++++++++++++++++- .../src/orchestrator/PreAggregations.ts | 7 ++ .../src/orchestrator/QueryOrchestrator.ts | 2 +- .../test/unit/PreAggregations.test.ts | 28 +++++ .../test/unit/QueryOrchestrator.jobs.test.ts | 54 +++++++++ .../src/compiler/CubeEvaluator.ts | 2 + .../src/core/RefreshScheduler.ts | 7 +- .../test/unit/RefreshScheduler.test.ts | 8 ++ 9 files changed, 217 insertions(+), 17 deletions(-) create mode 100644 packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.jobs.test.ts diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts index 1b5389d9e5c86..1dad077cf277a 100644 --- a/packages/cubejs-api-gateway/src/gateway.ts +++ b/packages/cubejs-api-gateway/src/gateway.ts @@ -1133,8 +1133,6 @@ class ApiGateway { .refreshScheduler() .getCachedBuildJobs(context, tokens); - const metaCache: Map = new Map(); - const response: PreAggJobStatusItem[] = await Promise.all( jobs.map(async ({ job, token }) => { if (!job) { @@ -1147,12 +1145,15 @@ class ApiGateway { const ctx = { ...context, ...job.context }; const orchestrator = await this.getAdapterApi(ctx); const compiler = await this.getCompilerApi(ctx); + // TODO(1.8): drop the fallback, no job posted by 1.7 can still be in the cache. + const dataSource = job.dataSource || (await compiler.preAggregations()) + .find(pa => pa.id === job.preagg)?.dataSource; const selector: PreAggsSelector = { cubes: [job.preagg.split('.')[0]], preAggregations: [job.preagg], contexts: [job.context], timezones: [job.timezone], - dataSources: [job.dataSource], + dataSources: [dataSource], }; if ( job.status.indexOf('done') === 0 || @@ -1170,6 +1171,7 @@ class ApiGateway { const status = await this.getPreAggJobQueueStatus( orchestrator, job, + dataSource, ); if (status) { // returning queued status @@ -1180,18 +1182,13 @@ class ApiGateway { selector, }; } else { - const metaCacheKey = JSON.stringify(ctx); - if (!metaCache.has(metaCacheKey)) { - metaCache.set(metaCacheKey, await compiler.metaConfigExtended(context, ctx)); - } - // checking and fetching result status const s = await this.getPreAggJobResultStatus( ctx.requestId, orchestrator, compiler, - metaCache.get(metaCacheKey), job, + dataSource, token, ); @@ -1225,10 +1222,11 @@ class ApiGateway { private async getPreAggJobQueueStatus( orchestrator: any, job: PreAggJob, + dataSource?: string, ): Promise { let inQueue = false; let status: string = 'n/a'; - const queuedList = await orchestrator.getPreAggregationQueueStates(job.dataSource); + const queuedList = await orchestrator.getPreAggregationQueueStates(dataSource); queuedList.forEach((item) => { if ( item.queryHandler && @@ -1264,19 +1262,18 @@ class ApiGateway { requestId: string, orchestrator: any, compiler: any, - metadata: any, job: PreAggJob, + dataSource: string | undefined, token: string, ): Promise { const preaggs = await compiler.preAggregations(); const preagg = preaggs.find(pa => pa.id === job.preagg); if (preagg) { - const cube = metadata.cubeDefinitions[preagg.cube]; const [, status]: [boolean, string] = await orchestrator.isPartitionExist( requestId, preagg.preAggregation.external, - cube.dataSource, + dataSource, compiler.preAggregationsSchema, job.target, job.key, diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts index 87d0825e0bdf7..07d3a17c003ae 100644 --- a/packages/cubejs-api-gateway/test/index.test.ts +++ b/packages/cubejs-api-gateway/test/index.test.ts @@ -1231,11 +1231,112 @@ describe('API Gateway', () => { }), }; - const status = await (apiGateway as any).getPreAggJobQueueStatus(orchestrator, job); + const status = await (apiGateway as any).getPreAggJobQueueStatus(orchestrator, job, job.dataSource); expect(orchestrator.getPreAggregationQueueStates).toHaveBeenCalledWith(job.dataSource); expect(status).toEqual('processing'); }); + + // https://github.com/cube-js/cube/issues/11615 + describe('job result status', () => { + const jobFor = (overrides: Partial = {}): PreAggJob => ({ + request: 'request-id', + context: { securityContext: {} }, + preagg: 'orders_test.main', + table: 'orders_test_main', + target: 'orders_test_main_20200101', + structure: 'structure-version', + content: 'content-version', + updated: 1, + key: [], + status: 'posted', + timezone: 'UTC', + dataSource: 'test_ds', + ...overrides, + }); + + const compiler = { + preAggregations: async () => ([ + { id: 'orders_dep.main', cube: 'orders_dep', dataSource: 'dep_ds', preAggregation: { external: false } }, + { id: 'orders_test.main', cube: 'orders_test', dataSource: 'model_ds', preAggregation: { external: true } }, + ]), + preAggregationsSchema: 'stb_pre_aggregations', + }; + + const resultStatus = async (job: PreAggJob, orchestrator: any) => { + const apiGateway = Object.create(ApiGateway.prototype); + return (apiGateway as any).getPreAggJobResultStatus( + job.request, + orchestrator, + compiler, + job, + job.dataSource, + 'job-token', + ); + }; + + test('is checked on the data source resolved for the job', async () => { + const orchestrator = { isPartitionExist: jest.fn(async () => [true, 'done']) }; + const job = jobFor(); + + await expect(resultStatus(job, orchestrator)).resolves.toEqual('done'); + + expect(orchestrator.isPartitionExist).toHaveBeenCalledWith( + job.request, + true, + 'test_ds', + compiler.preAggregationsSchema, + job.target, + job.key, + 'job-token', + ); + }); + + test('reports a pre-aggregation the model no longer has', async () => { + const orchestrator = { isPartitionExist: jest.fn() }; + + await expect( + resultStatus(jobFor({ preagg: 'orders_test.dropped' }), orchestrator) + ).resolves.toEqual('pre_agg_not_found'); + + expect(orchestrator.isPartitionExist).not.toHaveBeenCalled(); + }); + + // Everywhere, because a queue lookup left on the default data source finds nothing + // and a build that is still scheduled then reads as a missing partition. + // TODO(1.8): goes away with the fallback in preAggregationsJobsGET. + test('a job with no recorded data source resolves it from the model everywhere', async () => { + const apiGateway = Object.create(ApiGateway.prototype); + const job = jobFor({ dataSource: undefined as any, status: 'scheduled' }); + const orchestrator = { + getPreAggregationQueueStates: jest.fn(async () => []), + isPartitionExist: jest.fn(async () => [false, 'missing_partition']), + }; + apiGateway.refreshScheduler = () => ({ + getCachedBuildJobs: async () => [{ job, token: 'job-token' }], + }); + apiGateway.getAdapterApi = async () => orchestrator; + apiGateway.getCompilerApi = async () => compiler; + + const [item] = await (apiGateway as any).preAggregationsJobsGET( + { requestId: 'request-id' }, + ['job-token'], + ); + + expect(item.status).toEqual('missing_partition'); + expect(item.selector.dataSources).toEqual(['model_ds']); + expect(orchestrator.getPreAggregationQueueStates).toHaveBeenCalledWith('model_ds'); + expect(orchestrator.isPartitionExist).toHaveBeenCalledWith( + 'request-id', + true, + 'model_ds', + compiler.preAggregationsSchema, + job.target, + job.key, + 'job-token', + ); + }); + }); }); describe('healtchecks', () => { diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts index 8c2e824bfb77f..2ad4087fa81ff 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts @@ -150,6 +150,10 @@ export type LoadPreAggregationResult = { partitionRange?: QueryDateRange; isMultiTableUnion?: boolean; usageTargetTableNames?: Record; + type?: 'rollup' | 'originalSql'; + preAggregationId?: string; + dataSource?: string; + timezone?: string; }; export type PreAggregationTableToTempTable = [string, LoadPreAggregationResult]; @@ -574,6 +578,9 @@ export class PreAggregations { const usedPreAggregation = { ...loadResult, type: p.type, + preAggregationId: p.preAggregationId, + dataSource: p.dataSource || 'default', + timezone: p.timezone, }; if (!usedPreAggregation.isMultiTableUnion) { await this.addTableUsed(usedPreAggregation.targetTableName); diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts index c7cfaf8827542..de96a3c7ccbca 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts @@ -252,7 +252,7 @@ export class QueryOrchestrator { // /cubejs-system/v1/pre-aggregations/jobs endpoint). if (queryBody.isJob) { return preAggregationsTablesToTempTables.map((pa) => ({ - preAggregation: queryBody.preAggregations[0].preAggregationId, + preAggregation: pa[1].preAggregationId || queryBody.preAggregations[0].preAggregationId, tableName: pa[0], ...pa[1], })); diff --git a/packages/cubejs-query-orchestrator/test/unit/PreAggregations.test.ts b/packages/cubejs-query-orchestrator/test/unit/PreAggregations.test.ts index 17fd28ffa6e17..9a2da35351881 100644 --- a/packages/cubejs-query-orchestrator/test/unit/PreAggregations.test.ts +++ b/packages/cubejs-query-orchestrator/test/unit/PreAggregations.test.ts @@ -337,6 +337,34 @@ describe('PreAggregations', () => { expect(result[0][1].targetTableName).toMatch(/stb_pre_aggregations.orders_number_and_count20191101_kjypcoio_5yftl5il/); expect(result[0][1].lastUpdatedAt).toEqual(12345000); }); + + // A jobed build gets back a flat list of entries and has to tell them apart. + // https://github.com/cube-js/cube/issues/11615 + test('each entry carries the identity of the descriptor it was built from', async () => { + const { preAggregationsTablesToTempTables: result } = await preAggregations!.loadAllPreAggregationsIfNeeded( + createBasicQuery({ + cacheMode: 'must-revalidate', + preAggregations: [{ + ...basicQuery.preAggregations[0], + preAggregationId: 'Orders.numberAndCount', + dataSource: 'orders_ds', + timezone: 'America/Los_Angeles', + }], + }) + ); + + expect(result[0][1]).toMatchObject({ + preAggregationId: 'Orders.numberAndCount', + dataSource: 'orders_ds', + timezone: 'America/Los_Angeles', + }); + }); + + test('an entry built without a named data source falls back to the default one', async () => { + const { preAggregationsTablesToTempTables: result } = await preAggregations!.loadAllPreAggregationsIfNeeded(basicQueryWithRenew); + + expect(result[0][1].dataSource).toEqual('default'); + }); }); describe('loadAllPreAggregationsIfNeeded with external rollup and writable source', () => { diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.jobs.test.ts b/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.jobs.test.ts new file mode 100644 index 0000000000000..be6578cd924b7 --- /dev/null +++ b/packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.jobs.test.ts @@ -0,0 +1,54 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { QueryOrchestrator } from '../../src'; + +// A jobed build returns one entry per built pre-aggregation — the requested partition plus +// every dependency it had to build first — and each entry becomes its own polling token. +// https://github.com/cube-js/cube/issues/11615 +describe('QueryOrchestrator jobed build', () => { + const entry = (tableName: string, overrides: Record = {}) => ([ + tableName, + { + targetTableName: `${tableName}_kjypcoio_5yftl5il_1593709044209`, + refreshKeyValues: [], + lastUpdatedAt: 1593709044209, + ...overrides, + }, + ]); + + const fetchJob = async (preAggregationsTablesToTempTables: any[], preAggregations: any[]) => { + const orchestrator = Object.create(QueryOrchestrator.prototype); + orchestrator.rollupOnlyMode = false; + orchestrator.preAggregations = { + loadAllPreAggregationsIfNeeded: async () => ({ + preAggregationsTablesToTempTables, + values: null, + }), + }; + + return orchestrator.fetchQuery({ isJob: true, preAggregations }); + }; + + test('labels every entry with its own pre-aggregation and data source', async () => { + const job = await fetchJob( + [ + entry('stb_pre_aggregations.orders_main', { preAggregationId: 'Orders.main', dataSource: 'orders_ds' }), + entry('stb_pre_aggregations.orders_rollup', { preAggregationId: 'Orders.rollup', dataSource: 'default', timezone: 'UTC' }), + ], + [{ preAggregationId: 'Orders.main' }, { preAggregationId: 'Orders.rollup' }], + ); + + expect(job).toMatchObject([ + { preAggregation: 'Orders.main', tableName: 'stb_pre_aggregations.orders_main', dataSource: 'orders_ds' }, + { preAggregation: 'Orders.rollup', tableName: 'stb_pre_aggregations.orders_rollup', dataSource: 'default', timezone: 'UTC' }, + ]); + }); + + test('falls back to the requested pre-aggregation for an entry without an id', async () => { + const job = await fetchJob( + [entry('stb_pre_aggregations.orders_rollup')], + [{ preAggregationId: 'Orders.rollup' }], + ); + + expect(job[0].preAggregation).toEqual('Orders.rollup'); + }); +}); diff --git a/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts b/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts index 2909d43bf8c2e..3beac37469f09 100644 --- a/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts +++ b/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts @@ -160,6 +160,7 @@ export type PreAggregationInfo = { preAggregationName: string, preAggregation: any, cube: string, + dataSource: string, references: PreAggregationReferences, refreshKey: unknown, indexesReferences: unknown, @@ -963,6 +964,7 @@ export class CubeEvaluator extends CubeSymbols { preAggregationName, preAggregation: preAggregations[preAggregationName], cube, + dataSource: this.evaluatedCubes[cube].dataSource || 'default', references: this.evaluatePreAggregationReferences(cube, preAggregations[preAggregationName]), refreshKey, indexesReferences: indexes && Object.keys(indexes).reduce((obj, indexName) => { diff --git a/packages/cubejs-server-core/src/core/RefreshScheduler.ts b/packages/cubejs-server-core/src/core/RefreshScheduler.ts index 67a10cbe26f8a..5ab684bbb83f5 100644 --- a/packages/cubejs-server-core/src/core/RefreshScheduler.ts +++ b/packages/cubejs-server-core/src/core/RefreshScheduler.ts @@ -790,8 +790,11 @@ export class RefreshScheduler { metadata: queryingOptions.metadata, isJob: true, }); - job[0].dataSource = partition.dataSource; - job[0].timezone = partition.timezone; + job.forEach((j: JobedPreAggregation) => { + j.dataSource = j.dataSource || partition.dataSource; + j.timezone = j.timezone || partition.timezone; + }); + return job; } ) diff --git a/packages/cubejs-server-core/test/unit/RefreshScheduler.test.ts b/packages/cubejs-server-core/test/unit/RefreshScheduler.test.ts index 8a274e0256f63..85ccd19175ced 100644 --- a/packages/cubejs-server-core/test/unit/RefreshScheduler.test.ts +++ b/packages/cubejs-server-core/test/unit/RefreshScheduler.test.ts @@ -683,6 +683,7 @@ describe('Refresh Scheduler', () => { external: false, }, cube: 'Foo', + dataSource: 'default', references: { dimensions: [], measures: ['Foo.count'], @@ -902,6 +903,13 @@ describe('Refresh Scheduler', () => { const buildJobs = await refreshScheduler.getCachedBuildJobs(ctx, jobs); const allTokensExist = jobs.every(token => buildJobs.some(job => job.token === token)); expect(allTokensExist).toBeTruthy(); + + // Not only the first entry: every entry of a posted job is its own poll token. + // https://github.com/cube-js/cube/issues/11615 + buildJobs.forEach(({ job }) => { + expect(job?.dataSource).toEqual('default'); + expect(['UTC', 'America/Los_Angeles']).toContain(job?.timezone); + }); }); test('Only `first` pre-aggregation', async () => { From 111eb3ec049b9be1fa2f31186a7f83f58b5a0291 Mon Sep 17 00:00:00 2001 From: waralexrom <108349432+waralexrom@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:17:22 +0200 Subject: [PATCH 2/2] feat(cubestore): cap concurrent websocket connections per user (#11667) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(cubestore): cap concurrent websocket connections per user Adds CUBESTORE_MAX_WS_CONNECTIONS_PER_USER, an optional per-node cap on how many concurrent WebSocket connections one authenticated user may hold. A client that opens connections faster than it closes them can otherwise consume the whole process file descriptor table, which affects every other user of that node. At the limit the user's oldest connection is closed to admit the new one, rather than the new one being refused: a client that still needs the closed connection reconnects, while one that had forgotten about it simply loses it. It is closed with a close frame, so the client sees a clean close and can reconnect on its own terms. A slot is taken at the websocket upgrade and released by a guard held for the lifetime of the connection task, so it comes back on any exit path: a close frame, a transport error, a panic, a cancelled task, or a connection that dies before the upgrade completes. Live connections are keyed by a monotonic id, so an evicting admission and the victim's own later drop cannot remove each other's entry, and the limit never overshoots. The limit defaults to 0, which disables the cap and leaves behaviour unchanged. * fix(cubestore): bound the evicted connection's close and name its reason Sending the close frame flushes, so a peer whose receive window is closed would park the connection task indefinitely with the descriptor still held. That matters more than a stray task: the entry leaves the counter when the connection is evicted rather than when its task ends, so the cap would read as satisfied while the process held more connections than it allows. The send is now bounded, and the bound is deliberately short — the steady-state overshoot is the eviction rate times this timeout, and a close frame is a handful of bytes, so a peer that cannot take them in two seconds is not going to. The frame now carries close code 1013 ("Try Again Later") and a reason, so a client can tell being recycled from a transport failure and reconnect rather than report an error. Adds an end-to-end test: with a cap of one, a second connection for the same authenticated user evicts the first, which receives that close code and reason, while a different user connects unaffected. It covers the parts the unit tests cannot reach — the select branch, the close path and the guard living in the connection task. The read is bounded, so a regression in the eviction path fails the test rather than hanging it. * refactor(cubestore): bound every take of the connection counter's lock The counter's critical section is a map lookup and an insert with nothing awaited inside it — `acquire` is not an `async fn`, so nothing can be — which is the shape the `tokio::sync::Mutex` docs point at a standard-library lock for. That makes a stall unreachable by reasoning about the callers; taking the lock with a bound makes it unreachable without needing that reasoning at all. Switches to parking_lot for `try_lock_for`. It also removes poisoning, which had to be recovered by hand before: there is nothing meaningful to recover for a map of live connections, and a poisoned unwrap in the destructor would have aborted the process. parking_lot is already in the tree at this version, so the lockfile gains one dependency edge and no crate. Both callers fail open, because the cap is a safety net rather than a correctness invariant: an unavailable counter admits the connection untracked instead of refusing or stalling it, and a finished connection leaves its entry behind to be evicted — that entry is the oldest one of its user by construction, and cancelling a finished connection's token is a no-op. Both paths log, so a bound that is ever reached is visible instead of quietly weakening the cap. The bound is 100ms rather than something tighter because a thread holding the lock can be descheduled for milliseconds under CPU pressure, and giving up then would stop counting connections exactly when the cap matters most. * refactor(cubestore): give the connection counter's lock a second to be taken The bound exists to make a stall impossible without reasoning about the callers, not to be tight, and the failure it guards against is one-sided: a thread holding the lock can be descheduled for a long time under CPU pressure, and giving up then would quietly stop counting connections exactly when the cap matters most. A second is far past anything a map lookup and an insert can take, and giving up is logged either way. --- .../configuration/environment-variables.mdx | 17 + rust/cubestore/Cargo.lock | 1 + rust/cubestore/cubestore/Cargo.toml | 2 + rust/cubestore/cubestore/src/config/mod.rs | 15 + rust/cubestore/cubestore/src/http/mod.rs | 427 +++++++++++++++++- 5 files changed, 460 insertions(+), 2 deletions(-) diff --git a/docs-mintlify/reference/configuration/environment-variables.mdx b/docs-mintlify/reference/configuration/environment-variables.mdx index 9dd1c9fd20256..ed782d1eb9fc5 100644 --- a/docs-mintlify/reference/configuration/environment-variables.mdx +++ b/docs-mintlify/reference/configuration/environment-variables.mdx @@ -1975,6 +1975,23 @@ The logging level for Cube Store. See also [`CUBEJS_LOG_LEVEL`](/reference/configuration/environment-variables#cubejs_log_level). +## `CUBESTORE_MAX_WS_CONNECTIONS_PER_USER` + +The maximum number of concurrent WebSocket connections a single authenticated +user may hold. At the limit, that user's oldest connection is closed to admit the +new one: a client that still needs it reconnects, while one that had forgotten +about it simply loses it. The limit is counted per Cube Store node, so a user +connecting to several nodes may hold up to this many on each. + +Use it to keep one client that leaks connections from exhausting the node's file +descriptors for everyone else. Set it comfortably above the number of connections +a client legitimately keeps open at once, or working connections will be recycled. +`0` disables the limit. + +| Possible Values | Default in Development | Default in Production | +| ------------------------- | ---------------------- | --------------------- | +| A non-negative integer | `0` | `0` | + ## `CUBESTORE_META_ADDR` The address/port pair for the Cube Store **router** node in the cluster. diff --git a/rust/cubestore/Cargo.lock b/rust/cubestore/Cargo.lock index 7046d5b330ca3..34d2321545297 100644 --- a/rust/cubestore/Cargo.lock +++ b/rust/cubestore/Cargo.lock @@ -1508,6 +1508,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-otlp", "opentelemetry_sdk", + "parking_lot", "parquet-format", "parse-size", "paste", diff --git a/rust/cubestore/cubestore/Cargo.toml b/rust/cubestore/cubestore/Cargo.toml index 830d005dad398..bed8b87f6f0c0 100644 --- a/rust/cubestore/cubestore/Cargo.toml +++ b/rust/cubestore/cubestore/Cargo.toml @@ -94,6 +94,8 @@ opentelemetry-otlp = { version = "0.26.0", default-features = false, features = opentelemetry-http = { version = "0.26.0", features = ["reqwest"] } lru = "0.18.2" moka = { version = "0.10.1", features = ["future"] } +# Already in the tree transitively at this version; used for a timed lock. +parking_lot = "0.12" ctor = "0.1.20" json = "0.12.4" futures-util = "0.3.17" diff --git a/rust/cubestore/cubestore/src/config/mod.rs b/rust/cubestore/cubestore/src/config/mod.rs index ffdb34b160ffa..d58ad3f79273d 100644 --- a/rust/cubestore/cubestore/src/config/mod.rs +++ b/rust/cubestore/cubestore/src/config/mod.rs @@ -633,6 +633,10 @@ pub trait ConfigObj: DIService { fn check_ws_orphaned_messages_interval_secs(&self) -> u64; + /// Concurrent websocket connections allowed per authenticated user. + /// 0 disables the limit. + fn max_ws_connections_per_user(&self) -> usize; + fn drop_ws_processing_messages_after_secs(&self) -> u64; fn drop_ws_complete_messages_after_secs(&self) -> u64; @@ -785,6 +789,7 @@ pub struct ConfigObjImpl { pub metadata_cache_time_to_idle_secs: u64, pub stream_replay_check_interval_secs: u64, pub check_ws_orphaned_messages_interval_secs: u64, + pub max_ws_connections_per_user: usize, pub drop_ws_processing_messages_after_secs: u64, pub drop_ws_complete_messages_after_secs: u64, pub skip_kafka_parsing_errors: bool, @@ -1178,6 +1183,10 @@ impl ConfigObj for ConfigObjImpl { self.check_ws_orphaned_messages_interval_secs } + fn max_ws_connections_per_user(&self) -> usize { + self.max_ws_connections_per_user + } + fn drop_ws_processing_messages_after_secs(&self) -> u64 { self.drop_ws_processing_messages_after_secs } @@ -1960,6 +1969,10 @@ impl Config { "CUBESTORE_CHECK_WS_ORPHANED_MESSAGES_INTERVAL", 30, ), + max_ws_connections_per_user: env_parse( + "CUBESTORE_MAX_WS_CONNECTIONS_PER_USER", + 0, + ), drop_ws_processing_messages_after_secs: env_parse( "CUBESTORE_DROP_WS_PROCESSING_MESSAGES_AFTER", 60 * 60, @@ -2192,6 +2205,7 @@ impl Config { gc_loop_interval: 60, stream_replay_check_interval_secs: 60, check_ws_orphaned_messages_interval_secs: 1, + max_ws_connections_per_user: 0, drop_ws_processing_messages_after_secs: 60, drop_ws_complete_messages_after_secs: 10, skip_kafka_parsing_errors: false, @@ -2974,6 +2988,7 @@ impl Config { Duration::from_secs(config.drop_ws_complete_messages_after_secs()), config.transport_max_message_size(), config.transport_max_frame_size(), + config.max_ws_connections_per_user(), ) }) .await; diff --git a/rust/cubestore/cubestore/src/http/mod.rs b/rust/cubestore/cubestore/src/http/mod.rs index a1de965ea3b8f..2f2893d0f8a1d 100644 --- a/rust/cubestore/cubestore/src/http/mod.rs +++ b/rust/cubestore/cubestore/src/http/mod.rs @@ -31,10 +31,11 @@ use log::error; use log::info; use log::trace; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::convert::TryFrom; use std::error::Error as StdError; use std::net::SocketAddr; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; use tempfile::NamedTempFile; use tokio::fs::File; @@ -51,6 +52,34 @@ use warp::reject::Reject; /// process because it is too large (RFC 6455 section 7.4.1, "Message Too Big"). const MESSAGE_TOO_BIG_CLOSE_CODE: u16 = 1009; +/// Close code the WebSocket protocol reserves for a peer that should come back +/// later (RFC 6455 section 7.4.1, "Try Again Later"). An evicted connection was +/// recycled to make room for a newer one, not broken, and a client that still +/// needs it should reconnect rather than report a transport failure. +const CONNECTION_EVICTED_CLOSE_CODE: u16 = 1013; + +/// How long the server waits for an evicted connection's close frame to go out. +/// +/// The connection's entry leaves the counter when it is evicted, not when its +/// task ends, so this is also how long its descriptor can outlive the slot it +/// accounted for: the steady-state overshoot is the eviction rate times this +/// bound. Sending is what can block — a peer whose receive window is closed +/// keeps the frame in our buffer indefinitely — and a close frame is a handful +/// of bytes, so a peer that cannot take them within this long is not going to. +const EVICTED_CLOSE_TIMEOUT: Duration = Duration::from_secs(2); + +/// How long a caller waits for the connection counter's lock before giving up. +/// +/// Nothing awaits inside the critical section — it is a map lookup and an +/// insert — so waiting at all is barely reachable by the code as written; the +/// bound is here so that knowing this cannot stall a connection needs no +/// reasoning about the callers. It is deliberately generous rather than tight: +/// a thread holding the lock can be descheduled for a long time under CPU +/// pressure, and giving up then would quietly stop counting connections +/// exactly when the cap matters most. Both callers log when they give up, so a +/// bound that is ever reached is visible rather than silent. +const COUNTER_LOCK_TIMEOUT: Duration = Duration::from_secs(1); + /// How much room the transport is given above the configured sizes. /// /// The size limit is enforced here rather than by the transport, so that an @@ -132,10 +161,166 @@ pub struct HttpServer { cancel_token: CancellationToken, max_message_size: usize, max_frame_size: usize, + ws_connections: Arc, } crate::di_service!(HttpServer, []); +/// Caps how many concurrent websocket connections one authenticated user may +/// hold. A client that leaks connections otherwise exhausts the process file +/// descriptor table for every other user sharing the node. +/// +/// At the limit the user's oldest connection is closed to admit the new one, +/// rather than the new one being refused: a client that still needs the closed +/// connection reconnects, while one that had forgotten about it simply loses it. +/// A limit of 0 disables the cap. +pub struct WsConnectionCounter { + limit: usize, + /// Live connections per user, keyed by a monotonic id so the first entry is the + /// oldest, holding the token that asks the connection task to close. + /// + /// Synchronous on purpose, not a `tokio` lock behind `util::lock::acquire_lock`: + /// entries are removed from `WsConnectionGuard::drop`, and a destructor cannot + /// `.await`. Releasing anywhere else would leak the slot on a panic, a cancelled + /// task, or a connection that dies before the upgrade completes. Nothing awaits + /// inside the critical section either — `acquire` is not an `async fn`, so one + /// cannot be added — which is the shape the `tokio::sync::Mutex` docs point at a + /// standard-library lock for. Every caller still takes it with a bound, so the + /// property does not rest on that reasoning holding: see COUNTER_LOCK_TIMEOUT. + users: parking_lot::Mutex>>, + next_id: AtomicU64, +} + +/// `None` when the lock could not be taken within [`COUNTER_LOCK_TIMEOUT`]. +/// Callers fail open: the cap is a safety net, so a counter that is briefly +/// unavailable must not be what refuses or stalls a connection. +fn try_lock_users( + users: &parking_lot::Mutex>>, +) -> Option>>> { + users.try_lock_for(COUNTER_LOCK_TIMEOUT) +} + +/// Frees the slot taken by [`WsConnectionCounter::acquire`] when the connection +/// task ends, however it ends, and carries the token that task waits on. +pub struct WsConnectionGuard { + counter: Arc, + user: Option, + id: u64, + cancel: CancellationToken, +} + +impl WsConnectionCounter { + pub fn new(limit: usize) -> Arc { + Arc::new(Self { + limit, + users: parking_lot::Mutex::new(HashMap::new()), + next_id: AtomicU64::new(0), + }) + } + + /// Unauthenticated connections and a disabled cap stay untracked. + pub fn acquire(self: &Arc, user: Option<&str>) -> WsConnectionGuard { + let cancel = CancellationToken::new(); + let user = match user { + Some(user) if self.limit > 0 => user, + _ => { + return WsConnectionGuard { + counter: self.clone(), + user: None, + id: 0, + cancel, + } + } + }; + + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + let evicted = { + let mut users = match try_lock_users(&self.users) { + Some(users) => users, + None => { + log::error!( + "Timed out locking the websocket connection counter; admitting an untracked connection (user: {})", + user, + ); + return WsConnectionGuard { + counter: self.clone(), + user: None, + id: 0, + cancel, + }; + } + }; + let entries = users.entry(user.to_string()).or_default(); + let evicted = if entries.len() >= self.limit { + let oldest = *entries.keys().next().expect("a full map has an entry"); + // Removed here rather than left to the victim's own `drop`, which + // runs later: otherwise this admission would overshoot the limit. + entries.remove(&oldest) + } else { + None + }; + entries.insert(id, cancel.clone()); + evicted + }; + + // Outside the lock. Cancelling only wakes the victim's task, so this cannot + // deadlock, but there is no reason to hold the lock across it. + if let Some(evicted) = evicted { + evicted.cancel(); + } + + WsConnectionGuard { + counter: self.clone(), + user: Some(user.to_string()), + id, + cancel, + } + } + + pub fn count(&self, user: &str) -> usize { + try_lock_users(&self.users) + .and_then(|users| users.get(user).map(BTreeMap::len)) + .unwrap_or(0) + } +} + +impl WsConnectionGuard { + /// Resolves when this connection has been chosen to make room for a newer one. + pub async fn evicted(&self) { + self.cancel.cancelled().await + } +} + +impl Drop for WsConnectionGuard { + fn drop(&mut self) { + let user = match &self.user { + Some(user) => user, + None => return, + }; + let mut users = match try_lock_users(&self.counter.users) { + Some(users) => users, + None => { + // The entry stays behind, but it is the oldest one of that user by + // construction, so the next admission evicts it: cancelling an + // already finished connection's token is a no-op. + log::error!( + "Timed out locking the websocket connection counter; leaving a finished connection's slot to be evicted (user: {})", + user, + ); + return; + } + }; + if let Some(entries) = users.get_mut(user) { + // Idempotent: an evicting `acquire` may have removed this entry already. + entries.remove(&self.id); + // Dropped when empty so a churn of one-off users can't grow the map. + if entries.is_empty() { + users.remove(user); + } + } + } +} + #[derive(Debug)] pub enum CubeRejection { NotAuthorized, @@ -187,6 +372,7 @@ impl HttpServer { drop_complete_messages_after: Duration, max_message_size: usize, max_frame_size: usize, + max_ws_connections_per_user: usize, ) -> Arc { Arc::new(Self { bind_address, @@ -197,6 +383,7 @@ impl HttpServer { drop_complete_messages_after, max_message_size, max_frame_size, + ws_connections: WsConnectionCounter::new(max_ws_connections_per_user), worker_loop: WorkerLoop::new("HttpServer message processing"), drop_orphaned_messages_loop: WorkerLoop::new("HttpServer drop orphaned messages"), cancel_token: CancellationToken::new(), @@ -238,19 +425,47 @@ impl HttpServer { let context_filter_to_move = context_filter.clone(); let max_frame_size = self.max_frame_size.clone(); let max_message_size = self.max_message_size.clone(); + let ws_connections = self.ws_connections.clone(); + let ws_connections_filter = warp::any().map(move || ws_connections.clone()); let query_route = warp::path!("ws") .and(context_filter_to_move) + .and(ws_connections_filter) .and(warp::ws::ws()) - .and_then(move |tx: mpsc::Sender<(mpsc::Sender>, SqlQueryContext, HttpMessage)>, sql_query_context: SqlQueryContext, ws: Ws| async move { + .and_then(move |tx: mpsc::Sender<(mpsc::Sender>, SqlQueryContext, HttpMessage)>, sql_query_context: SqlQueryContext, ws_connections: Arc, ws: Ws| async move { let tx_to_move = tx.clone(); let sql_query_context = sql_query_context.clone(); + let connection_guard = ws_connections.acquire(sql_query_context.user.as_deref()); let reply = ws.max_frame_size(max_frame_size.saturating_mul(TRANSPORT_SIZE_HEADROOM)).max_message_size(max_message_size.saturating_mul(TRANSPORT_SIZE_HEADROOM)).on_upgrade(async move |mut web_socket| { + // Lives as long as the connection task; dropping it frees the slot. + let connection_guard = connection_guard; let process_id = sql_query_context.process_id.as_deref().unwrap_or("None"); trace!("WebSocket connection established (process_id: {})", process_id); let (response_tx, mut response_rx) = mpsc::channel::>(10000); loop { tokio::select! { + _ = connection_guard.evicted() => { + log::warn!( + "Closing websocket connection to admit a newer one for the same user (process_id: {})", + process_id, + ); + // A close frame naming the reason rather than a bare drop, + // so the client can tell being recycled from a network + // blip. Bounded, because the descriptor is already + // unaccounted for: see EVICTED_CLOSE_TIMEOUT. + let close = web_socket.send(Message::close_with( + CONNECTION_EVICTED_CLOSE_CODE, + "connection evicted to admit a newer one for the same user", + )); + match tokio::time::timeout(EVICTED_CLOSE_TIMEOUT, close).await { + Ok(Ok(())) => {} + Ok(Err(e)) => error!("Websocket close send error: {:?}", e), + Err(_) => log::warn!( + "Timed out sending the close frame of an evicted websocket connection" + ), + } + break; + } Some(res) = response_rx.recv() => { trace!("Sending web socket response (process_id: {})", process_id); let send_res = web_socket.send(Message::binary(res.bytes())).await; @@ -1323,6 +1538,131 @@ mod tests { } } + #[test] + fn acquire_evicts_the_oldest_connection_to_admit_a_new_one() { + let counter = WsConnectionCounter::new(2); + let oldest = counter.acquire(Some("tenant-a")); + let newer = counter.acquire(Some("tenant-a")); + assert_eq!(counter.count("tenant-a"), 2); + + let newcomer = counter.acquire(Some("tenant-a")); + + assert!( + oldest.cancel.is_cancelled(), + "the oldest connection is the one asked to close" + ); + assert!(!newer.cancel.is_cancelled()); + assert!(!newcomer.cancel.is_cancelled()); + assert_eq!( + counter.count("tenant-a"), + 2, + "the victim leaves before the newcomer is inserted, so the limit never overshoots" + ); + + // The victim's own drop runs later and must not take another entry with it, + // which is why entries are keyed by a monotonic id. + drop(oldest); + assert_eq!(counter.count("tenant-a"), 2); + + // Eviction keeps following insertion order. + let latest = counter.acquire(Some("tenant-a")); + assert!( + newer.cancel.is_cancelled(), + "now the second connection is oldest" + ); + assert!(!newcomer.cancel.is_cancelled()); + assert!(!latest.cancel.is_cancelled()); + } + + #[test] + fn ws_connection_counter_caps_each_user_independently() { + let counter = WsConnectionCounter::new(1); + let a = counter.acquire(Some("tenant-a")); + let b = counter.acquire(Some("tenant-b")); + + assert!( + !a.cancel.is_cancelled() && !b.cancel.is_cancelled(), + "one connection each is within the cap" + ); + + let a_again = counter.acquire(Some("tenant-a")); + assert!(a.cancel.is_cancelled(), "tenant-a is at its cap"); + assert!( + !b.cancel.is_cancelled(), + "one user at its cap must not affect another" + ); + assert!(!a_again.cancel.is_cancelled()); + assert_eq!(counter.count("tenant-a"), 1); + assert_eq!(counter.count("tenant-b"), 1); + } + + #[test] + fn ws_connection_counter_forgets_users_that_dropped_to_zero() { + let counter = WsConnectionCounter::new(1); + + drop(counter.acquire(Some("tenant-a"))); + assert_eq!(counter.count("tenant-a"), 0); + assert!( + try_lock_users(&counter.users) + .expect("uncontended") + .is_empty(), + "a churn of one-off users must not grow the map" + ); + + // Both guards are bound to names: one left as a temporary dies at the end of + // its own statement and would free the slot it is meant to hold. + let _b = counter.acquire(Some("tenant-b")); + let _c = counter.acquire(Some("tenant-c")); + assert_eq!( + try_lock_users(&counter.users).expect("uncontended").len(), + 2 + ); + } + + #[test] + fn acquire_gives_up_on_a_stuck_lock_instead_of_waiting_on_it() { + let counter = WsConnectionCounter::new(1); + let blocker = counter.users.lock(); + + // From another thread: the lock is not reentrant, and the point is that a + // caller returns on its own schedule rather than the lock holder's. + let probe = Arc::clone(&counter); + let guard = std::thread::spawn(move || probe.acquire(Some("tenant-a"))) + .join() + .expect("acquire must return rather than wait for the lock"); + + assert!( + guard.user.is_none(), + "the cap is a safety net, so an unavailable counter admits the connection untracked rather than refusing or stalling it" + ); + assert!(!guard.cancel.is_cancelled()); + + drop(blocker); + assert_eq!(counter.count("tenant-a"), 0); + } + + #[test] + fn ws_connection_counter_leaves_untracked_what_it_cannot_attribute() { + let disabled = WsConnectionCounter::new(0); + let mut held = Vec::new(); + for _ in 0..1000 { + held.push(disabled.acquire(Some("tenant-a"))); + } + assert!( + held.iter().all(|guard| !guard.cancel.is_cancelled()), + "limit 0 disables the cap" + ); + assert_eq!(disabled.count("tenant-a"), 0); + + let counter = WsConnectionCounter::new(1); + let first = counter.acquire(None); + let second = counter.acquire(None); + assert!( + !first.cancel.is_cancelled() && !second.cancel.is_cancelled(), + "unauthenticated connections are not attributable and stay untracked" + ); + } + #[tokio::test] async fn upload_writes_the_whole_body_before_handing_off_the_path() -> Result<(), CubeError> { let dir = tempfile::tempdir()?; @@ -1736,6 +2076,7 @@ mod tests { Duration::from_millis(1000), config.transport_max_message_size(), config.transport_max_frame_size(), + config.max_ws_connections_per_user(), )); { let http_server = http_server.clone(); @@ -1939,6 +2280,7 @@ mod tests { Duration::from_millis(1000), config.transport_max_message_size(), config.transport_max_frame_size(), + config.max_ws_connections_per_user(), )); { let http_server = http_server.clone(); @@ -1976,6 +2318,84 @@ mod tests { Ok(()) } + /// An evicted connection is told why it is going away, so a client can tell + /// being recycled from a network blip, and the cap actually reaches the + /// connection rather than only the counter. + #[tokio::test] + async fn ws_connection_evicted_test() -> Result<(), CubeError> { + init_test_logger().await; + + let mut auth = MockSqlAuthService::new(); + auth.expect_authenticate().return_const(Ok(None)); + + let http_server = Arc::new(HttpServer::new( + "127.0.0.1:53036".to_string(), + Arc::new(auth), + Arc::new(SqlServiceMock { + message_counter: AtomicU64::new(0), + }), + Duration::from_millis(100), + Duration::from_millis(10000), + Duration::from_millis(1000), + 64 * 1024, + 64 * 1024, + // One connection per user, so the second one has to evict the first. + 1, + )); + { + let http_server = http_server.clone(); + cube_ext::spawn(async move { http_server.run_server().await }); + } + + tokio::time::sleep(Duration::from_secs(1)).await; + + fn connect_request(user: &str) -> Request { + let mut request = "ws://127.0.0.1:53036/ws".into_client_request().unwrap(); + request.headers_mut().insert( + "authorization", + HeaderValue::from_str(&Credentials::new(user, "").as_http_header()).unwrap(), + ); + request + } + + let (mut first, _) = connect_async(connect_request("tenant-a")) + .await + .expect("the first connection is within the cap"); + let (second, _) = connect_async(connect_request("tenant-a")) + .await + .expect("the second connection is admitted by evicting the first"); + + // Bounded, so a regression in the eviction path fails this test instead of + // hanging it: nothing else would ever wake this read. + let msg = tokio::time::timeout(Duration::from_secs(10), first.next()) + .await + .expect("the evicted connection must be closed promptly") + .expect("the evicted connection is closed, not left open") + .unwrap(); + match msg { + Message::Close(Some(frame)) => { + assert_eq!(u16::from(frame.code), CONNECTION_EVICTED_CLOSE_CODE); + assert!( + frame.reason.contains("evicted"), + "unexpected close reason: {}", + frame.reason + ); + } + msg => panic!("Close frame expected, got: {:?}", msg), + } + + // Another user is unaffected by tenant-a having been at its cap. + let (other, _) = connect_async(connect_request("tenant-b")) + .await + .expect("a different user has its own slot"); + + drop(other); + drop(second); + + http_server.stop_processing().await; + Ok(()) + } + /// An incoming message past the transport backstop is answered with the /// WebSocket "message too big" close code instead of the connection being /// dropped without a word, which the client can only read as a bare @@ -1999,6 +2419,7 @@ mod tests { Duration::from_millis(1000), max_message_size, max_message_size, + 0, // no websocket connection cap )); { let http_server = http_server.clone(); @@ -2077,6 +2498,7 @@ mod tests { Duration::from_millis(1000), max_message_size, max_frame_size, + 0, // no websocket connection cap )); { let http_server = http_server.clone(); @@ -2151,6 +2573,7 @@ mod tests { Duration::from_millis(1000), max_message_size, max_message_size, + 0, // no websocket connection cap )); { let http_server = http_server.clone();