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/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 () => { 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();