diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts index ad63f10ba1cac..8a222cd7ca755 100644 --- a/packages/cubejs-backend-shared/src/env.ts +++ b/packages/cubejs-backend-shared/src/env.ts @@ -1877,6 +1877,19 @@ const variables: Record any> = { cubeStoreNoHeartBeatTimeout: () => get('CUBEJS_CUBESTORE_NO_HEART_BEAT_TIMEOUT') .default('30') .asInt(), + /** + * Maximum size in bytes of a single message exchanged with Cube Store, both + * of a query sent to it and of a response received from it. + * + * It is the only limit that applies to responses, since Cube Store doesn't + * cap what it sends. For queries it is independent of, and by default looser + * than, CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE (64 MB), which is what Cube + * Store itself accepts: a query over that but under this one is refused by + * Cube Store rather than by this limit. + */ + cubeStoreMaxMessageSize: () => get('CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE') + .default(String(100 * 1024 * 1024)) + .asIntPositive(), cubeStoreRollingWindowJoin: () => get('CUBEJS_CUBESTORE_ROLLING_WINDOW_JOIN') .default('true') .asBoolStrict(), diff --git a/packages/cubejs-cubestore-driver/package.json b/packages/cubejs-cubestore-driver/package.json index 9c4b3a27912ad..c0538e0a3e374 100644 --- a/packages/cubejs-cubestore-driver/package.json +++ b/packages/cubejs-cubestore-driver/package.json @@ -22,8 +22,9 @@ "build": "rm -rf dist && npm run tsc", "tsc": "tsc", "watch": "tsc -w", - "lint": "eslint src/*.ts", - "lint:fix": "eslint --fix src/*.ts" + "lint": "eslint src/*.ts test/*.ts", + "lint:fix": "eslint --fix src/*.ts test/*.ts", + "unit": "jest --coverage" }, "dependencies": { "@cubejs-backend/base-driver": "1.7.16", diff --git a/packages/cubejs-cubestore-driver/src/WebSocketConnection.ts b/packages/cubejs-cubestore-driver/src/WebSocketConnection.ts index 6e70d08b2378b..766800142e005 100644 --- a/packages/cubejs-cubestore-driver/src/WebSocketConnection.ts +++ b/packages/cubejs-cubestore-driver/src/WebSocketConnection.ts @@ -4,7 +4,7 @@ import { v4 as uuidv4 } from 'uuid'; import { InlineTable } from '@cubejs-backend/base-driver'; import { getEnv, getProcessUid } from '@cubejs-backend/shared'; import { parseCubestoreResultMessage } from '@cubejs-backend/native'; -import { ConnectionError, QueryError } from './errors'; +import { ConnectionError, MessageTooLargeError, QueryError } from './errors'; import { BinaryValue, BoolValue, @@ -22,10 +22,36 @@ import { StringValue, } from '../codegen'; +// The WebSocket close code for a message that is too big to be processed: `ws` +// closes with it when an incoming message is over `maxPayload`, and a peer that +// refuses a message of ours is expected to close with it as well. +const MESSAGE_TOO_BIG_CLOSE_CODE = 1009; + +// The `ws` error code for an incoming message bigger than `maxPayload`. +const MAX_PAYLOAD_EXCEEDED_CODE = 'WS_ERR_UNSUPPORTED_MESSAGE_LENGTH'; + +function formatSize(bytes: number): string { + const units: [number, string][] = [[1024 * 1024, 'MB'], [1024, 'KB']]; + + for (const [unit, name] of units) { + if (bytes >= unit) { + return `${Math.round((bytes / unit) * 10) / 10} ${name}`; + } + } + + return `${bytes} bytes`; +} + interface SentMessage { resolve: (value: any) => void; reject: (reason?: any) => void; buffer: Uint8Array; + // How many times this message was re-sent over a freshly established + // connection. Used to give up instead of retrying forever. + resendCount: number; + // How many times this message was in flight when the connection died from a + // failure that can't be attributed to a single message. + fatalRounds: number; } export type QueryParameter = null | boolean | number | string | Buffer; @@ -40,7 +66,13 @@ interface CubeStoreWebSocket extends WebSocket { readyPromise: Promise; lastHeartBeat: Date; sentMessages: Record; - sendAsync: (message: Uint8Array) => Promise; + sendAsync: (message: Uint8Array, messageId?: number) => Promise; + // Set as soon as the 'close' handler has scheduled a re-send of the messages + // that are still in flight on this socket. + resendScheduled: boolean; + // A failure that killed this socket and that re-sending can't fix, so pending + // messages are rejected with it instead of being retried. + fatalError: Error | null; } export class WebSocketConnection { @@ -50,6 +82,8 @@ export class WebSocketConnection { protected readonly noHeartBeatTimeout: number; + protected readonly maxMessageSize: number; + protected currentConnectionTry: number; protected webSocket: CubeStoreWebSocket | null = null; @@ -65,6 +99,7 @@ export class WebSocketConnection { this.messageCounter = 1; this.maxConnectRetries = getEnv('cubeStoreMaxConnectRetries'); this.noHeartBeatTimeout = getEnv('cubeStoreNoHeartBeatTimeout'); + this.maxMessageSize = getEnv('cubeStoreMaxMessageSize'); this.currentConnectionTry = 0; this.connectionId = uuidv4(); } @@ -74,7 +109,7 @@ export class WebSocketConnection { const headers: Record = {}; headers['x-process-id'] = getProcessUid(); - const webSocket = new WebSocket(this.url, { headers }) as CubeStoreWebSocket; + const webSocket = new WebSocket(this.url, { headers, maxPayload: this.maxMessageSize }) as CubeStoreWebSocket; webSocket.on('upgrade', (response: any) => { this.cubeStoreVersion = response.headers['x-cubestore-version'] || null; }); @@ -91,23 +126,50 @@ export class WebSocketConnection { } }, 5000); - webSocket.sendAsync = async (message: Uint8Array) => new Promise((resolveSend, rejectSend) => { + webSocket.sendAsync = async (message: Uint8Array, messageId?: number) => new Promise((resolveSend) => { // If socket is closing this message should be resent - if (webSocket.readyState === WebSocket.OPEN) { - webSocket.send(message, (err) => { - if (err) { - rejectSend(new ConnectionError( - `CubeStore connection error: ${err.message}`, - err - )); - } else { - resolveSend(); - } - }); + if (webSocket.readyState !== WebSocket.OPEN) { + resolveSend(); + return; } + + webSocket.send(message, (err) => { + if (err) { + // The write failed (EPIPE/ECONNRESET when Cube Store dropped the + // connection). The message stays registered in `sentMessages`, so + // it's re-sent once this socket is closed and a new one is + // established -- failing it here would surface a spurious + // `write EPIPE` to the user for a perfectly retryable query. + this.handleSendError(webSocket, err, messageId); + } + + resolveSend(); + }); }); webSocket.on('open', () => resolve(webSocket)); webSocket.on('error', (err) => { + if ((err as any).code === MAX_PAYLOAD_EXCEEDED_CODE) { + // Cube Store answered with a message bigger than this connection + // accepts, and `ws` is tearing the connection down. Neither + // reconnecting nor retrying the query helps: the response would be + // just as big. Pending messages are rejected by the 'close' handler. + webSocket.fatalError = new MessageTooLargeError( + `Cube Store response size exceeds the maximum message size of ${formatSize(this.maxMessageSize)}. ` + + 'Reduce the amount of data the query returns, e.g. by adding filters or a limit, ' + + 'or raise CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE.', + err + ); + + if (webSocket === this.webSocket) { + this.webSocket = null; + } + + // No-op if the connection was already established. + reject(webSocket.fatalError); + + return; + } + this.currentConnectionTry += 1; if (this.currentConnectionTry < this.maxConnectRetries) { @@ -131,17 +193,85 @@ export class WebSocketConnection { } webSocket.lastHeartBeat = new Date(); }); - webSocket.on('close', () => { + webSocket.on('close', (code: number) => { clearInterval(pingInterval); if (Object.keys(webSocket.sentMessages).length) { + const fatalError = webSocket.fatalError || ( + // Cube Store refused a message that didn't fit into its limits. + code === MESSAGE_TOO_BIG_CLOSE_CODE ? new MessageTooLargeError( + 'Cube Store closed the connection: message size exceeds the maximum message size Cube Store accepts. ' + + 'Reduce the size of the query and of the inline tables it sends, or raise ' + + 'CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE on the Cube Store side.' + ) : null + ); + + if (fatalError) { + // The connection multiplexes messages, and an oversized one can't + // be attributed: `ws` drops the frame before its message id is + // read. A message that was alone in flight is certainly the one at + // fault. Otherwise every message gets one more round, which + // answers the innocent ones and usually leaves the offender alone + // on the connection, where the next round does attribute it. What + // is still in flight after that round is failed regardless, so an + // offender that keeps killing the connection before the others are + // answered can't turn into a re-send loop. + const pending = Object.keys(webSocket.sentMessages); + + // eslint-disable-next-line no-restricted-syntax + for (const key of pending) { + const sentMessage = webSocket.sentMessages[key]; + sentMessage.fatalRounds += 1; + + if (pending.length === 1 || sentMessage.fatalRounds > 1) { + delete webSocket.sentMessages[key]; + sentMessage.reject(fatalError); + } + } + + if (!Object.keys(webSocket.sentMessages).length) { + if (webSocket === this.webSocket) { + this.webSocket = null; + } + + return; + } + } + + webSocket.resendScheduled = true; + setTimeout(async () => { try { - const nextWebSocket = await this.initWebSocket(); + const nextWebSocket = await this.openSocket(); + const resent: [string, SentMessage][] = []; + + // Register the whole batch before writing any of it. Writing + // yields, and a socket that closes mid-batch must find every + // message of it in `sentMessages`: the ones not registered yet + // would end up on a socket whose 'close' has already been + // handled, with nobody left to write or to re-send them. // eslint-disable-next-line no-restricted-syntax for (const key of Object.keys(webSocket.sentMessages)) { - nextWebSocket.sentMessages[key] = webSocket.sentMessages[key]; - await nextWebSocket.sendAsync(webSocket.sentMessages[key].buffer); + const sentMessage = webSocket.sentMessages[key]; + + if (sentMessage.resendCount >= this.maxConnectRetries) { + sentMessage.reject(new ConnectionError( + `CubeStore connection lost: message wasn't delivered after ${sentMessage.resendCount} retries` + )); + } else { + sentMessage.resendCount += 1; + nextWebSocket.sentMessages[key] = sentMessage; + resent.push([key, sentMessage]); + } + } + + // eslint-disable-next-line no-restricted-syntax + for (const [key, sentMessage] of resent) { + // Skip what was answered, or failed, while the batch was + // being written. + if (nextWebSocket.sentMessages[key] === sentMessage) { + await nextWebSocket.sendAsync(sentMessage.buffer, Number(key)); + } } } catch (e) { // eslint-disable-next-line no-restricted-syntax @@ -182,6 +312,8 @@ export class WebSocketConnection { }); webSocket.sentMessages = {}; + webSocket.resendScheduled = false; + webSocket.fatalError = null; this.webSocket = webSocket; } @@ -192,26 +324,97 @@ export class WebSocketConnection { return 1000 * (this.currentConnectionTry + 1); } + /** + * Returns a socket that can still carry a message. + * + * `initWebSocket()` resolves as soon as a socket is open, but that socket may + * have been closed again by then. Registering a message on a closed socket + * would strand it: nothing writes it, and the re-send loop of that socket has + * already taken its snapshot, so the message would never settle. + */ + private async openSocket(): Promise { + // A closed socket is dropped from `this.webSocket` by its 'close' handler, + // so the next attempt establishes a new one. + for (let attempt = 0; attempt < 2; attempt++) { + const socket = await this.initWebSocket(); + + if (socket.readyState !== WebSocket.CLOSED) { + return socket; + } + } + + throw new ConnectionError('CubeStore connection is closed'); + } + + /** + * Handles a failed write to an already established socket, e.g. `write EPIPE` + * when Cube Store closed the connection between the `readyState` check and the + * actual write to the underlying TCP socket. + * + * Such a message is not lost: it stays registered in `sentMessages` and is + * re-sent by the 'close' handler over a freshly established connection, so it + * must not be rejected here. The socket is terminated to make sure that + * 'close' (and with it the re-send) really happens. + */ + private handleSendError(webSocket: CubeStoreWebSocket, err: Error, messageId?: number) { + if (webSocket.readyState !== WebSocket.CLOSED) { + if (webSocket.readyState === WebSocket.OPEN) { + // The socket is broken, but `ws` doesn't know it yet. Terminating it + // emits 'close', which re-sends everything still pending on it. + webSocket.terminate(); + } + + // 'close' is still to come and will re-send pending messages. + return; + } + + // The socket is already closed and the re-send loop is not going to pick + // this message up, so there's nothing left to wait for. + if (!webSocket.resendScheduled && messageId !== undefined) { + const sentMessage = webSocket.sentMessages[messageId]; + if (sentMessage) { + delete webSocket.sentMessages[messageId]; + sentMessage.reject(new ConnectionError( + `CubeStore connection error: ${err.message}`, + err + )); + } + } + } + private async sendMessage(messageId: number, buffer: Uint8Array): Promise { - const socket = await this.initWebSocket(); + if (buffer.length > this.maxMessageSize) { + // Cube Store would close the connection on such a message, which shows up + // as an unrelated `write EPIPE`, so report it before sending anything. + // This only catches what is over our own limit: Cube Store applies its + // own, by default stricter, CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE, and a + // message it refuses is reported once it closes the connection. + throw new MessageTooLargeError( + `Cube Store request size of ${formatSize(buffer.length)} exceeds the maximum message size of ` + + `${formatSize(this.maxMessageSize)}. Reduce the size of the query and of the inline tables it sends, ` + + 'or raise CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE together with CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE ' + + 'on the Cube Store side.' + ); + } + + const socket = await this.openSocket(); return new Promise((resolve, reject) => { + socket.sentMessages[messageId] = { + resolve, + reject, + buffer, + resendCount: 0, + fatalRounds: 0, + }; + + // If socket is closing this message should be resent if (socket.readyState === WebSocket.OPEN) { socket.send(buffer, (err) => { if (err) { - delete socket.sentMessages[messageId]; - reject(new ConnectionError( - `CubeStore connection error: ${err.message}`, - err - )); + this.handleSendError(socket, err, messageId); } }); } - - socket.sentMessages[messageId] = { - resolve, - reject, - buffer - }; }); } diff --git a/packages/cubejs-cubestore-driver/src/errors.ts b/packages/cubejs-cubestore-driver/src/errors.ts index df9f52529a375..f0713c42d2c19 100644 --- a/packages/cubejs-cubestore-driver/src/errors.ts +++ b/packages/cubejs-cubestore-driver/src/errors.ts @@ -13,6 +13,19 @@ export class ConnectionError extends CubeStoreError { } } +/** + * A message didn't fit into the size limit of the connection. Unlike other + * connection errors this one is not worth retrying: the same message would be + * rejected again. + */ +export class MessageTooLargeError extends ConnectionError { + public constructor(message: string, cause?: Error) { + super(message, cause); + + this.name = 'MessageTooLargeError'; + } +} + export class QueryError extends CubeStoreError { public constructor(message: string) { super(message); diff --git a/packages/cubejs-cubestore-driver/test/mock-cubestore-server.ts b/packages/cubejs-cubestore-driver/test/mock-cubestore-server.ts new file mode 100644 index 0000000000000..f22121d545a52 --- /dev/null +++ b/packages/cubejs-cubestore-driver/test/mock-cubestore-server.ts @@ -0,0 +1,164 @@ +import { AddressInfo, Socket } from 'net'; +import * as flatbuffers from 'flatbuffers'; +import WebSocket from 'ws'; + +import { + HttpCommand, + HttpError, + HttpMessage, + HttpQuery, + HttpQueryResult, + HttpQueryResultArrow, + HttpQueryResultData, +} from '../codegen'; + +export interface ReceivedMessage { + connectionIndex: number; + messageId: number; + query: string; +} + +export interface MockConnection { + index: number; + ws: WebSocket; + socket: Socket; +} + +export type MessageHandler = (message: ReceivedMessage, connection: MockConnection) => void; + +/** + * Cube Store answers a query either with a result set or with an error. Tests + * use the error variant, because it's the only answer that can be asserted + * without the native result parser, and it's enough to tell "the query reached + * Cube Store and was answered" from "the query failed on the transport". + */ +export function buildErrorMessage(messageId: number, error: string): Buffer { + const builder = new flatbuffers.Builder(1024); + const errorOffset = builder.createString(error); + const commandOffset = HttpError.createHttpError(builder, errorOffset); + const message = HttpMessage.createHttpMessage(builder, messageId, HttpCommand.HttpError, commandOffset, 0); + builder.finish(message); + + return Buffer.from(builder.asUint8Array()); +} + +/** + * A successful answer. Its payload is decoded by the native result parser, so + * tests that use it stub that parser out. + */ +export function buildResultMessage(messageId: number, data: Buffer = Buffer.alloc(0)): Buffer { + const builder = new flatbuffers.Builder(1024); + const dataOffset = HttpQueryResultArrow.createDataVector(builder, data); + const arrowOffset = HttpQueryResultArrow.createHttpQueryResultArrow(builder, dataOffset, true); + const commandOffset = HttpQueryResult.createHttpQueryResult( + builder, + HttpQueryResultData.HttpQueryResultArrow, + arrowOffset + ); + const message = HttpMessage.createHttpMessage(builder, messageId, HttpCommand.HttpQueryResult, commandOffset, 0); + builder.finish(message); + + return Buffer.from(builder.asUint8Array()); +} + +export function answeredBy(connectionIndex: number): string { + return `answered by connection #${connectionIndex}`; +} + +/** + * A minimal Cube Store look-alike: it speaks the same WebSocket + flatbuffers + * protocol, so the driver talks to it over real TCP sockets, which can then be + * broken in the exact ways a real Cube Store restart breaks them. + */ +export class MockCubeStoreServer { + public readonly connections: MockConnection[] = []; + + public readonly received: ReceivedMessage[] = []; + + /** + * Replies to every query with an error naming the connection that received + * it, so a test can tell which connection answered. Can be replaced to + * emulate a Cube Store that goes away instead of answering. + */ + public handler: MessageHandler = (message, connection) => { + connection.ws.send(buildErrorMessage(message.messageId, answeredBy(connection.index))); + }; + + protected constructor(protected readonly wss: WebSocket.Server) { + wss.on('connection', (ws: WebSocket, request: any) => { + const connection: MockConnection = { + index: this.connections.length, + ws, + socket: request.socket, + }; + this.connections.push(connection); + + ws.on('message', (raw: Buffer) => { + const httpMessage = HttpMessage.getRootAsHttpMessage(new flatbuffers.ByteBuffer(raw)); + const message: ReceivedMessage = { + connectionIndex: connection.index, + messageId: httpMessage.messageId(), + query: httpMessage.command(new HttpQuery())?.query() || '', + }; + this.received.push(message); + this.handler(message, connection); + }); + // Connections are torn down by the tests on purpose, nothing to report. + ws.on('error', () => { + // noop + }); + }); + } + + public static async start(): Promise { + const wss = new WebSocket.Server({ host: '127.0.0.1', port: 0 }); + await new Promise((resolve, reject) => { + wss.once('listening', resolve); + wss.once('error', reject); + }); + + return new MockCubeStoreServer(wss); + } + + public get url(): string { + return `ws://127.0.0.1:${(this.wss.address() as AddressInfo).port}`; + } + + public connection(index: number): MockConnection { + if (!this.connections[index]) { + throw new Error(`Connection #${index} was never established`); + } + + return this.connections[index]; + } + + public async waitForConnections(count: number, timeout: number = 20000): Promise { + await this.waitFor(() => this.connections.length >= count, timeout, `${count} connection(s)`); + } + + public async waitForMessages(count: number, timeout: number = 20000): Promise { + await this.waitFor(() => this.received.length >= count, timeout, `${count} message(s)`); + } + + protected async waitFor(condition: () => boolean, timeout: number, description: string): Promise { + const deadline = Date.now() + timeout; + + while (!condition()) { + if (Date.now() > deadline) { + throw new Error(`Timed out waiting for ${description}`); + } + + await new Promise((resolve) => { setTimeout(resolve, 25); }); + } + } + + public async stop(): Promise { + for (const connection of this.connections) { + connection.ws.terminate(); + } + + await new Promise((resolve) => { + this.wss.close(() => resolve()); + }); + } +} diff --git a/packages/cubejs-cubestore-driver/test/websocket-connection.test.ts b/packages/cubejs-cubestore-driver/test/websocket-connection.test.ts new file mode 100644 index 0000000000000..b122ef4d21391 --- /dev/null +++ b/packages/cubejs-cubestore-driver/test/websocket-connection.test.ts @@ -0,0 +1,354 @@ +import { Socket } from 'net'; + +import { WebSocketConnection } from '../src/WebSocketConnection'; +import { MessageTooLargeError, QueryError } from '../src/errors'; +import { QueryResultFormat } from '../codegen'; +import { + answeredBy, + buildErrorMessage, + buildResultMessage, + MockConnection, + MockCubeStoreServer, +} from './mock-cubestore-server'; + +const QUERY_RESULT = [{ answer: 42 }]; + +// Decoding a result set is a native addon and is orthogonal to the transport +// under test, so only that step is stubbed: the socket, the WebSocket framing +// and the flatbuffers protocol stay real. +jest.mock('@cubejs-backend/native', () => ({ + parseCubestoreResultMessage: jest.fn(async () => [{ answer: 42 }]), +})); + +const JEST_TIMEOUT = 60 * 1000; + +/** + * The error Node hands to a pending write when the peer is gone, as seen in + * `ConnectionError: CubeStore connection error: write EPIPE`. + */ +const epipe = () => Object.assign(new Error('write EPIPE'), { + code: 'EPIPE', + errno: -32, + syscall: 'write', +}); + +/** + * Waits until the frame the driver just handed to `ws` has reached the socket + * write buffer, where a corked socket holds it. + */ +const waitForBufferedWrite = async (socket: Socket) => { + const deadline = Date.now() + 5000; + + while (!socket.writableLength) { + if (Date.now() > deadline) { + throw new Error('Timed out waiting for a buffered write'); + } + + await new Promise((resolve) => { setImmediate(resolve); }); + } +}; + +/** + * Records which messages were registered on a socket when the re-send loop + * wrote the first message of a batch. `sendAsync` is only used to re-send, so + * its first call is that write. + */ +class ObservableConnection extends WebSocketConnection { + public registeredAtFirstResend: string[] | null = null; + + protected async initWebSocket(): Promise { + const socket: any = await super.initWebSocket(); + + if (!socket.sendAsyncObserved) { + socket.sendAsyncObserved = true; + + const { sendAsync } = socket; + socket.sendAsync = async (message: Uint8Array, messageId?: number) => { + if (this.registeredAtFirstResend === null) { + this.registeredAtFirstResend = Object.keys(socket.sentMessages); + } + + return sendAsync(message, messageId); + }; + } + + return socket; + } +} + +describe('WebSocketConnection', () => { + let server: MockCubeStoreServer; + let connection: WebSocketConnection | null = null; + + beforeEach(async () => { + server = await MockCubeStoreServer.start(); + connection = null; + }); + + afterEach(async () => { + connection?.close(); + await server.stop(); + }); + + const query = (sql: string) => connection!.query(sql, [], { responseFormat: QueryResultFormat.Legacy }); + + /** + * The mock answers every query with an error naming the connection that + * served it, so a rejection carrying that marker means the query made a full + * round trip. A rejection with anything else (`ConnectionError: ... write + * EPIPE`) means the driver dropped the query instead of delivering it. + */ + const expectAnsweredBy = async (promise: Promise, connectionIndex: number) => { + await expect(promise).rejects.toThrow(QueryError); + await expect(promise).rejects.toThrow(answeredBy(connectionIndex)); + }; + + // The socket the driver is writing to right now. + const clientSocket = (): Socket => (connection as any).webSocket._socket; + + it('resolves a query with the result Cube Store sent', async () => { + connection = new WebSocketConnection(server.url); + + server.handler = (message, mockConnection) => { + mockConnection.ws.send(buildResultMessage(message.messageId)); + }; + + await expect(query('SELECT 1')).resolves.toEqual(QUERY_RESULT); + + // And the same after the connection had to be re-established mid-query. + const socket = clientSocket(); + socket.cork(); + const promise = query('SELECT 2'); + await waitForBufferedWrite(socket); + socket.destroy(epipe()); + + await expect(promise).resolves.toEqual(QUERY_RESULT); + expect(server.received.map((message) => message.query)).toEqual(['SELECT 1', 'SELECT 2']); + }, JEST_TIMEOUT); + + it('resends a query when the write fails with EPIPE', async () => { + connection = new WebSocketConnection(server.url); + + // Establish the connection with a first, successfully answered query. + await expectAnsweredBy(query('SELECT 1'), 0); + + // Keep the outgoing frame in the socket write buffer, then break the socket + // the way Node does once Cube Store is gone: the buffered write fails with + // EPIPE while `ws` still reports the connection as OPEN. + const socket = clientSocket(); + socket.cork(); + const promise = query('SELECT 2'); + await waitForBufferedWrite(socket); + socket.destroy(epipe()); + + // The query never reached Cube Store, so it has to be resent over a new + // connection instead of failing with the write error. + await expectAnsweredBy(promise, 1); + + expect(server.received.map((message) => [message.connectionIndex, message.query])).toEqual([ + [0, 'SELECT 1'], + [1, 'SELECT 2'], + ]); + }, JEST_TIMEOUT); + + it('resends every query that was in flight when the write failed', async () => { + connection = new WebSocketConnection(server.url); + + await expectAnsweredBy(query('SELECT 1'), 0); + + const socket = clientSocket(); + socket.cork(); + const promises = [query('SELECT 2'), query('SELECT 3'), query('SELECT 4')]; + await waitForBufferedWrite(socket); + socket.destroy(epipe()); + + await Promise.all(promises.map((promise) => expectAnsweredBy(promise, 1))); + + expect( + server.received.filter((message) => message.connectionIndex === 1).map((message) => message.query).sort() + ).toEqual(['SELECT 2', 'SELECT 3', 'SELECT 4']); + }, JEST_TIMEOUT); + + it('registers the whole re-sent batch before writing any of it', async () => { + const observed = new ObservableConnection(server.url); + connection = observed; + + await expectAnsweredBy(query('SELECT 1'), 0); + + const socket = clientSocket(); + socket.cork(); + const promises = [query('SELECT 2'), query('SELECT 3'), query('SELECT 4')]; + await waitForBufferedWrite(socket); + socket.destroy(epipe()); + + await Promise.all(promises.map((promise) => expectAnsweredBy(promise, 1))); + + // Writing yields, so the rest of the batch has to be registered before the + // first message of it is written: a connection dying in between would + // otherwise never see them, and they would never settle. + expect(observed.registeredAtFirstResend).toEqual(['2', '3', '4']); + }, JEST_TIMEOUT); + + it('resends a query when the socket is no longer writable', async () => { + connection = new WebSocketConnection(server.url); + + await expectAnsweredBy(query('SELECT 1'), 0); + + // Half-close the socket: `ws` still reports OPEN, but the write fails. + clientSocket().end(); + + await expectAnsweredBy(query('SELECT 2'), 1); + }, JEST_TIMEOUT); + + it('resends a query when Cube Store closes the connection without answering', async () => { + connection = new WebSocketConnection(server.url); + + server.handler = (message, mockConnection) => { + if (mockConnection.index === 0) { + // Cube Store went away in the middle of the query. + mockConnection.ws.close(); + return; + } + + mockConnection.ws.send(buildErrorMessage(message.messageId, answeredBy(mockConnection.index))); + }; + + await expectAnsweredBy(query('SELECT 1'), 1); + }, JEST_TIMEOUT); + + describe('message size limit', () => { + const MAX_MESSAGE_SIZE = 1024 * 1024; + + beforeEach(() => { + process.env.CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE = String(MAX_MESSAGE_SIZE); + }); + + afterEach(() => { + delete process.env.CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE; + }); + + it('reports a response that is over the limit instead of retrying it', async () => { + connection = new WebSocketConnection(server.url); + + server.handler = (message, mockConnection) => { + mockConnection.ws.send(Buffer.alloc(MAX_MESSAGE_SIZE * 2)); + }; + + const promise = query('SELECT 1'); + + await expect(promise).rejects.toThrow(MessageTooLargeError); + await expect(promise).rejects.toThrow( + 'Cube Store response size exceeds the maximum message size of 1 MB. ' + + 'Reduce the amount of data the query returns, e.g. by adding filters or a limit, ' + + 'or raise CUBEJS_CUBESTORE_MAX_MESSAGE_SIZE.' + ); + + // Re-running the query would only produce the same oversized response. + expect(server.received).toHaveLength(1); + }, JEST_TIMEOUT); + + it('resends the other queries in flight and attributes the limit to the offender', async () => { + connection = new WebSocketConnection(server.url); + + // Cube Store answers the small query before it is done producing the big + // one. Ordering the two sends rather than spacing them apart in time + // keeps the test independent of how fast the driver is scheduled: the + // answer to the small query is on the wire, and therefore processed, + // before the oversized frame that tears the connection down. + const answeredSmall = new Set(); + const deferredBig = new Map(); + const sendOversized = (mockConnection: MockConnection) => { + mockConnection.ws.send(Buffer.alloc(MAX_MESSAGE_SIZE * 2)); + }; + + server.handler = (message, mockConnection) => { + if (message.query === 'SELECT big') { + // On the first connection the small query is left in flight, so that + // the oversized response kills it along with the query it belongs to. + if (mockConnection.index === 0 || answeredSmall.has(mockConnection.index)) { + sendOversized(mockConnection); + } else { + deferredBig.set(mockConnection.index, mockConnection); + } + + return; + } + + if (mockConnection.index > 0) { + mockConnection.ws.send(buildErrorMessage(message.messageId, answeredBy(mockConnection.index))); + answeredSmall.add(mockConnection.index); + + const deferred = deferredBig.get(mockConnection.index); + if (deferred) { + deferredBig.delete(mockConnection.index); + sendOversized(deferred); + } + } + }; + + const big = query('SELECT big'); + // Asserted below, handled here so that a rejection arriving earlier than + // expected is reported as a failed assertion and not as an unhandled one. + big.catch(() => { + // noop + }); + const small = query('SELECT small'); + + // The small query is unrelated to the size limit: it gets resent and + // answered rather than failing with an error about a limit it never + // approached. + await expectAnsweredBy(small, 1); + + // Which leaves the offending query alone on the connection, where the + // oversized response can be attributed to it. + await expect(big).rejects.toThrow(MessageTooLargeError); + await expect(big).rejects.toThrow('Cube Store response size exceeds the maximum message size of 1 MB'); + }, JEST_TIMEOUT); + + it('reports a request that is over the limit without sending it', async () => { + connection = new WebSocketConnection(server.url); + + const promise = query(`SELECT ${'x'.repeat(MAX_MESSAGE_SIZE + 1)}`); + + await expect(promise).rejects.toThrow(MessageTooLargeError); + await expect(promise).rejects.toThrow( + /Cube Store request size of \d+(\.\d+)? MB exceeds the maximum message size of 1 MB/ + ); + + expect(server.connections).toHaveLength(0); + }, JEST_TIMEOUT); + + it('reports a request Cube Store refused as too big instead of retrying it', async () => { + connection = new WebSocketConnection(server.url); + + server.handler = (message, mockConnection) => { + // How Cube Store rejects a message that doesn't fit into its limits. + mockConnection.ws.close(1009, 'Message too big'); + }; + + const promise = query('SELECT 1'); + + await expect(promise).rejects.toThrow(MessageTooLargeError); + await expect(promise).rejects.toThrow( + 'Cube Store closed the connection: message size exceeds the maximum message size Cube Store accepts' + ); + + expect(server.received).toHaveLength(1); + }, JEST_TIMEOUT); + }); + + it('rejects a query when the connection cannot be re-established', async () => { + process.env.CUBEJS_CUBESTORE_MAX_CONNECT_RETRIES = '2'; + + try { + const { url } = server; + await server.stop(); + + connection = new WebSocketConnection(url); + + await expect(query('SELECT 1')).rejects.toThrow('CubeStore connection failed after 2 retries'); + } finally { + delete process.env.CUBEJS_CUBESTORE_MAX_CONNECT_RETRIES; + } + }, JEST_TIMEOUT); +});