Monitor function calls inside python contexts

Question

Using python I would like to:

• monitor execution of some 'selected' functions (I would like to record function name, execution time, etc.)
• store these info only when functions are executed inside a context
• if these contexts are nested, info should be recorded in all parent contexts

I end up with the following code which monitor functions decorated with @monitor_decorator:

from dataclasses import dataclass

@dataclass
class MonitorRecord:
    function: str
    time: float

class MonitorContext:
    def __init__(self):
        self._records: list[MonitorRecord] = []

    def add_record(self, record: MonitorRecord) -> None:
        self._records.append(record)

    def __enter__(self) -> 'MonitorContext':
        handlers.register(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        handlers.delete(self)
        return

class MonitorHandlers:
    def __init__(self):
        self._handlers: list[MonitorContext] = []

    def register(self, handler: MonitorContext) -> None:
        self._handlers.append(handler)

    def delete(self, handler: MonitorContext) -> None:
        self._handlers.remove(handler)

    def add_record(self, record: MonitorRecord) -> None:
        for h in self._handlers:
            h.add_record(record)

handlers = MonitorHandlers()

def monitor_decorator(f):
    def _(*args, **kwargs):
        start = time.time()
        f(*args, **kwargs)
        handlers.add_record(
            MonitorRecord(
                function=f.__name__,
                time=time.time() - start,
            )
        )
    return _

This code works in simple cases like this:

@monitor_decorator
def run():
    time.sleep(5)

with MonitorContext() as m1:
    run()
    with MonitorContext() as m2:
        run()
        run()

print(len(m1._records))
print(len(m2._records))

Output:

3
2

But doesn't work with multithreading:

@monitor_decorator
def run():
    time.sleep(5)


def nested():
    with MonitorContext() as m:
        run()
    print(len(m._records))

with MonitorContext() as m1:
    threads = [threading.Thread(target=nested)
                   for i in range(10)]
    [t.start() for t in threads]
    [t.join() for t in threads]
print(len(m1.records))

Output:

1
2
3
4
5
6
7
9
8
10
10

Indeed, in this case, each thread adds its own context to the global variable handlers and the final result is that contexts inside a thread are monitoring other threads.

The problem arises from the global variable handlers, but so far I haven't found a simple and elegant alternative.

Answer 1

You can maintain a seperate handler list for the main thread and for other threads that don't share context. And then let other threads add records to the main thread context.

This solution works well in the case that you start all your threads with a shared context from the main thread. But does not work well in the case that a non main thread starts a context that then creates a second thread. Partly because python does not have a concept of a parent thread. only that of a main thread, a non daemon thread. and daemon threads.

The below implementation uses locks because iterating over a list of handlers and or deleting an element from a list is not thread safe. This might have performance implications.

import threading
from collections import UserList

class LocalList(threading.local, UserList):
    pass

class MonitorHandlers:
    def __init__(self):
        self._lock = threading.Lock()
        with self._lock:
            self._mainhandlers: list[MonitorContext] = []
        self._handlers: list[MonitorContext] = LocalList()

    def register(self, handler: MonitorContext) -> None:
        if threading.main_thread().ident == threading.get_ident():
            with self._lock:
                self._mainhandlers.append(handler)
        else:
            self._handlers.append(handler)

    def delete(self, handler: MonitorContext) -> None:
        if threading.main_thread().ident == threading.get_ident():
            with self._lock:
                self._mainhandlers.remove(handler)
        else:
            self._handlers.remove(handler)

    def add_record(self, record: MonitorRecord) -> None:
        for h in self._handlers:
            h.add_record(record)
        with self._lock:
            for h in self._mainhandlers:
                h.add_record(record)