Repository: cocotb/cocotb

Issue:
Scheduler._task_context performance
While looking at @themperek's profiling log in #961, I noticed that `Scheduler._task_context` is taking a lot of time.

Investigating a little, the [contextlib docs](https://docs.python.org/3.5/library/contextlib.html?highlight=contextlib#contextlib.contextmanager) details say that using the `@contextmanager` decorator creates a new one-shot context manager on each call of the function, which involves some extra layers of wrappers due to the generic nature of **contextlib**.

This is used every time a task is scheduled. Investigate possible performance improvements from re-writing task context manager as a simpler dedicated context manager class.


Discussion:
Is it not possible to reuse the one-shot object? We could try just dropping a `@cache` decorator in there an see if that helps.
We need a new context manager each time in order to pass in the new task.

It might be faster to use a stack and push/pop the current task. It's worth investigating performance on the scheduling hot path.
That sucks. The fact we have to have a stack of scheduled tasks also sucks. You are suggesting we have a single context object that when `__enter__`ing we push and `__exit__`ing we pop so we aren't continuously making new context manager objects?
Are you sure this isn't a red herring? `_task_context` might be showing up as also including all the time where the slow code is _within_ the context.
Here are the relevant lines for `Scheduler.schedule` from the log. I'm not very familiar with `pprofile`.

```
Line #|      Hits|         Time| Time per hit|      %|Source code
------+----------+-------------+-------------+-------+-----------
   794|    337112|      3.14394|  9.32609e-06|  1.89%|        with self._task_context(coroutine):
(call)|    337112|      5.76703|  1.71072e-05|  3.47%|# /home/themperek/miniconda/lib/python3.7/contextlib.py:237 helper
(call)|    337112|      5.59323|  1.65916e-05|  3.37%|# /home/themperek/miniconda/lib/python3.7/contextlib.py:107 __enter__
```
And also it shows **contextlib.py** as using ~8% of the simulation time:
```
File: /home/themperek/miniconda/lib/python3.7/contextlib.py
File duration: 13.9257s (8.39%)
```
It is being executed on every scheduling operation. It looks like `contextlib` overhead in this case is 0.5s of that time, with the body of the manager taking 2.5s. While it's maybe not that contextlib is slow, there could be a more efficient implementation.

EDIT: And of course this all goes away on a scheduler rewrite/refactor.

Resolving patch:
diff --git a/src/cocotb/scheduler.py b/src/cocotb/scheduler.py
--- a/src/cocotb/scheduler.py
+++ b/src/cocotb/scheduler.py
@@ -40,7 +40,6 @@
 import warnings
 from collections import OrderedDict
 from collections.abc import Coroutine
-from contextlib import contextmanager
 from typing import Any, Callable, Union
 
 import cocotb
@@ -791,16 +790,6 @@ def _trigger_from_any(self, result) -> Trigger:
             "Did you forget to decorate with @cocotb.coroutine?"
         )
 
-    @contextmanager
-    def _task_context(self, task):
-        """Context manager for the currently running task."""
-        old_task = self._current_task
-        self._current_task = task
-        try:
-            yield
-        finally:
-            self._current_task = old_task
-
     def _schedule(self, task, trigger=None):
         """Schedule a task to execute.
 
@@ -809,7 +798,10 @@ def _schedule(self, task, trigger=None):
             trigger (cocotb.triggers.Trigger): The trigger that caused this
                 task to be scheduled.
         """
-        with self._task_context(task):
+        if self._current_task is not None:
+            raise InternalError("_schedule() called while another Task is executing")
+        try:
+            self._current_task = task
             if trigger is None:
                 send_outcome = outcomes.Value(None)
             else:
@@ -861,6 +853,8 @@ def _schedule(self, task, trigger=None):
                     if state == external_state.EXITED:
                         self._pending_threads.remove(ext)
                         self._pending_events.append(ext.event)
+        finally:
+            self._current_task = None
 
     def _finish_test(self, exc):
         self._abort_test(exc)


Test patch: