SERVER-67046 Create system memory monitor

6 files changed, 126 insertions, 108 deletions

diff --git a/site_scons/site_tools/build_metrics.py b/site_scons/site_tools/build_metrics.py
deleted file mode 100644
index 0ab660a45e9..00000000000
--- a/site_scons/site_tools/build_metrics.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# Copyright 2020 MongoDB Inc.
-#
-# Permission is hereby granted, free of charge, to any person obtaining
-# a copy of this software and associated documentation files (the
-# "Software"), to deal in the Software without restriction, including
-# without limitation the rights to use, copy, modify, merge, publish,
-# distribute, sublicense, and/or sell copies of the Software, and to
-# permit persons to whom the Software is furnished to do so, subject to
-# the following conditions:
-#
-# The above copyright notice and this permission notice shall be included
-# in all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
-# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
-# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
-# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
-# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
-# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
-# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
-#
-"""Configure the build to track build performance."""
-
-import time
-import os
-import sys
-import atexit
-import json
-
-import psutil
-
-_SEC_TO_NANOSEC_FACTOR = 1000000000.0
-
-_BUILD_METRIC_DATA = {}
-
-
-def get_build_metric_dict():
-    global _BUILD_METRIC_DATA
-    return _BUILD_METRIC_DATA
-
-
-import atexit
-
-
-# This section is an excerpt of the original
-# https://stackoverflow.com/a/63029332/1644736
-class CaptureAtexits:
-    def __init__(self):
-        self.captured = []
-
-    def __eq__(self, other):
-        self.captured.append(other)
-        return False
-
-
-def finalize_build_metrics(env):
-    metrics = get_build_metric_dict()
-    metrics['end_time'] = time.time_ns()
-
-    build_metrics_file = env.GetOption('build-metrics')
-    if build_metrics_file == '-':
-        json.dump(metrics, sys.stdout, indent=4, sort_keys=True)
-    else:
-        with open(build_metrics_file, 'w') as f:
-            json.dump(metrics, f, indent=4, sort_keys=True)
-
-
-def add_meta_data(env, key, value):
-    get_build_metric_dict()[key] = value
-
-
-def generate(env, **kwargs):
-
-    # This will force our at exit to the of the stack ensuring
-    # that it is the last thing called when exiting.
-    c = CaptureAtexits()
-    atexit.unregister(c)
-    for func in c.captured:
-        atexit.unregister(func)
-    atexit.register(finalize_build_metrics, env)
-    for func in c.captured:
-        atexit.register(func)
-
-    env.AddMethod(get_build_metric_dict, "GetBuildMetricDictionary")
-    env.AddMethod(add_meta_data, "AddBuildMetricsMetaData")
-
-    metrics = get_build_metric_dict()
-    p = psutil.Process(os.getpid())
-
-    metrics['start_time'] = int(p.create_time() * _SEC_TO_NANOSEC_FACTOR)
-    metrics['scons_command'] = " ".join([sys.executable] + sys.argv)
-
-
-def exists(env):
-    return True
diff --git a/site_scons/site_tools/build_metrics/__init__.py b/site_scons/site_tools/build_metrics/__init__.py
index 0f3bd6c078d..f29c3328931 100644
--- a/site_scons/site_tools/build_metrics/__init__.py
+++ b/site_scons/site_tools/build_metrics/__init__.py
@@ -30,14 +30,10 @@ import time
 from jsonschema import validate
 import psutil
 
-_SEC_TO_NANOSEC_FACTOR = 1000000000.0
-
-_BUILD_METRIC_DATA = {}
-
+from .util import add_meta_data, get_build_metric_dict
+import build_metrics.memory
 
-def get_build_metric_dict():
-    global _BUILD_METRIC_DATA
-    return _BUILD_METRIC_DATA
+_SEC_TO_NANOSEC_FACTOR = 1000000000.0
 
 
 # This section is an excerpt of the original
@@ -54,6 +50,8 @@ class CaptureAtexits:
 def finalize_build_metrics(env):
     metrics = get_build_metric_dict()
     metrics['end_time'] = time.time_ns()
+    for m in _METRICS_COLLECTORS:
+        m.finalize()
 
     with open(os.path.join(os.path.dirname(__file__), "build_metrics_format.schema")) as f:
         validate(metrics, json.load(f))
@@ -66,11 +64,11 @@ def finalize_build_metrics(env):
             json.dump(metrics, f, indent=4, sort_keys=True)
 
 
-def add_meta_data(env, key, value):
-    get_build_metric_dict()[key] = value
+_METRICS_COLLECTORS = []
 
 
 def generate(env, **kwargs):
+    global _METRICS_COLLECTORS
 
     # This will force our at exit to the of the stack ensuring
     # that it is the last thing called when exiting.
@@ -91,6 +89,8 @@ def generate(env, **kwargs):
     metrics['start_time'] = int(p.create_time() * _SEC_TO_NANOSEC_FACTOR)
     metrics['scons_command'] = " ".join([sys.executable] + sys.argv)
 
+    _METRICS_COLLECTORS = [memory.MemoryMonitor(psutil.Process().memory_info().vms)]
+
 
 def exists(env):
     return True
diff --git a/site_scons/site_tools/build_metrics/build_metrics_format.schema b/site_scons/site_tools/build_metrics/build_metrics_format.schema
index 5b05b1491c9..4b6fec50068 100644
--- a/site_scons/site_tools/build_metrics/build_metrics_format.schema
+++ b/site_scons/site_tools/build_metrics/build_metrics_format.schema
@@ -1,10 +1,38 @@
 {
+    "$defs": {
+        "timestamp": {
+            "type": "integer",
+            "description": "Nanoseconds since Unix epoch"
+        },
+        "memory": {
+            "type": "integer",
+            "description": "Virtual memory used in bytes"
+        }
+    },
     "type" : "object",
     "properties" : {
-        "start_time" : {"type" : "integer"},
-        "end_time" : {"type" : "integer"},
+        "start_time" : { "$ref": "#/$defs/timestamp" },
+        "end_time" : { "$ref": "#/$defs/timestamp" },
         "evg_id" : {"type" : "string"},
         "variant" : {"type" : "string"},
-        "scons_command" : {"type" : "string"}
+        "scons_command" : {"type" : "string"},
+        "system_memory": {
+            "type": "object",
+            "properties": {
+                "mem_over_time": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "timestamp": { "$ref": "#/$defs/timestamp" },
+                            "memory": { "$ref": "#/$defs/memory" }
+                        }
+                    }
+                },
+                "max": { "$ref": "#/$defs/memory" },
+                "arithmetic_mean": {"type": "number"},
+                "start_mem": { "$ref": "#/$defs/memory" }
+            }
+        }
     }
-}
\ No newline at end of file
+}
diff --git a/site_scons/site_tools/build_metrics/memory.py b/site_scons/site_tools/build_metrics/memory.py
new file mode 100644
index 00000000000..efa8320eae3
--- /dev/null
+++ b/site_scons/site_tools/build_metrics/memory.py
@@ -0,0 +1,60 @@
+import threading
+import time
+import psutil
+import sys
+from .util import timestamp_now, get_build_metric_dict
+from .protocol import BuildMetricsCollector
+
+
+class MemoryMonitor(BuildMetricsCollector):
+    INTERVAL = 0.1  # seconds
+
+    def __init__(self, starting_memory_adjustment=0):
+        self._stop = False
+        metrics = get_build_metric_dict()
+        metrics['system_memory'] = {
+            "mem_over_time": [],
+            "start_mem": used_memory() - starting_memory_adjustment,
+        }
+
+        self._thread = threading.Thread(target=self.memory_monitor, daemon=True)
+        self._thread.start()
+
+    def finalize(self):
+        self._stop = True
+        self._record_data_point()
+
+        metrics = get_build_metric_dict()
+        sys_mem = metrics["system_memory"]
+
+        mean = 0
+        max_ = 0
+        count = 1
+        for val in sys_mem["mem_over_time"]:
+            max_ = max(val["memory"], max_)
+            # iterative mean calculation algorithm from https://stackoverflow.com/a/1934266
+            mean += (val["memory"] - mean) / count
+            count += 1
+
+        sys_mem["arithmetic_mean"] = mean
+        sys_mem["max"] = max_
+
+    def memory_monitor(self):
+        while not self._stop:
+            time.sleep(self.INTERVAL)
+            if self._stop:
+                break
+
+            self._record_data_point()
+
+    def _record_data_point(self):
+        used_mem = used_memory()
+        now_time = timestamp_now()
+
+        metrics = get_build_metric_dict()
+        metrics["system_memory"]["mem_over_time"].append(
+            {"timestamp": now_time, "memory": used_mem})
+
+
+def used_memory():
+    return psutil.virtual_memory().used
diff --git a/site_scons/site_tools/build_metrics/protocol.py b/site_scons/site_tools/build_metrics/protocol.py
new file mode 100644
index 00000000000..b0158ce7ef3
--- /dev/null
+++ b/site_scons/site_tools/build_metrics/protocol.py
@@ -0,0 +1,8 @@
+from typing import Protocol
+from abc import abstractmethod
+
+
+class BuildMetricsCollector(Protocol):
+    @abstractmethod
+    def finalize(self):
+        raise NotImplementedError
diff --git a/site_scons/site_tools/build_metrics/util.py b/site_scons/site_tools/build_metrics/util.py
new file mode 100644
index 00000000000..967bc5f33cc
--- /dev/null
+++ b/site_scons/site_tools/build_metrics/util.py
@@ -0,0 +1,17 @@
+import datetime
+import time
+
+_BUILD_METRIC_DATA = {}
+
+
+def get_build_metric_dict():
+    global _BUILD_METRIC_DATA
+    return _BUILD_METRIC_DATA
+
+
+def add_meta_data(env, key, value):
+    get_build_metric_dict()[key] = value
+
+
+def timestamp_now() -> int:
+    return time.time_ns()