how to gracefully shutdown websocket clients using tokio - rust

I am implementing a websocket chat application where I want to gracefully shutdown all clients when the server stop because of the ctrl+c signal.
I am listening to incoming events using mio poll and tokens. Any new socket connection is registered with mio poll and any event received on the socket is successfully captured on the polling.
My initial idea was to use tokio::select with listen events and shutdown as 2 branches. But, I guess this design needs some modification to enable graceful shutdown.
// get_events function
fn get_events(poll, conn) -> Option<mio::Event>{
let shutdown = tokio::signal::ctrl_c();
let res = tokio::select!{
res = get_poll_events(poll) => { // returns events asynchronously, its an async method
// ... handler when events are received
res
},
_ = shutdown => {
// ... handler when ctrl+c signal is received
// send close connection message to TcpStream
return None;
}
};
Some(res)
}
// main function
let poll = Poll::new();
let shared_poll = Arc::new(Mutex::new(poll));
let conn: Arc<Mutex<HashMap<Token, WebSocketClient>>> = Arc::new(Mutex::new(HashMap::new()));
let (shutdown_notifier, shutdown_receiver) = mpsc::channel(10);
loop {
let res = WebSocketServer::get_events(shared_poll.clone(), conn.clone()).await;
if let None = res {
drop(shutdown_notifier); // drop the original mpsc::Sender, cloned in all clients.
let _ = shutdown_receiver.recv().await; // mpsc::Receiver only returns error when all clients are dropped, thus dropping the senders inside them.
break; // stop the program
}
// use event to register a new client with mpsc::sender
// if readable; accept incoming message
// broadcast message to other subscribers
// ... other tasks
...
}
On Ctrl+c signal received the shutdown handler gets activated. Post that, close connection message is sent on the TcpStream to all active clients. At this point, everything is as expected. All the clients receive close messages. But the real issue starts here, since the shutdown handler executes, the get_poll_events branch of tokio::select gets dropped. Meaning, no more polling for the incoming events on the registered sources (TcpStreams).
Ideally, the server should only close when it has received back the close message response from all the clients. The clients communicate back on the TcpStream, but since there is no active mechanism to listen to these events, I am unable to capture them and hence not able to drop the clients. However, instead of sending close message to clients if I dropped all the clients manually, then there was no need to listen to close message response and things worked albeit an incorrect implementation.
Tokio::select isn't an ideal choice as far as I can guess, but I am unable to come up with a solution on how to implement this case, where the server is still active to listen to all the close message responses from clients. It can then close when all the clients have gone out of scope in the received close messages.
What would be a way to achieve this functionality? TIA.

Related

Properly 2 TcpStream

I am working on a TCP proxy server where a client connects to it and it starts a new connection to a backend and forward all packets (bidirectional). The function that handles this part looks like this.
pub fn route(source: TcpStream, worker: Worker) -> Result<()> {
// `source` and `destination` are both `TcpStream`
let mut source = source;
let mut destination = TcpStream::connect(worker.address)?;
let mut source_copy = source.try_clone()?;
let mut destination_copy = destination.try_clone()?;
let src2dst = std::thread::spawn(move || {
std::io::copy(&mut source_copy, &mut destination_copy).unwrap();
});
let dst2src = std::thread::spawn(move || {
std::io::copy(&mut destination, &mut source).unwrap();
// source.shutdown(Shutdown::Both).unwrap(); <--- if this line is commented, it will stuck
});
src2dst.join().unwrap();
dst2src.join().unwrap();
Ok(())
}
However, in its current form, this function will stuck. In particular, if I shutdown the source when destination stops writing, it won't block. But I am still not sure why this works (or why it does not). I am currently using it only to proxy HTTP traffic, and it does not appear to have problems. But I am not sure if it works for generic TCP. What is the proper way to do this?
std::io::copy reads from the reader until it gets an end-of-file condition. The way to signal an end-of-file condition on a TCP connection is to send a FIN packet, which is realized in software by calling shutdown to shut down the write half of the connection.
HTTP allows reusing a single connection to send multiple requests, one after the other. There is enough information in a request for the recipient to determine where the request ends, and likewise for responses. In your case, it appears the backend shuts down its write half of the connection, which causes your io::copy to return. If the HTTP response has Connection: close in its headers, then you must shut down the write half, or else the client will hang waiting for the response body because it never receives a FIN packet.
TL;DR: When the copy from A to B returns normally, you should shut down the write half of B to forward the end-of-file condition (FIN packet). This goes for the copies in both directions, and is valid for TCP in general. Shutting down the read half of sockets doesn't appear to be necessary (it doesn't generate any TCP packets).

How to share a TCP socket object between parent and (forked) child?

I have an application in which my TCP server module (parent) listens for 'connection' events and receives some data on the created socket to perform a handshake with the remote client. Once the handshake is performed, the server needs to send the socket object to a forked child, which will also send and receive data to the socket, do some stuff and finally send result to parent and be killed. For some reasons, I need to keep the socket object in the parent for further data processing not performed in the child, after the child has finished.
I've managed to send the socket to the child using the subprocess.send() method but, this way, the socket handle becomes null in the parent. I tried setting the keepOpen option to true and it almost worked, since I can send the socket and still work with it in the parent, but It seems not to work properly, because incoming data is not always received by the child 'data' event listener.
I also tried to removeListener for the 'data' event from the parent, prior to sending the socket to the child, but this made no difference, data is still being lost at some point on some occasions (on some others it is correctly received after an unexpected delay...). This code extract illustrates what I'm trying to do:
const net = require('net');
const server = net.createServer();
const cp = require('child_process');
server.on('connection', (socket) => {
socket.on('data', (data) => {
// Perform handhsake
const child = cp.fork('child.js');
child.on('message', (result) => {
console.log('CHILD finished processing: ', result);
child.kill('SIGHUP');
// Do more stuff with socket
});
child.send('socket', socket);
// (At this point, socket handle is null)
});
});
server.listen(PORT)
I'm new to nodejs, I assume there might be errors in the code. Thanks.

socket.io how to send multiple messages sequentially?

I'm using socket.io like this
Client:
socket.on('response', function(i){
console.log(i);
});
socket.emit('request', whateverdata);
Server:
socket.on('request', function(whateverdata){
for (i=0; i<10000; i++){
console.log(i);
socket.emit('response', i);
}
console.log("done!");
});
I need output like this when putting the two terminals side by side:
Server Client
0 0
1 1
. (etc) .
. .
9998 9998
9999 9999
done!
But instead I am getting this:
Server Client
0
1
. (etc)
.
9998
9999
done!
0
1
.
. (etc)
9998
9999
Why?
Shouldn't Socket.IO / Node emit the message immediately, not wait for the loop to complete before emitting any of them?
Notes:
The for loop is very long and computationally slow.
This question is referring to the socket.io library, not websockets in general.
Due to latency, waiting for confirmation from the client before sending each response is not possible
The order that the messages are received is not important, only that they are received as quickly as possible
The server emits them all in a loop and it takes a small bit of time for them to get to the client and get processed by the client in another process. This should not be surprising.
It is also possible that the single-threaded nature of Javascript in node.js prevents the emits from actually getting sent until your Javascript loop finishes. That would take detailed examination of socket.io code to know for sure if that is an issue. As I said before if you want to 1,1 then 2,2 then 3,3 instead of 1,2,3 sent, then 1,2,3 received you have to write code to force that.
If you want the client to receive the first before the server sends the 2nd, then you have to make the client send a response to the first and have the server not send the 2nd until it receives the response from the first. This is all async networking. You don't control the order of events in different processes unless you write specific code to force a particular sequence.
Also, how do you have client and server in the same console anyway? Unless you are writing out precise timestamps, you wouldn't be able to tell exactly what event came before the other in two separate processes.
One thing you could try is to send 10, then do a setTimeout(fn, 1) to send the next 10 and so on. That would give JS a chance to breathe and perhaps process some other events that are waiting for you to finish to allow the packets to get sent.
There's another networking issue too. By default TCP tries to batch up your sends (at the lowest TCP level). Each time you send, it sets a short timer and doesn't actually send until that timer fires. If more data arrives before the timer fires, it just adds that data to the "pending" packet and sets the timer again. This is referred to as the Nagle's algorithm. You can disable this "feature" on a per-socket basis with socket.setNoDelay(). You have to call that on the actual TCP socket.
I am seeing some discussion that Nagle's algorithm may already be turned off for socket.io (by default). Not sure yet.
In stepping through the process of socket.io's .emit(), there are some cases where the socket is marked as not yet writable. In those cases, the packets are added to a buffer and will be processed "later" on some future tick of the event loop. I cannot see exactly what puts the socket temporarily in this state, but I've definitely seen it happen in the debugger. When it's that way, a tight loop of .emit() will just buffer and won't send until you let other events in the event loop process. This is why doing setTimeout(fn, 0) every so often to keep sending will then let the prior packets process. There's some other event that needs to get processed before socket.io makes the socket writable again.
The issue occurs in the flush() method in engine.io (the transport layer for socket.io). Here's the code for .flush():
Socket.prototype.flush = function () {
if ('closed' !== this.readyState &&
this.transport.writable &&
this.writeBuffer.length) {
debug('flushing buffer to transport');
this.emit('flush', this.writeBuffer);
this.server.emit('flush', this, this.writeBuffer);
var wbuf = this.writeBuffer;
this.writeBuffer = [];
if (!this.transport.supportsFraming) {
this.sentCallbackFn.push(this.packetsFn);
} else {
this.sentCallbackFn.push.apply(this.sentCallbackFn, this.packetsFn);
}
this.packetsFn = [];
this.transport.send(wbuf);
this.emit('drain');
this.server.emit('drain', this);
}
};
What happens sometimes is that this.transport.writable is false. And, when that happens, it does not send the data yet. It will be sent on some future tick of the event loop.
From what I can tell, it looks like the issue may be here in the WebSocket code:
WebSocket.prototype.send = function (packets) {
var self = this;
for (var i = 0; i < packets.length; i++) {
var packet = packets[i];
parser.encodePacket(packet, self.supportsBinary, send);
}
function send (data) {
debug('writing "%s"', data);
// always creates a new object since ws modifies it
var opts = {};
if (packet.options) {
opts.compress = packet.options.compress;
}
if (self.perMessageDeflate) {
var len = 'string' === typeof data ? Buffer.byteLength(data) : data.length;
if (len < self.perMessageDeflate.threshold) {
opts.compress = false;
}
}
self.writable = false;
self.socket.send(data, opts, onEnd);
}
function onEnd (err) {
if (err) return self.onError('write error', err.stack);
self.writable = true;
self.emit('drain');
}
};
Where you can see that the .writable property is set to false when some data is sent until it gets confirmation that the data has been written. So, when rapidly sending data in a loop, it may not be letting the event come through that signals that the data has been successfully sent. When you do a setTimeout() to let some things in the event loop get processed that confirmation event comes through and the .writable property gets set to true again so data can again be sent immediately.
To be honest, socket.io is built of so many abstract layers across dozens of modules that it's very difficult code to debug or analyze on GitHub so it's hard to be sure of the exact explanation. I did definitely see the .writable flag as false in the debugger which did cause a delay so this seems like a plausible explanation to me. I hope this helps.

Handle new TCP connections synchronously

I know nodejs is asynchronous by nature and it is preferable to use that way, but I have a use case where we need to handle incoming TCP connections in synchronous way. Once a new connections received we need to connect to some other TCP server and perform some book keeping stuff etc and then handle some other connection. Since number of connections are limited, it is fine to handle this in synchronous way.
Looking for an elegant way to handle this scenario.
net.createServer(function(sock) {
console.log('Received a connection - ');
var sock = null;
var testvar = null;
sock = new net.Socket();
sock.connect(PORT, HOST, function() {
console.log('Connected to server - ');
});
//Other listeners
}
In the above code if two connections received simultaneously the output may be (since asynchronous nature):
Received a connection
Receive a connection
Connected to server
Connected to server
But the expectation is:
Received a connection
Connected to server
Receive a connection
Connected to server
What is the proper way of ding this?
One solution is implement a queue kind of solution with emitting 'done' or 'complete' events to handle next connection.
For this we may have to take the connection callback out of the createServer call. How to handle scoping of connection and other variables (testvar) in this case?
In this case what happens to the data/messages if received on connections which are in queue but not yet processed and not yet 'data' listener is registered.?
Any other better solutions will be helpful.
I think it is important to separate the concepts of synchronous code vs serial code. You want to process each request serially, but that can still be accomplished while handling each request asynchronously. For your case, the easiest way would probably be to have a queue of requests to handle instead.
var inProgress = false;
var queue = [];
net.createServer(function(sock){
queue.push(sock);
processQueue();
});
function processQueue(){
if (inProgress || queue.length === 0) return;
inProgress = true;
handleSockSerial(queue.shift(), function(){
inProgress = false;
processQueue();
});
}
function handleSockSerial(sock, callback){
// Do all your stuff and then call 'callback' when you are done.
}
Note, as long as you are using node >= 0.10, the data coming in from the socket will be buffered until you read the data.

nodeJS zeroMQ: why cant send message if sock.send not in setInterval

I was using zeroMQ in nodeJS. But it seems that while sending the data from producer to worker, if I do not put it in setInterval, then it does not send the data to the worker. My example code is as follows:
producer.js
===========
var zmq = require('zmq')
, sock = zmq.socket('push');
sock.bindSync('tcp://127.0.0.1:3000');
console.log('Producer bound to port 3000');
//sock.send("hello");
var i = 0;
//1. var timer = setInterval(function() {
var str = "hello";
console.log('sending work', str, i++);
sock.send(str);
//2. clearTimeout(timer);
//3. }, 150);
sock.on('message', function(msg) {
console.log("Got A message, [%s], [%s]", msg);
});
So in the above code, if I add back the lines commented in 1, 2 and 3, then I do receive the message to the worker side, else it does not work.
Can anyone throw light why to send message I need to put it in setInterval? Or am I doing something wrong way?
The problem is hidden in the zmq bindings for node.js . I've just spent some time digging into it and it basically does this on send():
Enqueue the message
Flush buffers
Now the problem is in the flushing part, because it does
Check if the output socket is ready, otherwise return
Flush the enqueued messages
In your code, because you call bind and immediately send, there is no worker connected at the moment of the call, because they simply didn't have enough time to notice. So the message is enqueued and we are waiting for some workers to appear. Now the interesting part - where do we check for new workers? In the send function itself! So unless we call send() later, when there are actually some workers connected, our messages are never flushed and they are enqueued forever. And that is why setInterval works, because workers have enough time to notice and connect and you periodically check if there are any.
You can find the interesting part at https://github.com/JustinTulloss/zeromq.node/blob/master/lib/index.js#L277 .
Cheers ;-)

Resources