Use separate timer for function call; not thread safe

Author: kesmit13

singlestoredb/functions/ext/asgi.py

@@ -319,7 +319,6 @@ def build_udf_endpoint(
 
         async def do_func(
             cancel_event: threading.Event,
-            finished_event: threading.Event,
             timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
@@ -333,7 +332,6 @@ async def do_func(
                         out.append(await func(*row))
                     else:
                         out.append(func(*row))
-            finished_event.set()
             return row_ids, list(zip(out))
 
         return do_func
@@ -367,7 +365,6 @@ def build_vector_udf_endpoint(
 
     async def do_func(
         cancel_event: threading.Event,
-        finished_event: threading.Event,
         timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
@@ -391,7 +388,6 @@ async def do_func(
                 else:
                     out = func()
 
-        finished_event.set()
         cancel_on_event(cancel_event)
 
         # Single masked value
@@ -434,7 +430,6 @@ def build_tvf_endpoint(
 
         async def do_func(
             cancel_event: threading.Event,
-            finished_event: threading.Event,
             timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
@@ -452,7 +447,6 @@ async def do_func(
                         res = func(*row)
                     out.extend(as_list_of_tuples(res))
                     out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
-            finished_event.set()
             return out_ids, out
 
         return do_func
@@ -485,7 +479,6 @@ def build_vector_tvf_endpoint(
 
     async def do_func(
         cancel_event: threading.Event,
-        finished_event: threading.Event,
         timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
@@ -517,8 +510,6 @@ async def do_func(
                 else:
                     func_res = func()
 
-        finished_event.set()
-
         res = get_dataframe_columns(func_res)
 
         cancel_on_event(cancel_event)
@@ -981,6 +972,12 @@ async def __call__(
                 datetime.timezone.utc,
             ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
         )
+        call_timer = Timer(
+            id=request_id,
+            timestamp=datetime.datetime.now(
+                datetime.timezone.utc,
+            ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
+        )
 
         assert scope['type'] == 'http'
 
@@ -997,6 +994,7 @@ async def __call__(
         func_endpoint = self.endpoints.get(func_name)
 
         timer.metadata['function'] = func_name.decode('utf-8') if func_name else ''
+        call_timer.metadata['function'] = timer.metadata['function']
 
         func = None
         func_info: Dict[str, Any] = {}
@@ -1037,23 +1035,18 @@ async def __call__(
                 result = []
 
                 cancel_event = threading.Event()
-                finished_event = threading.Event()
-
-                # Async functions don't need to set the finished event
-                if func_info['is_async']:
-                    finished_event.set()
 
                 with timer('parse_input'):
                     inputs = input_handler['load'](  # type: ignore
                         func_info['colspec'], b''.join(data),
                     )
 
                 func_task = asyncio.create_task(
-                    func(cancel_event, finished_event, timer, *inputs)
+                    func(cancel_event, call_timer, *inputs)
                     if func_info['is_async']
                     else to_thread(
                         lambda: asyncio.run(
-                            func(cancel_event, finished_event, timer, *inputs),
+                            func(cancel_event, call_timer, *inputs),
                         ),
                     ),
                 )
@@ -1073,9 +1066,6 @@ async def __call__(
 
                 await cancel_all_tasks(pending)
 
-                # Make sure threads finish before we proceed
-                finished_event.wait()
-
                 for task in done:
                     if task is disconnect_task:
                         cancel_event.set()
@@ -1163,6 +1153,9 @@ async def __call__(
             out['body'] = body
             await send(out)
 
+        for k, v in call_timer.metrics.items():
+            timer.metrics[k] = v
+
         timer.finish()
 
     def _create_link(

singlestoredb/functions/ext/timer.py

@@ -2,7 +2,7 @@
 import time
 from typing import Any
 from typing import Dict
-from typing import List
+from typing import Optional
 
 from . import utils
 
@@ -46,93 +46,52 @@ class Timer:
     print(timer.metrics)
     # {'receive_data': 0.1, 'parse_input': 0.2, 'inner_operation': 0.05,
     #  'call_function': 0.35, 'total': 0.65}
+
     """
 
     def __init__(self, **kwargs: Any) -> None:
-        """
-        Initialize the Timer.
-
-        Parameters
-        ----------
-        metrics : Dict[str, float]
-            Dictionary to store the timing results
-
-        """
         self.metadata: Dict[str, Any] = kwargs
         self.metrics: Dict[str, float] = dict()
-        self._stack: List[Dict[str, Any]] = []
+        self.entries: Dict[str, float] = dict()
+        self._current_key: Optional[str] = None
         self.start_time = time.perf_counter()
 
     def __call__(self, key: str) -> 'Timer':
-        """
-        Set the key for the next context manager usage.
-
-        Parameters
-        ----------
-        key : str
-            The key to store the execution time under
-
-        Returns
-        -------
-        Timer
-            Self, to be used as context manager
-
-        """
         self._current_key = key
         return self
 
     def __enter__(self) -> 'Timer':
-        """Enter the context manager and start timing."""
-        if not hasattr(self, '_current_key'):
+        if self._current_key is None:
             raise ValueError(
                 "No key specified. Use timer('key_name') as context manager.",
             )
-
-        # Push current timing info onto stack
-        timing_info = {
-            'key': self._current_key,
-            'start_time': time.perf_counter(),
-        }
-        self._stack.append(timing_info)
-
-        # Clear current key for next use
-        delattr(self, '_current_key')
-
+        self.entries[self._current_key] = time.perf_counter()
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        """Exit the context manager and store the elapsed time."""
-        if not self._stack:
-            return
-
-        # Pop the current timing from stack
-        timing_info = self._stack.pop()
-        elapsed = time.perf_counter() - timing_info['start_time']
-        self.metrics.setdefault(timing_info['key'], 0)
-        self.metrics[timing_info['key']] += elapsed
+        key = self._current_key
+        if key and key in self.entries:
+            start = self.entries.pop(key)
+            elapsed = time.perf_counter() - start
+            self.metrics[key] = elapsed
+        self._current_key = None
 
     async def __aenter__(self) -> 'Timer':
-        """Async enter for async context manager support."""
         return self.__enter__()
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        """Async exit for async context manager support."""
         self.__exit__(exc_type, exc_val, exc_tb)
 
+    def reset(self) -> None:
+        self.metrics.clear()
+        self.entries.clear()
+        self._current_key = None
+
     def finish(self) -> None:
         """Finish the current timing context and store the elapsed time."""
-        if self._stack:
-            raise RuntimeError('finish() called within a `with` block.')
-
         self.metrics['total'] = time.perf_counter() - self.start_time
-
         self.log_metrics()
 
-    def reset(self) -> None:
-        """Clear all stored times and reset the stack."""
-        self.metrics.clear()
-        self._stack.clear()
-
     def log_metrics(self) -> None:
         if self.metadata.get('function'):
             result = dict(type='function_metrics', **self.metadata, **self.metrics)