From b4ab719bfef5d9c6b47dc854553e379b1f14af42 Mon Sep 17 00:00:00 2001
From: bojiang <bojiang_@outlook.com>
Date: Thu, 15 Sep 2022 17:59:22 +0800
Subject: [PATCH] amend

---
 bentoml/_internal/server/instruments.py | 65 +++++++++++++------------
 1 file changed, 35 insertions(+), 30 deletions(-)

diff --git a/bentoml/_internal/server/instruments.py b/bentoml/_internal/server/instruments.py
index 1d1c530ffe3..06e8df23590 100644
--- a/bentoml/_internal/server/instruments.py
+++ b/bentoml/_internal/server/instruments.py
@@ -131,36 +131,41 @@ async def __call__(
 
         async def wrapped_send(message: "ext.ASGIMessage") -> None:
             if message["type"] == "http.response.start":
-                status_code = message["status"]
-
-                # instrument request total count
-                self.legacy_metrics_request_total.labels(
-                    endpoint=endpoint,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                ).inc()
-                self.metrics_request_total.labels(
-                    endpoint=endpoint,
-                    service_name=component_context.bento_name,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                ).inc()
-
-                # instrument request duration
-                assert START_TIME_VAR.get() != 0
-                total_time = max(default_timer() - START_TIME_VAR.get(), 0)
-                self.legacy_metrics_request_duration.labels(  # type: ignore
-                    endpoint=endpoint,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                ).observe(total_time)
-                self.metrics_request_duration.labels(  # type: ignore
-                    endpoint=endpoint,
-                    service_name=component_context.bento_name,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                ).observe(total_time)
-            START_TIME_VAR.set(0)
+                STATUS_VAR.set(message["status"])
+            elif message["type"] == "http.response.body":
+                if ("more_body" not in message) or not message["more_body"]:
+                    assert START_TIME_VAR.get() != 0
+                    assert STATUS_VAR.get() != 0
+
+                    # instrument request total count
+                    self.legacy_metrics_request_total.labels(
+                        endpoint=endpoint,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                    ).inc()
+                    self.metrics_request_total.labels(
+                        endpoint=endpoint,
+                        service_name=component_context.bento_name,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                    ).inc()
+
+                    # instrument request duration
+                    total_time = max(default_timer() - START_TIME_VAR.get(), 0)
+                    self.legacy_metrics_request_duration.labels(  # type: ignore
+                        endpoint=endpoint,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                    ).observe(total_time)
+                    self.metrics_request_duration.labels(  # type: ignore
+                        endpoint=endpoint,
+                        service_name=component_context.bento_name,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                    ).observe(total_time)
+
+                    START_TIME_VAR.set(0)
+                    STATUS_VAR.set(0)
             await send(message)
 
         with self.legacy_metrics_request_in_progress.labels(