Refactor calculating execution time

- Implement a context manager logging execution time.
- Use that context manager in daemon's handle_DATA method.

lacre/daemon.py

@@ -3,6 +3,7 @@
 import logging
 import lacre
 from lacre.text import DOUBLE_EOL_BYTES
+from lacre.stats import time_logger
 import lacre.config as conf
 import sys
 from aiosmtpd.controller import Controller
@@ -10,7 +11,6 @@ from aiosmtpd.smtp import Envelope
 import asyncio
 import email
 from email.policy import SMTPUTF8
-import time
 
 # Load configuration and init logging, in this order.  Only then can we load
 # the last Lacre module, i.e. lacre.mailgate.
@@ -34,43 +34,40 @@ class MailEncryptionProxy:
 
     async def handle_DATA(self, server, session, envelope: Envelope):
         """Accept a message and either encrypt it or forward as-is."""
-        start = time.process_time()
-        try:
-            keys = self._keyring.freeze_identities()
-            LOG.debug('Parsing message: %s', self._beginning(envelope))
-            message = email.message_from_bytes(envelope.original_content, policy=SMTPUTF8)
-            LOG.debug('Parsed into %s: %s', type(message), repr(message))
+        with time_logger('Message delivery', LOG):
+            try:
+                keys = self._keyring.freeze_identities()
+                LOG.debug('Parsing message: %s', self._beginning(envelope))
+                message = email.message_from_bytes(envelope.original_content, policy=SMTPUTF8)
+                LOG.debug('Parsed into %s: %s', type(message), repr(message))
 
-            if message.defects:
-                # Sometimes a weird message cannot be encoded back and
-                # delivered, so before bouncing such messages we at least
-                # record information about the issues.  Defects are identified
-                # by email.* package.
-                LOG.warning("Issues found: %d; %s", len(message.defects), repr(message.defects))
+                if message.defects:
+                    # Sometimes a weird message cannot be encoded back and
+                    # delivered, so before bouncing such messages we at least
+                    # record information about the issues.  Defects are identified
+                    # by email.* package.
+                    LOG.warning("Issues found: %d; %s", len(message.defects), repr(message.defects))
 
-            if conf.flag_enabled('daemon', 'log_headers'):
-                LOG.info('Message headers: %s', self._extract_headers(message))
+                if conf.flag_enabled('daemon', 'log_headers'):
+                    LOG.info('Message headers: %s', self._extract_headers(message))
 
-            send = xport.SendFrom(envelope.mail_from)
-            for operation in gate.delivery_plan(envelope.rcpt_tos, message, keys):
-                LOG.debug(f"Sending mail via {operation!r}")
-                try:
-                    new_message = operation.perform(message)
-                    send(new_message, operation.recipients())
-                except EncryptionException:
-                    # If the message can't be encrypted, deliver cleartext.
-                    LOG.exception('Unable to encrypt message, delivering in cleartext')
-                    if not isinstance(operation, KeepIntact):
-                        self._send_unencrypted(operation, message, envelope, send)
-                    else:
-                        LOG.error(f'Cannot perform {operation}')
+                send = xport.SendFrom(envelope.mail_from)
+                for operation in gate.delivery_plan(envelope.rcpt_tos, message, keys):
+                    LOG.debug(f"Sending mail via {operation!r}")
+                    try:
+                        new_message = operation.perform(message)
+                        send(new_message, operation.recipients())
+                    except EncryptionException:
+                        # If the message can't be encrypted, deliver cleartext.
+                        LOG.exception('Unable to encrypt message, delivering in cleartext')
+                        if not isinstance(operation, KeepIntact):
+                            self._send_unencrypted(operation, message, envelope, send)
+                        else:
+                            LOG.error(f'Cannot perform {operation}')
 
-        except:
-            LOG.exception('Unexpected exception caught, bouncing message')
-            return xport.RESULT_ERROR
-
-        ellapsed = (time.process_time() - start) * 1000
-        LOG.info(f'Message delivered in {ellapsed:.2f} ms')
+            except:
+                LOG.exception('Unexpected exception caught, bouncing message')
+                return xport.RESULT_ERROR
 
         return xport.RESULT_OK
 
@@ -95,6 +92,9 @@ class MailEncryptionProxy:
             'cte'      : message['Content-Transfer-Encoding']
             }
 
+    def _seconds_between(self, start_ms, end_ms) -> float:
+        return (end_ms - start_ms) * 1000
+
 
 def _init_controller(keys: kcache.KeyRing, max_body_bytes=None, tout: float = 5):
     proxy = MailEncryptionProxy(keys)

lacre/stats.py

@@ -0,0 +1,27 @@
+"""Insights into Lacre's inner workings."""
+
+import time
+import logging
+
+class ExecutionTimeLogger:
+    """Context-manager that measures how much time some operation took and logs it."""
+
+    def __init__(self, message: str, logger: logging.Logger):
+        self._message = message
+        self._log = logger
+        self._start = None
+
+    def __enter__(self):
+        self._start = time.process_time()
+
+    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
+        end = time.process_time()
+        ellapsed = (end - self._start) * 1000
+
+        if exc_type:
+            self._log.error('%s took %d ms, raised exception %s', self._message, ellapsed, exc_type)
+        else:
+            self._log.info('%s took %d ms', self._message, ellapsed)
+
+def time_logger(msg: str, logger: logging.Logger):
+    return ExecutionTimeLogger(msg, logger)