fix(express): listen for finish event on response for async express…

… layer open-telemetry#107

plugins/node/opentelemetry-plugin-express/src/express.ts

@@ -201,6 +201,13 @@ export class ExpressPlugin extends BasePlugin<typeof express> {
           span.end(startTime);
           spanHasEnded = true;
         }
+        // listener for response.on('finish')
+        const onResponseFinish = () => {
+          if (spanHasEnded === false) {
+            spanHasEnded = true;
+            span.end(startTime);
+          }
+        };
         // verify we have a callback
         const args = Array.from(arguments);
         const callbackIdx = args.findIndex(arg => typeof arg === 'function');
@@ -209,6 +216,7 @@ export class ExpressPlugin extends BasePlugin<typeof express> {
             if (spanHasEnded === false) {
               span.end();
               spanHasEnded = true;
+              req.res?.removeListener('finish', onResponseFinish);
             }
             if (!(req.route && arguments[0] instanceof Error)) {
               (req[_LAYERS_STORE_PROPERTY] as string[]).pop();
@@ -218,12 +226,13 @@ export class ExpressPlugin extends BasePlugin<typeof express> {
           };
         }
         const result = original.apply(this, arguments);
-        // If the callback is never called, we need to close the span.
-        setImmediate(() => {
-          if (spanHasEnded === false) {
-            span.end(startTime);
-          }
-        });
+        /**
+         * As this point if the callback wasn't called, that means either the
+         * layer is asynchronous (so it will call the callback later on) or that
+         * the layer directly end the http response, so we'll hook into the "finish"
+         * event to handle the later case.
+         */
+        req.res?.once('finish', onResponseFinish);
         return result;
       };
     });

plugins/node/opentelemetry-plugin-express/test/express.test.ts

@@ -90,8 +90,10 @@ describe('Express Plugin', () => {
       });
       const router = express.Router();
       app.use('/toto', router);
-      router.get('/:id', (req, res, next) => {
-        return res.status(200).end();
+      router.get('/:id', (req, res) => {
+        setTimeout(() => {
+          return res.status(200).end();
+        }, 5);
       });
       const server = http.createServer(app);
       await new Promise(resolve => server.listen(0, resolve));