Python Memory Management — Event Bus That Ate 8GB

import sys
import tracemalloc

# Start tracing memory allocations
tracemalloc.start()

# --- Demonstrate size classes and sys.getsizeof ---

# Small integers are cached by CPython (-5 to 256)
small_int = 42
large_int = 1000

print(f"Size of integer 42:    {sys.getsizeof(small_int)} bytes")
print(f"Size of integer 1000:  {sys.getsizeof(large_int)} bytes")
print(f"Size of empty list:    {sys.getsizeof([])} bytes")
print(f"Size of empty dict:    {sys.getsizeof({}) } bytes")
print(f"Size of empty str:     {sys.getsizeof('')} bytes")
print()

# --- Show that small ints are the SAME object in memory ---
# CPython caches integers from -5 to 256 to avoid repeated allocation
a = 256
b = 256
print(f"a = 256, b = 256 -> same object? {a is b}")  # True — cached

c = 257
d = 257
print(f"c = 257, d = 257 -> same object? {c is d}")  # False — not cached
print()

# --- Demonstrate pymalloc vs raw malloc boundary ---
# Objects <= 512 bytes use pymalloc pools; larger use malloc directly
small_bytes = bytes(100)   # 100 bytes -> pymalloc
large_bytes = bytes(600)   # 600 bytes -> malloc directly

print(f"Size of 100-byte object: {sys.getsizeof(small_bytes)} bytes (pymalloc pool)")
print(f"Size of 600-byte object: {sys.getsizeof(large_bytes)} bytes (raw malloc)")
print()

# --- Snapshot: see what tracemalloc recorded ---
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')

print("Top 3 memory allocations in this script:")
for stat in top_stats[:3]:
    print(f"  {stat}")

tracemalloc.stop()

import gc
import sys
import ctypes

# ── PART 1: Observe reference counts directly ──────────────────────────────

class TrackedNode:
    """A simple node we'll use to build a reference cycle."""
    def __init__(self, label):
        self.label = label
        self.partner = None  # Will point to another TrackedNode

    def __del__(self):
        # This fires when the object is actually destroyed
        print(f"  [destructor] TrackedNode '{self.label}' was freed")

# Create a single node and watch the refcount
node_alpha = TrackedNode("alpha")
# getrefcount always reports +1 because the function argument itself is a reference
print(f"Refcount of node_alpha (just created): {sys.getrefcount(node_alpha) - 1}")

alias = node_alpha  # Second binding — refcount goes to 2
print(f"Refcount after creating alias:         {sys.getrefcount(node_alpha) - 1}")

del alias           # Remove one binding — refcount drops to 1
print(f"Refcount after deleting alias:         {sys.getrefcount(node_alpha) - 1}")
print()

# ── PART 2: Create an unreachable cycle and prove GC finds it ──────────────

# Disable automatic GC so we can control exactly when it runs
gc.disable()

node_one = TrackedNode("one")
node_two = TrackedNode("two")

# Wire them into a cycle: one -> two -> one
node_one.partner = node_two
node_two.partner = node_one

# Now remove the only external references to both nodes
# Reference counting CANNOT free these — each has refcount 1 from the other
print("Deleting external references to node_one and node_two...")
del node_one
del node_two
print("(No destructor fired yet — cycle keeps both alive)")
print()

# Manually check what the GC considers unreachable
unreachable_count = gc.collect()  # Collect all generations
print(f"GC collected {unreachable_count} unreachable objects")
print()

# ── PART 3: Inspect GC generations ────────────────────────────────────────

gc.enable()

print("GC generation thresholds:", gc.get_threshold())
print("GC generation counts:    ", gc.get_count())
# Thresholds: (700, 10, 10) means:
#   gen0 collects every 700 allocations
#   gen1 collects every 10 gen0 collections
#   gen2 collects every 10 gen1 collections

import weakref
import sys
import gc

# ══════════════════════════════════════════════════════════════
# PART 1: WeakValueDictionary as a memory-safe cache
# ══════════════════════════════════════════════════════════════

class ExpensiveResource:
    """Simulates an object that's costly to create (DB connection, parsed config)."""
    def __init__(self, resource_id):
        self.resource_id = resource_id

    def __repr__(self):
        return f"ExpensiveResource(id={self.resource_id})"

# A cache where entries vanish automatically when nothing else holds them
resource_cache = weakref.WeakValueDictionary()

# Create a resource and store it in the cache
db_connection = ExpensiveResource(resource_id="db-primary")
resource_cache["db-primary"] = db_connection

print(f"Cache hit:  {resource_cache.get('db-primary')}")
print(f"Cache size: {len(resource_cache)}")
print()

# When the strong reference disappears, the cache entry cleans itself up
del db_connection
gc.collect()  # Force cleanup for demo purposes

print(f"After del:  {resource_cache.get('db-primary')}")
print(f"Cache size: {len(resource_cache)}")
print()

# ══════════════════════════════════════════════════════════════
# PART 2: Breaking a parent-child cycle with weakref.ref
# ══════════════════════════════════════════════════════════════

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.children = []
        self._parent_ref = None  # Will hold a weak reference, not a strong one

    def add_child(self, child_node):
        child_node._parent_ref = weakref.ref(self)  # Weak — child won't keep parent alive
        self.children.append(child_node)            # Strong — parent keeps children alive

    @property
    def parent(self):
        # Dereference the weak ref; returns None if parent was collected
        if self._parent_ref is None:
            return None
        return self._parent_ref()  # Calling a weakref returns the object or None

    def __repr__(self):
        return f"TreeNode({self.value})"

root = TreeNode("root")
child = TreeNode("child")
root.add_child(child)

print(f"child.parent = {child.parent}")
print(f"root.children = {root.children}")
print()

# ══════════════════════════════════════════════════════════════
# PART 3: __slots__ memory savings — measured
# ══════════════════════════════════════════════════════════════

class RegularPoint:
    """Standard class — every instance carries a full __dict__."""
    def __init__(self, x_coord, y_coord, z_coord):
        self.x_coord = x_coord
        self.y_coord = y_coord
        self.z_coord = z_coord

class SlottedPoint:
    """Slots class — fixed-size C array, no __dict__ overhead."""
    __slots__ = ('x_coord', 'y_coord', 'z_coord')

    def __init__(self, x_coord, y_coord, z_coord):
        self.x_coord = x_coord
        self.y_coord = y_coord
        self.z_coord = z_coord

regular = RegularPoint(1.0, 2.0, 3.0)
slotted = SlottedPoint(1.0, 2.0, 3.0)

regular_size = sys.getsizeof(regular) + sys.getsizeof(regular.__dict__)
slotted_size = sys.getsizeof(slotted)  # No __dict__ to add

print(f"RegularPoint size (object + __dict__): {regular_size} bytes")
print(f"SlottedPoint size (no __dict__):        {slotted_size} bytes")
print(f"Memory saved per instance:              {regular_size - slotted_size} bytes")
print()

# Scale that up to a realistic data pipeline with 1M points
num_instances = 1_000_000
savings_mb = (regular_size - slotted_size) * num_instances / (1024 ** 2)
print(f"Projected saving across {num_instances:,} instances: {savings_mb:.1f} MB")

import tracemalloc
import gc
import collections
import linecache

# ── Helper: pretty-print a tracemalloc diff ────────────────────────────────

def display_top_allocations(snapshot, key_type='lineno', limit=5):
    """Print the top N memory consumers from a tracemalloc snapshot."""
    stats = snapshot.statistics(key_type)
    print(f"{'Rank':<5} {'Size':>10} {'Count':>8}  Location")
    print("-" * 60)
    for rank, stat in enumerate(stats[:limit], start=1):
        frame = stat.traceback[0]
        # Fetch the actual source line for context
        source_line = linecache.getline(frame.filename, frame.lineno).strip()
        print(f"{rank:<5} {stat.size / 1024:>8.1f} KB {stat.count:>8}  "
              f"{frame.filename}:{frame.lineno}")
        print(f"      {'':>10} {'':>8}  -> {source_line}")
    print()

# ── Simulate a leaking registry (classic production pattern) ───────────────

class EventBus:
    """
    A naive event bus that never deregisters listeners.
    This is the #1 cause of Python service memory leaks.
    """
    _listeners: dict = {}

    @classmethod
    def register(cls, event_name, handler_func):
        cls._listeners.setdefault(event_name, []).append(handler_func)

    @classmethod
    def listener_count(cls):
        return sum(len(v) for v in cls._listeners.values())

# ── Take baseline snapshot ─────────────────────────────────────────────────

tracemalloc.start(depth=5)  # depth=5 captures 5 frames of stack context
gc.collect()                 # Clean slate before baseline

baseline_snapshot = tracemalloc.take_snapshot()
baseline_gc_counts = collections.Counter(
    type(obj).__name__ for obj in gc.get_objects()
)

print("=== Simulating 500 request cycles (leaking handlers each time) ===")

# Simulate a web server handling requests — each 'request' registers a
# new handler but the old ones are never removed
for request_number in range(500):
    def handle_user_event(event_data, req=request_number):
        """Handler closure — captures req, keeping it alive in the bus."""
        return f"request {req} handled {event_data}"

    EventBus.register("user.login", handle_user_event)

print(f"EventBus now holds {EventBus.listener_count()} handlers")
print()

# ── Take leak snapshot and diff ────────────────────────────────────────────

leak_snapshot = tracemalloc.take_snapshot()
leak_gc_counts = collections.Counter(
    type(obj).__name__ for obj in gc.get_objects()
)

print("=== Top memory allocations AFTER the leak ===")
display_top_allocations(leak_snapshot, limit=4)

print("=== Object count changes (GC-tracked objects) ===")
for type_name, count in (leak_gc_counts - baseline_gc_counts).most_common(5):
    print(f"  +{count:>6}  {type_name}")
print()

# ── Show the diff between snapshots ───────────────────────────────────────
print("=== Snapshot diff (new allocations since baseline) ===")
diff_stats = leak_snapshot.compare_to(baseline_snapshot, 'lineno')
for stat in diff_stats[:4]:
    print(stat)

tracemalloc.stop()

# ── The fix: use WeakSet so the bus doesn't prevent GC ────────────────────
print()
print("=== Fix: use weakref.WeakSet for listener registry ===")
import weakref

class SafeEventBus:
    _listeners: dict = {}

    @classmethod
    def register(cls, event_name, handler_func):
        if event_name not in cls._listeners:
            cls._listeners[event_name] = weakref.WeakSet()
        cls._listeners[event_name].add(handler_func)

    @classmethod
    def listener_count(cls):
        return sum(len(list(v)) for v in cls._listeners.values())

print("SafeEventBus uses WeakSet — handlers are released when they go out of scope.")