Skip to content

Commit 4b9ad77

Browse files
durrandurran
authored
feat(NODE-4139): streaming protocol message changes (#3256)
1 parent c496c25 commit 4b9ad77

File tree

6 files changed

+183
-106
lines changed

6 files changed

+183
-106
lines changed

‎src/cmap/connection.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,15 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
293293
this[kHello] = response;
294294
}
295295

296+
// Set the whether the message stream is for a monitoring connection.
297+
set isMonitoringConnection(value: boolean) {
298+
this[kMessageStream].isMonitoringConnection = value;
299+
}
300+
301+
get isMonitoringConnection(): boolean {
302+
return this[kMessageStream].isMonitoringConnection;
303+
}
304+
296305
get serviceId(): ObjectId | undefined {
297306
return this.hello?.serviceId;
298307
}

‎src/cmap/message_stream.ts

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,13 +53,19 @@ export class MessageStream extends Duplex {
5353
maxBsonMessageSize: number;
5454
/** @internal */
5555
[kBuffer]: BufferPool;
56+
/** @internal */
57+
isMonitoringConnection = false;
5658

5759
constructor(options: MessageStreamOptions = {}) {
5860
super(options);
5961
this.maxBsonMessageSize = options.maxBsonMessageSize || kDefaultMaxBsonMessageSize;
6062
this[kBuffer] = new BufferPool();
6163
}
6264

65+
get buffer(): BufferPool {
66+
return this[kBuffer];
67+
}
68+
6369
override _write(chunk: Buffer, _: unknown, callback: Callback<Buffer>): void {
6470
this[kBuffer].append(chunk);
6571
processIncomingData(this, callback);
@@ -162,15 +168,36 @@ function processIncomingData(stream: MessageStream, callback: Callback<Buffer>)
162168
opCode: message.readInt32LE(12)
163169
};
164170

171+
const monitorHasAnotherHello = () => {
172+
if (stream.isMonitoringConnection) {
173+
// Can we read the next message size?
174+
if (buffer.length >= 4) {
175+
const sizeOfMessage = buffer.peek(4).readInt32LE();
176+
if (sizeOfMessage <= buffer.length) {
177+
return true;
178+
}
179+
}
180+
}
181+
return false;
182+
};
183+
165184
let ResponseType = messageHeader.opCode === OP_MSG ? BinMsg : Response;
166185
if (messageHeader.opCode !== OP_COMPRESSED) {
167186
const messageBody = message.slice(MESSAGE_HEADER_SIZE);
168-
stream.emit('message', new ResponseType(message, messageHeader, messageBody));
169187

170-
if (buffer.length >= 4) {
188+
// If we are a monitoring connection message stream and
189+
// there is more in the buffer that can be read, skip processing since we
190+
// want the last hello command response that is in the buffer.
191+
if (monitorHasAnotherHello()) {
171192
processIncomingData(stream, callback);
172193
} else {
173-
callback();
194+
stream.emit('message', new ResponseType(message, messageHeader, messageBody));
195+
196+
if (buffer.length >= 4) {
197+
processIncomingData(stream, callback);
198+
} else {
199+
callback();
200+
}
174201
}
175202

176203
return;
@@ -198,12 +225,19 @@ function processIncomingData(stream: MessageStream, callback: Callback<Buffer>)
198225
return;
199226
}
200227

201-
stream.emit('message', new ResponseType(message, messageHeader, messageBody));
202-
203-
if (buffer.length >= 4) {
228+
// If we are a monitoring connection message stream and
229+
// there is more in the buffer that can be read, skip processing since we
230+
// want the last hello command response that is in the buffer.
231+
if (monitorHasAnotherHello()) {
204232
processIncomingData(stream, callback);
205233
} else {
206-
callback();
234+
stream.emit('message', new ResponseType(message, messageHeader, messageBody));
235+
236+
if (buffer.length >= 4) {
237+
processIncomingData(stream, callback);
238+
} else {
239+
callback();
240+
}
207241
}
208242
});
209243
}

‎src/sdam/monitor.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ export class Monitor extends TypedEventEmitter<MonitorEvents> {
8888
[kMonitorId]?: InterruptibleAsyncInterval;
8989
[kRTTPinger]?: RTTPinger;
9090

91+
get connection(): Connection | undefined {
92+
return this[kConnection];
93+
}
94+
9195
constructor(server: Server, options: MonitorOptions) {
9296
super();
9397

@@ -310,6 +314,10 @@ function checkServer(monitor: Monitor, callback: Callback<Document | null>) {
310314
}
311315

312316
if (conn) {
317+
// Tell the connection that we are using the streaming protocol so that the
318+
// connection's message stream will only read the last hello on the buffer.
319+
conn.isMonitoringConnection = true;
320+
313321
if (isInCloseState(monitor)) {
314322
conn.destroy({ force: true });
315323
return;

‎test/tools/utils.ts

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import { EJSON } from 'bson';
2+
import * as BSON from 'bson';
23
import { expect } from 'chai';
34
import { inspect, promisify } from 'util';
45

6+
import { OP_MSG } from '../../src/cmap/wire_protocol/constants';
7+
import { Document } from '../../src/index';
58
import { Logger } from '../../src/logger';
69
import { deprecateOptions, DeprecateOptionsConfig } from '../../src/utils';
710
import { runUnifiedSuite } from './unified-spec-runner/runner';
@@ -343,6 +346,24 @@ export class TestBuilder {
343346
}
344347
}
345348

349+
export function generateOpMsgBuffer(document: Document): Buffer {
350+
const header = Buffer.alloc(4 * 4 + 4);
351+
352+
const typeBuffer = Buffer.alloc(1);
353+
typeBuffer[0] = 0;
354+
355+
const docBuffer = BSON.serialize(document);
356+
357+
const totalLength = header.length + typeBuffer.length + docBuffer.length;
358+
359+
header.writeInt32LE(totalLength, 0);
360+
header.writeInt32LE(0, 4);
361+
header.writeInt32LE(0, 8);
362+
header.writeInt32LE(OP_MSG, 12);
363+
header.writeUInt32LE(0, 16);
364+
return Buffer.concat([header, typeBuffer, docBuffer]);
365+
}
366+
346367
export class UnifiedTestSuiteBuilder {
347368
private _description = 'Default Description';
348369
private _schemaVersion = '1.0';

‎test/unit/cmap/message_stream.test.js

Lines changed: 100 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
'use strict';
2-
const Readable = require('stream').Readable;
3-
const Writable = require('stream').Writable;
2+
const { on, once } = require('events');
3+
const { Readable, Writable } = require('stream');
4+
45
const { MessageStream } = require('../../../src/cmap/message_stream');
56
const { Msg } = require('../../../src/cmap/commands');
67
const expect = require('chai').expect;
78
const { LEGACY_HELLO_COMMAND } = require('../../../src/constants');
9+
const { generateOpMsgBuffer } = require('../../tools/utils');
810

911
function bufferToStream(buffer) {
1012
const stream = new Readable();
@@ -18,117 +20,117 @@ function bufferToStream(buffer) {
1820
return stream;
1921
}
2022

21-
describe('Message Stream', function () {
22-
describe('reading', function () {
23-
[
24-
{
25-
description: 'valid OP_REPLY',
26-
data: Buffer.from(
27-
'370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
28-
'hex'
29-
),
30-
documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
31-
},
32-
{
33-
description: 'valid multiple OP_REPLY',
34-
expectedMessageCount: 4,
35-
data: Buffer.from(
36-
'370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
37-
'370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
38-
'370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
39-
'370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
40-
'hex'
41-
),
42-
documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
43-
},
44-
{
45-
description: 'valid OP_REPLY (partial)',
46-
data: [
47-
Buffer.from('37', 'hex'),
48-
Buffer.from('0000', 'hex'),
49-
Buffer.from(
50-
'000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
51-
'hex'
52-
)
53-
],
54-
documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
55-
},
56-
57-
{
58-
description: 'valid OP_MSG',
59-
data: Buffer.from(
60-
'370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000',
61-
'hex'
62-
),
63-
documents: [{ $db: 'admin', [LEGACY_HELLO_COMMAND]: 1 }]
64-
},
65-
{
66-
description: 'valid multiple OP_MSG',
67-
expectedMessageCount: 4,
68-
data: Buffer.from(
69-
'370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
70-
'370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
71-
'370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
72-
'370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000',
73-
'hex'
74-
),
75-
documents: [{ $db: 'admin', [LEGACY_HELLO_COMMAND]: 1 }]
76-
},
77-
78-
{
79-
description: 'Invalid message size (negative)',
80-
data: Buffer.from('ffffffff', 'hex'),
81-
error: 'Invalid message size: -1'
82-
},
83-
{
84-
description: 'Invalid message size (exceeds maximum)',
85-
data: Buffer.from('01000004', 'hex'),
86-
error: 'Invalid message size: 67108865, max allowed: 67108864'
87-
}
88-
].forEach(test => {
89-
it(test.description, function (done) {
90-
const error = test.error;
91-
const expectedMessageCount = test.expectedMessageCount || 1;
92-
const inputStream = bufferToStream(test.data);
93-
const messageStream = new MessageStream();
23+
describe('MessageStream', function () {
24+
context('when the stream is for a monitoring connection', function () {
25+
const response = { isWritablePrimary: true };
26+
const lastResponse = { ok: 1 };
27+
let firstHello;
28+
let secondHello;
29+
let thirdHello;
30+
let partial;
31+
32+
beforeEach(function () {
33+
firstHello = generateOpMsgBuffer(response);
34+
secondHello = generateOpMsgBuffer(response);
35+
thirdHello = generateOpMsgBuffer(lastResponse);
36+
partial = Buffer.alloc(5);
37+
partial.writeInt32LE(100, 0);
38+
});
9439

95-
let messageCount = 0;
96-
messageStream.on('message', msg => {
97-
messageCount++;
98-
if (error) {
99-
done(new Error(`expected error: ${error}`));
100-
return;
101-
}
40+
it('only reads the last message in the buffer', async function () {
41+
const inputStream = bufferToStream(Buffer.concat([firstHello, secondHello, thirdHello]));
42+
const messageStream = new MessageStream();
43+
messageStream.isMonitoringConnection = true;
44+
45+
inputStream.pipe(messageStream);
46+
const messages = await once(messageStream, 'message');
47+
const msg = messages[0];
48+
msg.parse();
49+
expect(msg).to.have.property('documents').that.deep.equals([lastResponse]);
50+
// Make sure there is nothing left in the buffer.
51+
expect(messageStream.buffer.length).to.equal(0);
52+
});
10253

103-
msg.parse();
54+
it('does not read partial messages', async function () {
55+
const inputStream = bufferToStream(
56+
Buffer.concat([firstHello, secondHello, thirdHello, partial])
57+
);
58+
const messageStream = new MessageStream();
59+
messageStream.isMonitoringConnection = true;
60+
61+
inputStream.pipe(messageStream);
62+
const messages = await once(messageStream, 'message');
63+
const msg = messages[0];
64+
msg.parse();
65+
expect(msg).to.have.property('documents').that.deep.equals([lastResponse]);
66+
// Make sure the buffer wasn't read to the end.
67+
expect(messageStream.buffer.length).to.equal(5);
68+
});
69+
});
10470

105-
if (test.documents) {
106-
expect(msg).to.have.property('documents').that.deep.equals(test.documents);
107-
}
71+
context('when the stream is not for a monitoring connection', function () {
72+
context('when the messages are valid', function () {
73+
const response = { isWritablePrimary: true };
74+
let firstHello;
75+
let secondHello;
76+
let thirdHello;
77+
let messageCount = 0;
78+
79+
beforeEach(function () {
80+
firstHello = generateOpMsgBuffer(response);
81+
secondHello = generateOpMsgBuffer(response);
82+
thirdHello = generateOpMsgBuffer(response);
83+
});
10884

109-
if (messageCount === expectedMessageCount) {
110-
done();
111-
}
112-
});
85+
it('reads all messages in the buffer', async function () {
86+
const inputStream = bufferToStream(Buffer.concat([firstHello, secondHello, thirdHello]));
87+
const messageStream = new MessageStream();
11388

114-
messageStream.on('error', err => {
115-
if (error == null) {
116-
done(err);
89+
inputStream.pipe(messageStream);
90+
for await (const messages of on(messageStream, 'message')) {
91+
messageCount++;
92+
const msg = messages[0];
93+
msg.parse();
94+
expect(msg).to.have.property('documents').that.deep.equals([response]);
95+
// Test will not complete until 3 messages processed.
96+
if (messageCount === 3) {
11797
return;
11898
}
99+
}
100+
});
101+
});
119102

120-
expect(err).to.have.property('message').that.equals(error);
103+
context('when the messages are invalid', function () {
104+
context('when the message size is negative', function () {
105+
it('emits an error', async function () {
106+
const inputStream = bufferToStream(Buffer.from('ffffffff', 'hex'));
107+
const messageStream = new MessageStream();
121108

122-
done();
109+
inputStream.pipe(messageStream);
110+
const errors = await once(messageStream, 'error');
111+
const err = errors[0];
112+
expect(err).to.have.property('message').that.equals('Invalid message size: -1');
123113
});
114+
});
124115

125-
inputStream.pipe(messageStream);
116+
context('when the message size exceeds the bson maximum', function () {
117+
it('emits an error', async function () {
118+
const inputStream = bufferToStream(Buffer.from('01000004', 'hex'));
119+
const messageStream = new MessageStream();
120+
121+
inputStream.pipe(messageStream);
122+
const errors = await once(messageStream, 'error');
123+
const err = errors[0];
124+
expect(err)
125+
.to.have.property('message')
126+
.that.equals('Invalid message size: 67108865, max allowed: 67108864');
127+
});
126128
});
127129
});
128130
});
129131

130-
describe('writing', function () {
131-
it('should write a message to the stream', function (done) {
132+
context('when writing to the message stream', function () {
133+
it('pushes the message', function (done) {
132134
const readableStream = new Readable({ read() {} });
133135
const writeableStream = new Writable({
134136
write: (chunk, _, callback) => {

‎test/unit/sdam/monitor.test.js

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,10 @@ describe('monitoring', function () {
115115
monitor = new Monitor(server, {});
116116

117117
monitor.on('serverHeartbeatFailed', () => done(new Error('unexpected heartbeat failure')));
118-
monitor.on('serverHeartbeatSucceeded', () => done());
118+
monitor.on('serverHeartbeatSucceeded', () => {
119+
expect(monitor.connection.isMonitoringConnection).to.be.true;
120+
done();
121+
});
119122
monitor.connect();
120123
});
121124

0 commit comments

Comments
 (0)