dev: measure response data amount

HeyPuter · Jan 14, 2025 · 7fa0822 · 7fa0822
1 parent 113ed31
commit 7fa0822
Showing 1 changed file with 72 additions and 28 deletions.
diff --git a/src/backend/src/middleware/measure.js b/src/backend/src/middleware/measure.js
@@ -2,43 +2,87 @@ const express = require('express');
 const { PassThrough } = require('stream');
 const { pausing_tee } = require('../util/streamutil');
 
-function measure () {
-    return async (req, res, next) => {
-        if ( ! req.readable ) {
-            return next();
-        }
+const _intercept_req = ({ data, req }) => {
+    if ( ! req.readable ) {
+        return next();
+    }
 
-        let sz_incoming = 0;
-        // let sz_outgoing = 0; // future
+    try {
+        const [req_monitor, req_pass] = pausing_tee(req, 2);
 
-        try {
-            const [req_monitor, req_pass] = pausing_tee(req, 2);
+        req_monitor.on('data', (chunk) => {
+            data.sz_incoming += chunk.length;
+        });
 
-            req_monitor.on('data', (chunk) => {
-                sz_incoming += chunk.length;
+        const replaces = ['readable', 'pipe', 'on', 'once', 'removeListener'];
+        for ( const replace of replaces ) {
+            const replacement = req_pass[replace]
+            Object.defineProperty(req, replace, {
+                get () {
+                    if ( typeof replacement === 'function' ) {
+                        return replacement.bind(req_pass);
+                    }
+                    return replacement;
+                }
             });
+        }
+    } catch (e) {
+        console.error(e);
+        return next();
+    }
+};
 
-            const replaces = ['readable', 'pipe', 'on', 'once', 'removeListener'];
-            for ( const replace of replaces ) {
-                const replacement = req_pass[replace]
-                Object.defineProperty(req, replace, {
-                    get () {
-                        if ( typeof replacement === 'function' ) {
-                            return replacement.bind(req_pass);
-                        }
-                        return replacement;
-                    }
-                });
+const _intercept_res = ({ data, res }) => {
+    if ( ! res.writable ) {
+        return next();
+    }
+
+    try {
+        const org_write = res.write;
+        const org_end = res.end;
+
+        // Override the `write` method
+        res.write = function (chunk, ...args) {
+          if (Buffer.isBuffer(chunk)) {
+            data.sz_outgoing += chunk.length;
+          } else if (typeof chunk === 'string') {
+            data.sz_outgoing += Buffer.byteLength(chunk);
+          }
+          return org_write.apply(res, [chunk, ...args]);
+        };
+
+        // Override the `end` method
+        res.end = function (chunk, ...args) {
+          if (chunk) {
+            if (Buffer.isBuffer(chunk)) {
+              data.sz_outgoing += chunk.length;
+            } else if (typeof chunk === 'string') {
+              data.sz_outgoing += Buffer.byteLength(chunk);
             }
-        } catch (e) {
-            console.error(e);
-            return next();
-        }
+          }
+          const result = org_end.apply(res, [chunk, ...args]);
+          return result;
+        };
+    } catch (e) {
+        console.error(e);
+        return next();
+    }
+};
+
+function measure () {
+    return async (req, res, next) => {
+        const data = {
+            sz_incoming: 0,
+            sz_outgoing: 0,
+        };
+
+        _intercept_req({ data, req });
+        _intercept_res({ data, res });
 
         // Wait for the request to finish processing
         res.on('finish', () => {
-            console.log(`Incoming Data: ${sz_incoming} bytes`);
-            // console.log(`Outgoing Data: ${sz_outgoing} bytes`); // future
+            console.log(`Incoming Data: ${data.sz_incoming} bytes`);
+            console.log(`Outgoing Data: ${data.sz_outgoing} bytes`); // future
         });
 
         next();