Added new commands to memray attach

Letting memray attach supports deactivating an attached tracker manually, or after a specified heap size is reached, or after a specified time limit has elapsed.

Signed-off-by: Ivona Stojanovic <[email protected]>

src/memray/commands/attach.py

@@ -12,22 +12,49 @@
 import subprocess
 import sys
 import threading
+from enum import IntEnum
 
 import memray
 from memray._errors import MemrayCommandError
 
 from .live import LiveCommand
 from .run import _get_free_port
 
+
+class TrackingMode(IntEnum):
+    ACTIVATE = 1
+    DEACTIVATE = 2
+    UNTIL_HEAP_SIZE = 3
+    FOR_DURATION = 4
+
+
 GDB_SCRIPT = pathlib.Path(__file__).parent / "_attach.gdb"
 LLDB_SCRIPT = pathlib.Path(__file__).parent / "_attach.lldb"
 RTLD_DEFAULT = memray._memray.RTLD_DEFAULT
 RTLD_NOW = memray._memray.RTLD_NOW
 PAYLOAD = """
 import atexit
+import time
+import threading
+import resource
+import logging
+import pathlib
+from typing import Set
+from typing import Union
+from contextlib import suppress
 
 import memray
+from memray.commands.attach import TrackingMode
+
+ACTIVE_THREADS: Set[Union[threading.Thread, threading.Timer]] = set()
+TRACKER_LOCK = threading.Lock()
+
+logger = logging.getLogger(__name__)
 
+def _get_current_heap_size() -> int:
+    usage = resource.getrusage(resource.RUSAGE_SELF)
+    rss_bytes = usage.ru_maxrss * 1024  # Convert from KB to bytes
+    return rss_bytes
 
 def deactivate_last_tracker():
     tracker = getattr(memray, "_last_tracker", None)
@@ -37,23 +64,64 @@ def deactivate_last_tracker():
     memray._last_tracker = None
     tracker.__exit__(None, None, None)
 
+    for thread in ACTIVE_THREADS:
+        with suppress(AttributeError):
+            thread.cancel()  # type: ignore
+
+    ACTIVE_THREADS.clear()
+
+def activate_tracker():
+    deactivate_last_tracker()
+    tracker = {tracker_call}
+    try:
+        tracker.__enter__()
+    except:
+        # Prevent the exception from keeping the tracker alive.
+        # This way resources are cleaned up ASAP.
+        del tracker
+        raise
+
+    memray._last_tracker = tracker
+
+def track_until_heap_size(heap_size):
+    activate_tracker()
+
+    def check_heap_size() -> None:
+        while getattr(memray, "_last_tracker", None) is not None:
+            current_heap_size = _get_current_heap_size()
+            if current_heap_size >= heap_size:
+                logger.warning(
+                    "Heap size of %d reached (%d), disabling tracking",
+                    heap_size,
+                    current_heap_size,
+                )
+                deactivate_last_tracker()
+                break
+            time.sleep(1)
+
+    thread = threading.Thread(target=check_heap_size)
+    with TRACKER_LOCK:
+        ACTIVE_THREADS.add(thread)
+    thread.start()
+
+def track_for_duration(duration=5):
+    activate_tracker()
+    thread = threading.Timer(duration, deactivate_last_tracker)
+    thread.start()
+    ACTIVE_THREADS.add(thread)
 
 if not hasattr(memray, "_last_tracker"):
     # This only needs to be registered the first time we attach.
     atexit.register(deactivate_last_tracker)
 
-deactivate_last_tracker()
-
-tracker = {tracker_call}
-try:
-    tracker.__enter__()
-except:
-    # Prevent the exception from keeping the tracker alive.
-    # This way resources are cleaned up ASAP.
-    del tracker
-    raise
-
-memray._last_tracker = tracker
+if {mode} == TrackingMode.ACTIVATE:
+    activate_tracker()
+elif {mode} == TrackingMode.DEACTIVATE:
+    deactivate_last_tracker()
+elif {mode} == TrackingMode.UNTIL_HEAP_SIZE:
+    track_until_heap_size({heap_size})
+elif {mode} == TrackingMode.FOR_DURATION:
+    track_for_duration({duration})
 """
 
 
@@ -282,6 +350,22 @@ def prepare_parser(self, parser: argparse.ArgumentParser) -> None:
             choices=["auto", "gdb", "lldb"],
         )
 
+        parser.add_argument(
+            "--mode",
+            help="Activate or deactivate tracking.",
+            type=str,
+            default="activate",
+            choices=["activate", "deactivate"],
+        )
+
+        parser.add_argument(
+            "--heap_size", type=int, help="Heap size to track until (in bytes)"
+        )
+
+        parser.add_argument(
+            "--duration", type=int, help="Duration to track for (in seconds)"
+        )
+
         parser.add_argument(
             "-v",
             "--verbose",
@@ -297,6 +381,18 @@ def prepare_parser(self, parser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
         verbose = args.verbose
+        mode = TrackingMode.ACTIVATE
+        duration = None
+        heap_size = None
+
+        if args.mode == "deactivate":
+            mode = TrackingMode.DEACTIVATE
+        elif args.heap_size:
+            mode = TrackingMode.UNTIL_HEAP_SIZE
+            heap_size = args.heap_size
+        elif args.duration:
+            mode = TrackingMode.FOR_DURATION
+            duration = args.duration
 
         if args.method == "auto":
             # Prefer gdb on Linux but lldb on macOS
@@ -351,7 +447,14 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
 
             client = server.accept()[0]
 
-        client.sendall(PAYLOAD.format(tracker_call=tracker_call).encode("utf-8"))
+        client.sendall(
+            PAYLOAD.format(
+                tracker_call=tracker_call,
+                mode=mode,
+                heap_size=heap_size,
+                duration=duration,
+            ).encode("utf-8")
+        )
         client.shutdown(socket.SHUT_WR)
 
         if not live_port: