From d3fd35547287a30ad5f9c9d6390f38a84c1e34b6 Mon Sep 17 00:00:00 2001
From: John Hensley <john@freedom.press>
Date: Tue, 31 Mar 2020 13:46:52 -0400
Subject: [PATCH] Add utils.chronometer for measuring block execution time

---
 securedrop_client/storage.py | 17 +++++++++++++----
 securedrop_client/utils.py   | 20 +++++++++++++++++++-
 tests/api_jobs/test_sync.py  |  2 +-
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/securedrop_client/storage.py b/securedrop_client/storage.py
index 667859faa..cc60a3281 100644
--- a/securedrop_client/storage.py
+++ b/securedrop_client/storage.py
@@ -34,6 +34,8 @@
 from securedrop_client.crypto import CryptoError, GpgHelper
 from securedrop_client.db import (DraftReply, Source, Message, File, Reply, ReplySendStatus,
                                   ReplySendStatusCodes, User)
+from securedrop_client.utils import chronometer
+
 from sdclientapi import API
 from sdclientapi import Source as SDKSource
 from sdclientapi import Submission as SDKSubmission
@@ -120,10 +122,17 @@ def update_local_storage(session: Session,
     # The following update_* functions may change the database state.
     # Because of that, each get_local_* function needs to be called just before
     # its respective update_* function.
-    update_sources(gpg, remote_sources, get_local_sources(session), session, data_dir)
-    update_files(remote_files, get_local_files(session), session, data_dir)
-    update_messages(remote_messages, get_local_messages(session), session, data_dir)
-    update_replies(remote_replies, get_local_replies(session), session, data_dir)
+    with chronometer(logger, "update_sources"):
+        update_sources(gpg, remote_sources, get_local_sources(session), session, data_dir)
+
+    with chronometer(logger, "update_files"):
+        update_files(remote_files, get_local_files(session), session, data_dir)
+
+    with chronometer(logger, "update_messages"):
+        update_messages(remote_messages, get_local_messages(session), session, data_dir)
+
+    with chronometer(logger, "update_replies"):
+        update_replies(remote_replies, get_local_replies(session), session, data_dir)
 
 
 def update_source_key(gpg: GpgHelper, local_source: Source, remote_source: SDKSource) -> None:
diff --git a/securedrop_client/utils.py b/securedrop_client/utils.py
index 9f872a091..515a29303 100644
--- a/securedrop_client/utils.py
+++ b/securedrop_client/utils.py
@@ -1,5 +1,10 @@
-import os
+import logging
 import math
+import os
+import time
+
+from contextlib import contextmanager
+from typing import Generator
 
 
 def safe_mkdir(sdc_home: str, relative_path: str = None) -> None:
@@ -67,3 +72,16 @@ def humanize_filesize(filesize: int) -> str:
         return '{}KB'.format(math.floor(filesize / 1024))
     else:
         return '{}MB'.format(math.floor(filesize / 1024 ** 2))
+
+
+@contextmanager
+def chronometer(logger: logging.Logger, description: str) -> Generator:
+    """
+    Measures the execution time of its block.
+    """
+    try:
+        start = time.perf_counter()
+        yield
+    finally:
+        elapsed = time.perf_counter() - start
+        logger.debug(f"{description} duration: {elapsed:.4f}s")
diff --git a/tests/api_jobs/test_sync.py b/tests/api_jobs/test_sync.py
index 05ae25d33..8e2c0afea 100644
--- a/tests/api_jobs/test_sync.py
+++ b/tests/api_jobs/test_sync.py
@@ -136,5 +136,5 @@ def test_MetadataSyncJob_only_import_new_source_keys(mocker, homedir, session, s
 
     assert mock_get_remote_data.call_count == 2
 
-    log_msg = storage_logger.debug.call_args_list[1][0][0]
+    log_msg = storage_logger.debug.call_args_list[5][0][0]
     assert log_msg == 'Source key data is unchanged'