From 14c6eaaa7b193aa73b9e2d7dcd441e1f3d872f6e Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Fri, 18 Oct 2024 17:22:54 -0400
Subject: [PATCH] Wait for docker container to start, draw horizontal lines,
 use 0-based start time for containers instead of real wall time, and add a
 graph title.

---
 docker/scripts/graph_docker_memory.py | 56 ++++++++++++++++++++++-----
 1 file changed, 46 insertions(+), 10 deletions(-)

diff --git a/docker/scripts/graph_docker_memory.py b/docker/scripts/graph_docker_memory.py
index fef7082..2de5b71 100755
--- a/docker/scripts/graph_docker_memory.py
+++ b/docker/scripts/graph_docker_memory.py
@@ -1,13 +1,14 @@
 #!/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
+import re
+import subprocess
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from time import sleep
+from typing import Collection, Final, NewType, Tuple
 
 ContainerId = NewType("ContainerId", str)
 ContainerName = NewType("ContainerName", str)
@@ -30,11 +31,16 @@ def main():
     logging.basicConfig(level=logging.INFO)
     samples: list[Sample] = []
     labels: dict[ContainerId, ContainerName] = {}
+    first_pass = True
     # First wait for any docker container to exist.
     while True:
         sample, labels_in_sample = take_sample()
         if labels_in_sample:
             break
+        if first_pass:
+            first_pass = False
+            logging.info("Waiting for a docker container to exist to start recording.")
+        sleep(1)
     # And then record memory until no containers exist.
     while True:
         sample, labels_in_sample = take_sample()
@@ -44,18 +50,47 @@ def main():
         labels = {**labels, **labels_in_sample}
         sleep(SAMPLE_INTERVAL_SECONDS)
     if labels:
-        write_plot(samples, labels)
+        # Draws a red horizontal line at 32 GiB since that is the memory limit for cloud run.
+        write_plot(
+            samples,
+            labels,
+            horizontal_lines=[(32 * 1024**3, "red", "Cloud Run Max Memory")],
+        )
 
 
-def write_plot(samples: Collection[Sample], labels: dict[ContainerId, ContainerName]):
+def write_plot(
+    samples: Collection[Sample],
+    labels: dict[ContainerId, ContainerName],
+    *,
+    horizontal_lines: Collection[Tuple[int, str, str | None]] = [],
+):
+    starting_time_per_container = {
+        container_id: min(
+            (sample.instant for sample in samples if container_id in sample.stats)
+        )
+        for container_id in labels.keys()
+    }
     print(
         """set terminal svg background '#FFFFFF'
+set title 'Docker Memory Usage'
 set xdata time
 set timefmt '%s'
-set format x '%H:%M:%S'
+set format x '%tH:%tM:%tS'
+# Please note this is in SI units (base 10), not IEC (base 2). So, for example, this would show a Gigabyte, not a Gibibyte.
+set format y '%.0s%cB'
 set datafile separator "|"
 """
     )
+    for y_value, color, label in horizontal_lines:
+        print(
+            f'''set arrow from graph 0, first {y_value} to graph 1, first {y_value} nohead linewidth 2 linecolor rgb "{color}"'''
+        )
+        if label is not None:
+            print(f"""set label "{label}" at graph 0, first {y_value} offset 1,-0.5""")
+
+    # Include the horizontal lines in the range
+    if len(horizontal_lines) > 0:
+        print(f"""set yrange [*:{max(x[0] for x in horizontal_lines)}<*]""")
     line_definitions = ", ".join(
         [
             f""""-" using 1:2 title '{gnuplot_escape(name)}' with lines"""
@@ -64,12 +99,13 @@ set datafile separator "|"
     )
     print("plot", line_definitions)
     for container_id in sorted(labels.keys()):
+        start_time = int(starting_time_per_container[container_id].timestamp())
         for sample in sorted(samples, key=lambda x: x.instant):
             if container_id in sample.stats:
                 print(
                     "|".join(
                         [
-                            str(int(sample.instant.timestamp())),
+                            str(int((sample.instant).timestamp()) - start_time),
                             str(sample.stats[container_id].memory_usage_bytes),
                         ]
                     )