# HG changeset patch
# User Mads Kiilerich <mads@kiilerich.com>
# Date 2019-10-09 12:23:21
# Node ID 19418a4c6c6119fa3b3176bc957bbc8f7374a50f
# Parent  b42ee1bdf0820e3dd5da2d0bbf624acb83d256b0

wsgi: make WSGI wrapper follow the size of the result and log when it finished

The WSGI protocol allow for explicit writes and for iterator yields, and we
thus have to do accounting in both places.

diff --git a/kallithea/lib/middleware/wrapper.py b/kallithea/lib/middleware/wrapper.py
--- a/kallithea/lib/middleware/wrapper.py
+++ b/kallithea/lib/middleware/wrapper.py
@@ -38,12 +38,27 @@ log = logging.getLogger(__name__)
 
 class Meter:
 
-    def __init__(self):
+    def __init__(self, start_response):
+        self._start_response = start_response
         self._start = time.time()
+        self._size = 0
 
     def duration(self):
         return time.time() - self._start
 
+    def start_response(self, status, response_headers, exc_info=None):
+        write = self._start_response(status, response_headers, exc_info)
+        def metered_write(s):
+            self.measure(s)
+            write(s)
+        return metered_write
+
+    def measure(self, chunk):
+        self._size += len(chunk)
+
+    def size(self):
+        return self._size
+
 
 class ResultIter:
 
@@ -58,11 +73,12 @@ class ResultIter:
 
     def next(self):
         chunk = self._next()
+        self._meter.measure(chunk)
         return chunk
 
     def close(self):
         self._result_close()
-        log.info("%s responded after %.3fs", self._description, self._meter.duration())
+        log.info("%s responded after %.3fs with %s bytes", self._description, self._meter.duration(), self._meter.size())
 
 
 class RequestWrapper(object):
@@ -72,13 +88,13 @@ class RequestWrapper(object):
         self.config = config
 
     def __call__(self, environ, start_response):
-        meter = Meter()
+        meter = Meter(start_response)
         description = "Request from %s for %s" % (
             _get_ip_addr(environ),
             safe_unicode(_get_access_path(environ)),
         )
         try:
-            result = self.application(environ, start_response)
+            result = self.application(environ, meter.start_response)
         finally:
             log.info("%s responding after %.3fs", description, meter.duration())
         return ResultIter(result, meter, description)