Skip to content

Commit a65a047

Browse files
darrachequesnedarrachequesne
committed
fix: wait for all packets to be sent before closing the WebSocket connection
This reverts commit [1], which was included in `engine.io@5.1.0` and `socket.io@4.1.0`. The WebSocket connection was closed before all packets were written out, so for example when calling `socket.disconnect(true)` on the Socket.IO server (which disconnect from all namespaces and close the connection), the client would receive only the first disconnect packet and kept trying to reconnect to the other namespaces. The only difference with the previous implementation (pre 5.1.0) is that the "drain" event gets only called once at the end, and not after each packet. [1]: ad5306a Related: #648
1 parent ed87609 commit a65a047

File tree

4 files changed

+71
-67
lines changed

4 files changed

+71
-67
lines changed

‎lib/socket.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -560,10 +560,18 @@ export class Socket extends EventEmitter {
560560
this.readyState = "closing";
561561

562562
if (this.writeBuffer.length) {
563-
this.once("drain", this.closeTransport.bind(this, discard));
563+
debug(
564+
"there are %d remaining packets in the buffer, waiting for the 'drain' event",
565+
this.writeBuffer.length
566+
);
567+
this.once("drain", () => {
568+
debug("all packets have been sent, closing the transport");
569+
this.closeTransport(discard);
570+
});
564571
return;
565572
}
566573

574+
debug("the buffer is empty, closing the transport right away", discard);
567575
this.closeTransport(discard);
568576
}
569577

@@ -574,6 +582,7 @@ export class Socket extends EventEmitter {
574582
* @api private
575583
*/
576584
private closeTransport(discard) {
585+
debug("closing the transport (discard? %s)", discard);
577586
if (discard) this.transport.discard();
578587
this.transport.close(this.onClose.bind(this, "forced close"));
579588
}

‎lib/transports-uws/websocket.ts‎

Lines changed: 23 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -53,36 +53,32 @@ export class WebSocket extends Transport {
5353
* @api private
5454
*/
5555
send(packets) {
56-
const packet = packets.shift();
57-
if (typeof packet === "undefined") {
58-
this.writable = true;
59-
this.emit("drain");
60-
return;
61-
}
56+
this.writable = false;
6257

63-
// always creates a new object since ws modifies it
64-
const opts: { compress?: boolean } = {};
65-
if (packet.options) {
66-
opts.compress = packet.options.compress;
67-
}
58+
for (let i = 0; i < packets.length; i++) {
59+
const packet = packets[i];
60+
const isLast = i + 1 === packets.length;
6861

69-
const send = (data) => {
70-
const isBinary = typeof data !== "string";
71-
const compress =
72-
this.perMessageDeflate &&
73-
Buffer.byteLength(data) > this.perMessageDeflate.threshold;
62+
const send = (data) => {
63+
const isBinary = typeof data !== "string";
64+
const compress =
65+
this.perMessageDeflate &&
66+
Buffer.byteLength(data) > this.perMessageDeflate.threshold;
7467

75-
debug('writing "%s"', data);
76-
this.writable = false;
68+
debug('writing "%s"', data);
69+
this.socket.send(data, isBinary, compress);
7770

78-
this.socket.send(data, isBinary, compress);
79-
this.send(packets);
80-
};
71+
if (isLast) {
72+
this.writable = true;
73+
this.emit("drain");
74+
}
75+
};
8176

82-
if (packet.options && typeof packet.options.wsPreEncoded === "string") {
83-
send(packet.options.wsPreEncoded);
84-
} else {
85-
this.parser.encodePacket(packet, this.supportsBinary, send);
77+
if (packet.options && typeof packet.options.wsPreEncoded === "string") {
78+
send(packet.options.wsPreEncoded);
79+
} else {
80+
this.parser.encodePacket(packet, this.supportsBinary, send);
81+
}
8682
}
8783
}
8884

@@ -94,7 +90,7 @@ export class WebSocket extends Transport {
9490
doClose(fn) {
9591
debug("closing");
9692
fn && fn();
97-
// call fn first since socket.close() immediately emits a "close" event
98-
this.socket.close();
93+
// call fn first since socket.end() immediately emits a "close" event
94+
this.socket.end();
9995
}
10096
}

‎lib/transports/websocket.ts‎

Lines changed: 37 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -61,47 +61,48 @@ export class WebSocket extends Transport {
6161
* @api private
6262
*/
6363
send(packets) {
64-
const packet = packets.shift();
65-
if (typeof packet === "undefined") {
66-
this.writable = true;
67-
this.emit("drain");
68-
return;
69-
}
64+
this.writable = false;
7065

71-
// always creates a new object since ws modifies it
72-
const opts: { compress?: boolean } = {};
73-
if (packet.options) {
74-
opts.compress = packet.options.compress;
75-
}
66+
for (let i = 0; i < packets.length; i++) {
67+
const packet = packets[i];
68+
const isLast = i + 1 === packets.length;
7669

77-
const send = (data) => {
78-
if (this.perMessageDeflate) {
79-
const len =
80-
"string" === typeof data ? Buffer.byteLength(data) : data.length;
81-
if (len < this.perMessageDeflate.threshold) {
82-
opts.compress = false;
83-
}
70+
// always creates a new object since ws modifies it
71+
const opts: { compress?: boolean } = {};
72+
if (packet.options) {
73+
opts.compress = packet.options.compress;
8474
}
85-
debug('writing "%s"', data);
86-
this.writable = false;
8775

88-
this.socket.send(data, opts, (err) => {
89-
if (err) return this.onError("write error", err.stack);
90-
this.send(packets);
91-
});
92-
};
76+
const onSent = (err) => {
77+
if (err) {
78+
return this.onError("write error", err.stack);
79+
} else if (isLast) {
80+
this.writable = true;
81+
this.emit("drain");
82+
}
83+
};
9384

94-
if (packet.options && typeof packet.options.wsPreEncoded === "string") {
95-
send(packet.options.wsPreEncoded);
96-
} else if (this._canSendPreEncodedFrame(packet)) {
97-
// the WebSocket frame was computed with WebSocket.Sender.frame()
98-
// see https://github.com/websockets/ws/issues/617#issuecomment-283002469
99-
this.socket._sender.sendFrame(packet.options.wsPreEncodedFrame, (err) => {
100-
if (err) return this.onError("write error", err.stack);
101-
this.send(packets);
102-
});
103-
} else {
104-
this.parser.encodePacket(packet, this.supportsBinary, send);
85+
const send = (data) => {
86+
if (this.perMessageDeflate) {
87+
const len =
88+
"string" === typeof data ? Buffer.byteLength(data) : data.length;
89+
if (len < this.perMessageDeflate.threshold) {
90+
opts.compress = false;
91+
}
92+
}
93+
debug('writing "%s"', data);
94+
this.socket.send(data, opts, onSent);
95+
};
96+
97+
if (packet.options && typeof packet.options.wsPreEncoded === "string") {
98+
send(packet.options.wsPreEncoded);
99+
} else if (this._canSendPreEncodedFrame(packet)) {
100+
// the WebSocket frame was computed with WebSocket.Sender.frame()
101+
// see https://github.com/websockets/ws/issues/617#issuecomment-283002469
102+
this.socket._sender.sendFrame(packet.options.wsPreEncodedFrame, onSent);
103+
} else {
104+
this.parser.encodePacket(packet, this.supportsBinary, send);
105+
}
105106
}
106107
}
107108

‎test/server.js‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1797,9 +1797,7 @@ describe("server", () => {
17971797
conn.send("a");
17981798
conn.send("b");
17991799
conn.send("c");
1800-
setTimeout(() => {
1801-
conn.close();
1802-
}, 50);
1800+
conn.close();
18031801
});
18041802

18051803
socket.on("open", () => {

0 commit comments

Comments
 (0)