Merge pull request #946 from immobiliare/feat-flush-udp

fix: ensure udp messages are sent on close

Author: dnlup

src/socket.ts

@@ -3,25 +3,34 @@ import net, { Socket as SocketTCP, isIP } from 'net';
 import { URL } from 'url';
 import buildLookupFunction from './dns-cache';
 import { debuglog, DebugLoggerFunction } from 'util';
-export abstract class Socket {
+import EventEmitter, { once } from 'events';
+
+/**
+ * @emits idle The socket has no more pending messages. We use this event
+ * to be sure the socket is not active anymore before closing it.
+ */
+export abstract class Socket extends EventEmitter {
     protected hostname: string;
     protected port: number;
     protected connected: boolean;
     protected debug: typeof console.log;
     protected onError: (error: Error) => void;
+    protected _pendingMessages: number;
 
     protected constructor(
         url: URL,
         onError: (error: Error) => void = () => undefined,
         debug: DebugLoggerFunction = debuglog('dats')
     ) {
+        super();
         if (!url.port) {
             throw new Error('A port is required');
         }
         if (!url.hostname) {
             throw new Error('The hostname is required');
         }
 
+        this._pendingMessages = 0;
         this.onError = onError.bind(null);
         this.debug = debug;
         this.hostname = url.hostname;
@@ -33,6 +42,20 @@ export abstract class Socket {
         return this.connected;
     }
 
+    /**
+     * Number of messages that are still being sent.
+     */
+    get pendingMessages() {
+        return this._pendingMessages;
+    }
+    /**
+     * Whether the socket has pending messages or not.
+     * Used to check if the socket is still active.
+     */
+    get idle() {
+        return this._pendingMessages === 0;
+    }
+
     abstract connect(): Promise<boolean>;
 
     abstract send(data: string): void;
@@ -102,17 +125,34 @@ export class SocketTcp extends Socket {
 
     send(data: string): void {
         if (!this.connected || !data) return;
-        this.socket.write(data + '\n', (error) => error && this.onError(error));
+        this._pendingMessages += 1;
+        this.socket.write(data + '\n', (err) => {
+            if (this._pendingMessages) {
+                this._pendingMessages -= 1;
+            }
+            if (err) {
+                try {
+                    this.onError(err);
+                } catch (e) {
+                    this.debug('Exception on this.onError function', e);
+                }
+            }
+            if (this._pendingMessages === 0) {
+                this.emit('idle');
+            }
+        });
     }
 
-    close(): Promise<void> {
+    async close(): Promise<void> {
         if (this.closing) return Promise.resolve();
+        if (!this.idle) {
+            await once(this, 'idle');
+        }
         this.connected = false;
         this.closing = true;
         this.socket.removeListener('close', this.reconnectCb);
         this.socket.end();
         this.socket.destroy();
-        return Promise.resolve();
     }
 }
 
@@ -147,7 +187,6 @@ export class SocketUdp extends Socket {
     ) {
         super(url, onError, debug);
         this.socket = null;
-
         // Removed parenthesis if host name is ipv6 IP.
         if (this.hostname.startsWith('[') && this.hostname.endsWith(']')) {
             /* istanbul ignore next */
@@ -181,19 +220,31 @@ export class SocketUdp extends Socket {
 
     send(data: string): void {
         if (!this.connected || !data) return;
-
-        return this.socket.send(
-            data,
-            this.port,
-            this.hostname,
-            (err) => err && this.onError(err)
-        );
+        this._pendingMessages += 1;
+        this.socket.send(data, this.port, this.hostname, (err) => {
+            if (this._pendingMessages) {
+                this._pendingMessages -= 1;
+            }
+            if (err) {
+                try {
+                    this.onError(err);
+                } catch (e) {
+                    this.debug('Exception on this.onError function', e);
+                }
+            }
+            if (this._pendingMessages === 0) {
+                this.emit('idle');
+            }
+        });
     }
 
-    close(): Promise<void> {
-        if (!this.connected) return Promise.resolve();
-        return new Promise((res) => {
-            this.socket.close(res);
+    async close(): Promise<void> {
+        if (!this.connected) return;
+        if (!this.idle) {
+            await once(this, 'idle');
+        }
+        await new Promise((res) => {
+            this.socket.close(res as () => void);
             this.connected = false;
         });
     }

test/index.spec.ts

@@ -398,6 +398,28 @@ test('close with promise multiple times', async (t) => {
     t.pass();
 });
 
+test('close with queued messages should wait', async (t) => {
+    const host = new URL(`udp://127.0.0.1:${t.context.address.port}`);
+    const namespace = 'ns1.${hostname}.${pid}';
+    const client = new Client({ host, namespace });
+    t.plan(6);
+    client.connect();
+    client.counter('some');
+    client.counter('some');
+    client.counter('some');
+    client.counter('some');
+    //@ts-expect-error Just for tests
+    t.is(client.socket.pendingMessages, 4);
+    //@ts-expect-error Just for tests
+    const emit = sinon.spy(client.socket, 'emit');
+    await t.notThrowsAsync(client.close() as Promise<void>);
+    //@ts-expect-error Just for tests
+    t.is(client.socket.pendingMessages, 0);
+    t.notThrows(() => client.timing('time', 1));
+    t.true(emit.calledOnceWith('idle'));
+    t.pass();
+});
+
 test('getSupportedTypes test', (t) => {
     const host = new URL(`udp://127.0.0.1:${t.context.address.port}`);
     const client = new Client({

test/tcp.spec.ts

@@ -1,4 +1,3 @@
-import os from 'os';
 import sinon from 'sinon';
 import anyTest, { TestInterface } from 'ava';
 import Client from '../src/index';
@@ -110,6 +109,28 @@ test('TCP does not send if connection is closed', async (t) => {
     const sendFun = sinon.spy((socket as any).socket, 'write');
     socket.send('something');
     t.true(sendFun.notCalled);
+    t.is(socket.pendingMessages, 0);
+    t.is(socket.idle, true);
+});
+
+test('TCP close should wait', async (t) => {
+    const host = new URL(
+        `tcp://127.0.0.1:${(t.context.addressTcp as AddressInfo).port || 0}`
+    );
+    const socket = new SocketTcp(host);
+    await socket.connect();
+    socket.send('something');
+    socket.send('something');
+    socket.send('something');
+    socket.send('something');
+    t.is(socket.pendingMessages, 4);
+    t.false(socket.idle);
+    const emit = sinon.spy(socket, 'emit');
+    await socket.close();
+    socket.send('something');
+    t.is(socket.pendingMessages, 0);
+    t.is(socket.idle, true);
+    t.true(emit.calledOnceWith('idle'));
 });
 
 test('TCP _connect function doesnt create new connections if there is another connection active', async (t) => {