Split tracing middleware into two

vtex · Oct 2, 2023 · 8d7d6a9 · 8d7d6a9
1 parent 68d26b2
commit 8d7d6a9
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 44 deletions.
diff --git a/src/service/metrics/requestMetricsMiddleware.ts b/src/service/metrics/requestMetricsMiddleware.ts
@@ -0,0 +1,70 @@
+import { finished as onStreamFinished } from 'stream'
+import { hrToMillisFloat } from '../../utils'
+import { ServiceContext } from '../worker/runtime/typings'
+import {
+  createConcurrentRequestsInstrument,
+  createRequestsResponseSizesInstrument,
+  createRequestsTimingsInstrument,
+  createTotalAbortedRequestsInstrument,
+  createTotalRequestsInstrument,
+  RequestsMetricLabels,
+} from '../tracing/metrics/instruments'
+
+
+export const addRequestMetricsMiddleware = () => {
+  const concurrentRequests = createConcurrentRequestsInstrument()
+  const requestTimings = createRequestsTimingsInstrument()
+  const totalRequests = createTotalRequestsInstrument()
+  const responseSizes = createRequestsResponseSizesInstrument()
+  const abortedRequests = createTotalAbortedRequestsInstrument()
+
+  return async function addRequestMetrics(ctx: ServiceContext, next: () => Promise<void>) {
+    const start = process.hrtime()
+    concurrentRequests.inc(1)
+
+    ctx.req.once('aborted', () =>
+      abortedRequests.inc({ [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName }, 1)
+    )
+
+    let responseClosed = false
+    ctx.res.once('close', () => (responseClosed = true))
+
+    try {
+      await next()
+    } finally {
+      const responseLength = ctx.response.length
+      if (responseLength) {
+        responseSizes.observe(
+          { [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName },
+          responseLength
+        )
+      }
+
+      totalRequests.inc(
+        {
+          [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName,
+          [RequestsMetricLabels.STATUS_CODE]: ctx.response.status,
+        },
+        1
+      )
+
+      const onResFinished = () => {
+        requestTimings.observe(
+          {
+            [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName,
+          },
+          hrToMillisFloat(process.hrtime(start))
+        )
+
+        concurrentRequests.dec(1)
+      }
+
+      if (responseClosed) {
+        onResFinished()
+      } else {
+        onStreamFinished(ctx.res, onResFinished)
+      }
+    }
+  }
+}
+
diff --git a/src/service/tracing/tracingMiddlewares.ts b/src/service/tracing/tracingMiddlewares.ts
@@ -7,35 +7,19 @@ import { RuntimeLogFields } from '../../tracing/LogFields'
 import { CustomHttpTags, OpentracingTags, VTEXIncomingRequestTags } from '../../tracing/Tags'
 import { UserLandTracer } from '../../tracing/UserLandTracer'
 import { cloneAndSanitizeHeaders } from '../../tracing/utils'
-import { hrToMillis, hrToMillisFloat } from '../../utils'
+import { hrToMillis } from '../../utils'
 import { ServiceContext } from '../worker/runtime/typings'
-import {
-  createConcurrentRequestsInstrument,
-  createRequestsResponseSizesInstrument,
-  createRequestsTimingsInstrument,
-  createTotalAbortedRequestsInstrument,
-  createTotalRequestsInstrument,
-  RequestsMetricLabels,
-} from './metrics/instruments'
 
 const PATHS_BLACKLISTED_FOR_TRACING = ['/_status', '/healthcheck']
 
 export const addTracingMiddleware = (tracer: Tracer) => {
-  const concurrentRequests = createConcurrentRequestsInstrument()
-  const requestTimings = createRequestsTimingsInstrument()
-  const totalRequests = createTotalRequestsInstrument()
-  const responseSizes = createRequestsResponseSizesInstrument()
-  const abortedRequests = createTotalAbortedRequestsInstrument()
 
   return async function addTracing(ctx: ServiceContext, next: () => Promise<void>) {
-    const start = process.hrtime()
-    concurrentRequests.inc(1)
     const rootSpan = tracer.extract(FORMAT_HTTP_HEADERS, ctx.request.headers) as undefined | SpanContext
     ctx.tracing = { tracer, currentSpan: undefined}
 
     if (!shouldTrace(ctx, rootSpan)) {
       await next()
-      concurrentRequests.dec(1)
       return
     }
 
@@ -47,9 +31,6 @@ export const addTracingMiddleware = (tracer: Tracer) => {
     const initialSamplingDecision = getTraceInfo(currentSpan).isSampled
 
     ctx.tracing = { currentSpan, tracer }
-    ctx.req.once('aborted', () =>
-      abortedRequests.inc({ [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName }, 1)
-    )
 
     let responseClosed = false
     ctx.res.once('close', () => (responseClosed = true))
@@ -60,22 +41,6 @@ export const addTracingMiddleware = (tracer: Tracer) => {
       ErrorReport.create({ originalError: err }).injectOnSpan(currentSpan, ctx.vtex?.logger)
       throw err
     } finally {
-      const responseLength = ctx.response.length
-      if (responseLength) {
-        responseSizes.observe(
-          { [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName },
-          responseLength
-        )
-      }
-
-      totalRequests.inc(
-        {
-          [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName,
-          [RequestsMetricLabels.STATUS_CODE]: ctx.response.status,
-        },
-        1
-      )
-
       const traceInfo = getTraceInfo(currentSpan)
       if (traceInfo?.isSampled) {
         if (!initialSamplingDecision) {
@@ -98,14 +63,6 @@ export const addTracingMiddleware = (tracer: Tracer) => {
       }
 
       const onResFinished = () => {
-        requestTimings.observe(
-          {
-            [RequestsMetricLabels.REQUEST_HANDLER]: ctx.requestHandlerName,
-          },
-          hrToMillisFloat(process.hrtime(start))
-        )
-
-        concurrentRequests.dec(1)
         currentSpan?.finish()
       }
 

diff --git a/src/service/worker/index.ts b/src/service/worker/index.ts
@@ -14,6 +14,7 @@ import { logOnceToDevConsole } from '../logger/console'
 import { LogLevel } from '../logger/logger'
 import { TracerSingleton } from '../tracing/TracerSingleton'
 import { addTracingMiddleware } from '../tracing/tracingMiddlewares'
+import { addRequestMetricsMiddleware } from '../metrics/requestMetricsMiddleware'
 import { addProcessListeners, logger } from './listeners'
 import {
   healthcheckHandler,
@@ -221,6 +222,7 @@ export const startWorker = (serviceJSON: ServiceJSON) => {
     .use(error)
     .use(prometheusLoggerMiddleware())
     .use(addTracingMiddleware(tracer))
+    .use(addRequestMetricsMiddleware())
     .use(addMetricsLoggerMiddleware())
     .use(concurrentRateLimiter(serviceJSON?.rateLimitPerReplica?.concurrent))
     .use(compress())