From 927c14d2384466d8483189b2ef2d854abcadbf2a Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Thu, 17 Oct 2024 19:41:33 -0400
Subject: [PATCH] Add a script that graphs memory usage of docker containers
 over time.

---
 docker/scripts/graph_docker_memory.py | 125 ++++++++++++++++++++++++++
 1 file changed, 125 insertions(+)
 create mode 100755 docker/scripts/graph_docker_memory.py

diff --git a/docker/scripts/graph_docker_memory.py b/docker/scripts/graph_docker_memory.py
new file mode 100755
index 0000000..1ee9bf6
--- /dev/null
+++ b/docker/scripts/graph_docker_memory.py
@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+from __future__ import annotations
+from datetime import datetime, timedelta
+from typing import NewType, Final, Collection, Tuple
+from dataclasses import dataclass
+from time import sleep
+import subprocess
+import json
+import re
+import logging
+
+ContainerId = NewType("ContainerId", str)
+ContainerName = NewType("ContainerName", str)
+
+SAMPLE_INTERVAL_SECONDS: Final[int] = 2
+
+
+@dataclass
+class Sample:
+    instant: datetime
+    stats: dict[ContainerId, Stats]
+
+
+@dataclass
+class Stats:
+    memory_usage_bytes: int
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+    samples: list[Sample] = []
+    labels: dict[ContainerId, ContainerName] = {}
+    while True:
+        sample, labels_in_sample = take_sample()
+        if not labels_in_sample:
+            break
+        samples.append(sample)
+        labels = {**labels, **labels_in_sample}
+        sleep(SAMPLE_INTERVAL_SECONDS)
+    if labels:
+        write_plot(samples, labels)
+
+
+def write_plot(samples: Collection[Sample], labels: dict[ContainerId, ContainerName]):
+    print(
+        """set terminal svg background '#FFFFFF'
+set xdata time
+set timefmt '%s'
+set format x '%H:%M:%S'
+set datafile separator "|"
+"""
+    )
+    line_definitions = ", ".join(
+        [
+            f""""-" using 1:2 title '{gnuplot_escape(name)}' with lines"""
+            for container_id, name in sorted(labels.items())
+        ]
+    )
+    print("plot", line_definitions)
+    for container_id in sorted(labels.keys()):
+        for sample in sorted(samples, key=lambda x: x.instant):
+            if container_id in sample.stats:
+                print(
+                    "|".join(
+                        [
+                            str(int(sample.instant.timestamp())),
+                            str(sample.stats[container_id].memory_usage_bytes),
+                        ]
+                    )
+                )
+        print("e")
+
+
+def gnuplot_escape(inp: str) -> str:
+    out = ""
+    for c in inp:
+        if c == "_":
+            out += "\\"
+        out += c
+    return out
+
+
+def take_sample() -> Tuple[Sample, dict[ContainerId, ContainerName]]:
+    labels: dict[ContainerId, ContainerName] = {}
+    stats: dict[ContainerId, Stats] = {}
+    docker_inspect = subprocess.run(
+        ["docker", "stats", "--no-stream", "--no-trunc", "--format", "json"],
+        stdout=subprocess.PIPE,
+    )
+    for container_stat in (
+        json.loads(l) for l in docker_inspect.stdout.decode("utf8").splitlines()
+    ):
+        if not container_stat["ID"]:
+            # When containers are starting up, they sometimes have no ID and "--" as the name.
+            continue
+        labels[ContainerId(container_stat["ID"])] = ContainerName(
+            container_stat["Name"]
+        )
+        memory_usage = parse_mem_usage(container_stat["MemUsage"])
+        stats[ContainerId(container_stat["ID"])] = Stats(
+            memory_usage_bytes=memory_usage
+        )
+    for container_id, container_stat in stats.items():
+        logging.info(
+            f"Recorded stat {labels[container_id]}: {container_stat.memory_usage_bytes} bytes"
+        )
+    return Sample(instant=datetime.now(), stats=stats), labels
+
+
+def parse_mem_usage(mem_usage: str) -> int:
+    parsed_mem_usage = re.match(
+        r"(?P<number>[0-9]+\.?[0-9]*)(?P<unit>[^\s]+)", mem_usage
+    )
+    if parsed_mem_usage is None:
+        raise Exception(f"Invalid Mem Usage: {mem_usage}")
+    number = float(parsed_mem_usage.group("number"))
+    unit = parsed_mem_usage.group("unit")
+    for multiplier, identifier in enumerate(["B", "KiB", "MiB", "GiB", "TiB"]):
+        if unit == identifier:
+            return int(number * (1024**multiplier))
+    raise Exception(f"Unrecognized unit: {unit}")
+
+
+if __name__ == "__main__":
+    main()