Fix heap stats and snapshots

Author: bpcreech

src/py_mini_racer/_context.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING, Any, NewType, Protocol, cast
 
 from py_mini_racer._dll import init_mini_racer
-from py_mini_racer._exc import JSEvalException, JSPromiseError, JSTimeoutException
+from py_mini_racer._exc import JSEvalException, JSPromiseError
 from py_mini_racer._types import (
     CancelableJSFunction,
     JSArray,
@@ -93,7 +93,7 @@ class Context:
 
     _dll: ctypes.CDLL
     _ctx: ContextType
-    _event_loop: asyncio.AbstractEventLoop
+    event_loop: asyncio.AbstractEventLoop
     _object_factory: ObjectFactory
     _next_async_callback_id: Iterator[int]
     _active_cancelable_mr_task_callbacks: dict[int, Callable[[ValueHandle], None]] = (
@@ -126,7 +126,7 @@ def handle_callback_from_v8(
         if callback_id == _UNCANCELABLE_TASK_CALLBACK_ID:
             self._non_cancelable_mr_task_results_queue.put(val_handle)
         else:
-            self._event_loop.call_soon_threadsafe(
+            self.event_loop.call_soon_threadsafe(
                 self._handle_callback_from_v8_on_event_loop, callback_id, val_handle
             )
 
@@ -179,9 +179,7 @@ async def await_promise(self, promise: JSPromise) -> PythonJSConvertedTypes:
             ),
         )
 
-        future: asyncio.Future[PythonJSConvertedTypes] = (
-            self._event_loop.create_future()
-        )
+        future: asyncio.Future[PythonJSConvertedTypes] = self.event_loop.create_future()
 
         def on_resolved(val_handle: ValueHandle) -> None:
             if future.cancelled():
@@ -315,24 +313,7 @@ def array_push(self, arr: JSArray, new_val: PythonJSConvertedTypes) -> None:
         )
 
     def are_we_running_on_the_mini_racer_event_loop(self) -> bool:
-        return get_running_loop_or_none() is self._event_loop
-
-    def run_coro(
-        self,
-        coro: Coroutine[Any, Any, PythonJSConvertedTypes],
-        timeout_sec: float | None = None,
-    ) -> PythonJSConvertedTypes:
-        assert not self.are_we_running_on_the_mini_racer_event_loop(), (
-            "Cannot run a blocking operation from our own event loop"
-        )
-
-        async def run() -> PythonJSConvertedTypes:
-            try:
-                return await asyncio.wait_for(coro, timeout=timeout_sec)
-            except asyncio.TimeoutError as e:
-                raise JSTimeoutException from e
-
-        return asyncio.run_coroutine_threadsafe(run(), self._event_loop).result()
+        return get_running_loop_or_none() is self.event_loop
 
     async def call_function_cancelable(
         self,
@@ -391,14 +372,22 @@ def was_soft_memory_limit_reached(self) -> bool:
     def low_memory_notification(self) -> None:
         self._dll.mr_low_memory_notification(self._ctx)
 
-    async def heap_stats(self) -> str:
-        return cast("str", await self._run_cancelable_mr_task(self._dll.mr_heap_stats))
+    def heap_stats(self) -> str:
+        return cast(
+            "str",
+            self._value_handle_to_python(
+                self._wrap_raw_handle(self._dll.mr_heap_stats(self._ctx))
+            ),
+        )
 
-    async def heap_snapshot(self) -> str:
+    def heap_snapshot(self) -> str:
         """Return a snapshot of the V8 isolate heap."""
 
         return cast(
-            "str", await self._run_cancelable_mr_task(self._dll.mr_heap_snapshot)
+            "str",
+            self._value_handle_to_python(
+                self._wrap_raw_handle(self._dll.mr_heap_snapshot(self._ctx))
+            ),
         )
 
     def value_count(self) -> int:
@@ -444,7 +433,7 @@ async def await_into_js_promise_resolvers(val_handle: ValueHandle) -> None:
                     err_maker(f"Error running Python function:\n{format_exc()}")
                 )
 
-        async with _make_task_set(self._event_loop) as task_set:
+        async with _make_task_set(self.event_loop) as task_set:
             with self._register_js_notification(
                 lambda val_handle: task_set.start_task(
                     await_into_js_promise_resolvers(val_handle)

src/py_mini_racer/_dll.py

@@ -108,16 +108,16 @@ def _build_dll_handle(dll_path: Path) -> ctypes.CDLL:  # noqa: PLR0915
 
     handle.mr_cancel_task.argtypes = [ctypes.c_uint64, ctypes.c_uint64]
 
-    handle.mr_heap_stats.argtypes = [ctypes.c_uint64, ctypes.c_uint64]
-    handle.mr_heap_stats.restype = ctypes.c_uint64
+    handle.mr_heap_stats.argtypes = [ctypes.c_uint64]
+    handle.mr_heap_stats.restype = RawValueHandle
 
     handle.mr_low_memory_notification.argtypes = [ctypes.c_uint64]
 
     handle.mr_make_js_callback.argtypes = [ctypes.c_uint64, ctypes.c_uint64]
     handle.mr_make_js_callback.restype = RawValueHandle
 
-    handle.mr_heap_snapshot.argtypes = [ctypes.c_uint64, ctypes.c_uint64]
-    handle.mr_heap_snapshot.restype = ctypes.c_uint64
+    handle.mr_heap_snapshot.argtypes = [ctypes.c_uint64]
+    handle.mr_heap_snapshot.restype = RawValueHandle
 
     handle.mr_get_identity_hash.argtypes = [ctypes.c_uint64, RawValueHandle]
     handle.mr_get_identity_hash.restype = RawValueHandle

src/py_mini_racer/_mini_racer.py

@@ -15,7 +15,7 @@
 
 from py_mini_racer._context import Context, ContextType, get_running_loop_or_none
 from py_mini_racer._dll import init_mini_racer, mr_callback_func
-from py_mini_racer._exc import WrongReturnTypeException
+from py_mini_racer._exc import JSTimeoutException, WrongReturnTypeException
 from py_mini_racer._objects import ObjectFactoryImpl
 from py_mini_racer._set_timeout import INSTALL_SET_TIMEOUT
 
@@ -146,20 +146,28 @@ def eval(
             # Système international d'unités use seconds.
             timeout_sec = timeout / 1000
 
-        assert self._ctx is not None
+        ctx = self._ctx
+        assert ctx is not None
 
-        if not self._ctx.are_we_running_on_the_mini_racer_event_loop():
-            return self._ctx.run_coro(
-                self._ctx.eval_cancelable(code), timeout_sec=timeout_sec
-            )
+        if not ctx.are_we_running_on_the_mini_racer_event_loop():
+
+            async def run() -> PythonJSConvertedTypes:
+                try:
+                    return await asyncio.wait_for(
+                        ctx.eval_cancelable(code), timeout=timeout_sec
+                    )
+                except asyncio.TimeoutError as e:
+                    raise JSTimeoutException from e
+
+            return asyncio.run_coroutine_threadsafe(run(), ctx.event_loop).result()
 
         assert timeout_sec is None, (
             "To apply a timeout in an async context, use "
             "`await asyncio.wait_for(mr.eval_cancelable(your_params), "
             "timeout=your_timeout)`"
         )
 
-        return self._ctx.eval(code)
+        return ctx.eval(code)
 
     async def eval_cancelable(self, code: str) -> PythonJSConvertedTypes:
         """Evaluate JavaScript code in the V8 isolate.
@@ -316,7 +324,13 @@ def heap_stats(self) -> Any:  # noqa: ANN401
         """Return the V8 isolate heap statistics."""
 
         assert self._ctx is not None
-        return self.json_impl.loads(self._ctx.run_coro(self._ctx.heap_stats()))
+        return self.json_impl.loads(self._ctx.heap_stats())
+
+    def heap_snapshot(self) -> Any:  # noqa: ANN401
+        """Return a snapshot of the V8 isolate heap."""
+
+        assert self._ctx is not None
+        return self.json_impl.loads(self._ctx.heap_snapshot())
 
 
 @contextmanager

src/py_mini_racer/_objects.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import ctypes
 from datetime import datetime, timezone
 from operator import index as op_index
@@ -15,6 +16,7 @@
     JSOOMException,
     JSParseException,
     JSTerminatedException,
+    JSTimeoutException,
     JSValueError,
 )
 from py_mini_racer._types import (
@@ -131,9 +133,19 @@ def __call__(
         timeout_sec: float | None = None,
     ) -> PythonJSConvertedTypes:
         if not self._ctx.are_we_running_on_the_mini_racer_event_loop():
-            return self._ctx.run_coro(
-                self._ctx.call_function_cancelable(self, *args, this=this), timeout_sec
-            )
+
+            async def run() -> PythonJSConvertedTypes:
+                try:
+                    return await asyncio.wait_for(
+                        self._ctx.call_function_cancelable(self, *args, this=this),
+                        timeout=timeout_sec,
+                    )
+                except asyncio.TimeoutError as e:
+                    raise JSTimeoutException from e
+
+            return asyncio.run_coroutine_threadsafe(
+                run(), self._ctx.event_loop
+            ).result()
 
         assert timeout_sec is None, (
             "To apply a timeout in an async context, use "
@@ -166,7 +178,15 @@ def get(self, *, timeout: float | None = None) -> PythonJSConvertedTypes:
             "In an async context, call `await promise` instead of promise.get()"
         )
 
-        return self._ctx.run_coro(self._ctx.await_promise(self), timeout_sec=timeout)
+        async def run() -> PythonJSConvertedTypes:
+            try:
+                return await asyncio.wait_for(
+                    self._ctx.await_promise(self), timeout=timeout
+                )
+            except asyncio.TimeoutError as e:
+                raise JSTimeoutException from e
+
+        return asyncio.run_coroutine_threadsafe(run(), self._ctx.event_loop).result()
 
     def __await__(self) -> Generator[Any, None, Any]:
         return self._ctx.await_promise(self).__await__()

src/v8_py_frontend/context.cc

@@ -125,20 +125,21 @@ void Context::CancelTask(uint64_t task_id) {
   cancelable_task_manager_.Cancel(task_id);
 }
 
-auto Context::HeapSnapshot(uint64_t callback_id) -> uint64_t {
-  return RunTask(
-      [this](v8::Isolate* isolate) {
-        return heap_reporter_.HeapSnapshot(isolate);
-      },
-      callback_id);
+auto Context::HeapSnapshot() -> BinaryValueHandle* {
+  return bv_registry_.Remember(isolate_manager_
+                                   .Run([this](v8::Isolate* isolate) mutable {
+                                     return heap_reporter_.HeapSnapshot(
+                                         isolate);
+                                   })
+                                   .get());
 }
 
-auto Context::HeapStats(uint64_t callback_id) -> uint64_t {
-  return RunTask(
-      [this](v8::Isolate* isolate) {
-        return heap_reporter_.HeapStats(isolate);
-      },
-      callback_id);
+auto Context::HeapStats() -> BinaryValueHandle* {
+  return bv_registry_.Remember(isolate_manager_
+                                   .Run([this](v8::Isolate* isolate) mutable {
+                                     return heap_reporter_.HeapStats(isolate);
+                                   })
+                                   .get());
 }
 
 auto Context::GetIdentityHash(BinaryValueHandle* obj_handle)

src/v8_py_frontend/context.h

@@ -41,8 +41,8 @@ class Context {
   template <typename... Params>
   auto AllocBinaryValue(Params&&... params) -> BinaryValueHandle*;
   void CancelTask(uint64_t task_id);
-  auto HeapSnapshot(uint64_t callback_id) -> uint64_t;
-  auto HeapStats(uint64_t callback_id) -> uint64_t;
+  auto HeapSnapshot() -> BinaryValueHandle*;
+  auto HeapStats() -> BinaryValueHandle*;
   auto Eval(BinaryValueHandle* code_handle,
 
             uint64_t callback_id) -> uint64_t;

src/v8_py_frontend/exports.cc

@@ -115,13 +115,13 @@ LIB_EXPORT void mr_cancel_task(uint64_t context_id, uint64_t task_id) {
   context->CancelTask(task_id);
 }
 
-LIB_EXPORT auto mr_heap_stats(uint64_t context_id,
-                              uint64_t callback_id) -> uint64_t {
+LIB_EXPORT auto mr_heap_stats(uint64_t context_id)
+    -> MiniRacer::BinaryValueHandle* {
   auto context = GetContext(context_id);
   if (!context) {
     return 0;
   }
-  return context->HeapStats(callback_id);
+  return context->HeapStats();
 }
 
 LIB_EXPORT void mr_set_hard_memory_limit(uint64_t context_id, size_t limit) {
@@ -273,13 +273,13 @@ LIB_EXPORT auto mr_call_function(uint64_t context_id,
                                callback_id);
 }
 
-LIB_EXPORT auto mr_heap_snapshot(uint64_t context_id,
-                                 uint64_t callback_id) -> uint64_t {
+LIB_EXPORT auto mr_heap_snapshot(uint64_t context_id)
+    -> MiniRacer::BinaryValueHandle* {
   auto context = GetContext(context_id);
   if (!context) {
     return 0;
   }
-  return context->HeapSnapshot(callback_id);
+  return context->HeapSnapshot();
 }
 
 LIB_EXPORT auto mr_value_count(uint64_t context_id) -> size_t {

src/v8_py_frontend/exports.h

@@ -209,8 +209,7 @@ LIB_EXPORT auto mr_array_push(uint64_t context_id,
 
 /** Cancel the given asynchronous task.
  *
- * (Such tasks are started by mr_eval, mr_call_function, mr_heap_stats, and
- * mr_heap_snapshot).
+ * (Such tasks are started by mr_eval and mr_call_function).
  **/
 LIB_EXPORT void mr_cancel_task(uint64_t context_id, uint64_t task_id);
 
@@ -248,26 +247,16 @@ LIB_EXPORT auto mr_call_function(uint64_t context_id,
 /** Get stats for the V8 heap.
  *
  * This function is intended for use in debugging only.
- *
- * This call is processed asynchronously and as such accepts a callback ID.
- * The callback ID and a MiniRacer::BinaryValueHandle* containing the
- * evaluation result are passed back to the callback upon completion. A task ID
- * is returned which can be passed back to mr_cancel_task to cancel evaluation.
  **/
-LIB_EXPORT auto mr_heap_stats(uint64_t context_id,
-                              uint64_t callback_id) -> uint64_t;
+LIB_EXPORT auto mr_heap_stats(uint64_t context_id)
+    -> MiniRacer::BinaryValueHandle*;
 
 /** Get a snapshot of the V8 heap.
  *
  * This function is intended for use in debugging only.
- *
- * This call is processed asynchronously and as such accepts a callback ID.
- * The callback ID and a MiniRacer::BinaryValueHandle* containing the
- * evaluation result are passed back to the callback upon completion. A task ID
- * is returned which can be passed back to mr_cancel_task to cancel evaluation.
  **/
-LIB_EXPORT auto mr_heap_snapshot(uint64_t context_id,
-                                 uint64_t callback_id) -> uint64_t;
+LIB_EXPORT auto mr_heap_snapshot(uint64_t context_id)
+    -> MiniRacer::BinaryValueHandle*;
 
 // NOLINTEND(bugprone-easily-swappable-parameters)
 

tests/test_heap.py

@@ -11,3 +11,12 @@ def test_heap_stats() -> None:
     assert mr.heap_stats()["total_heap_size"] > 0
 
     assert_no_v8_objects(mr)
+
+
+def test_heap_snapshot() -> None:
+    mr = MiniRacer()
+
+    assert mr.heap_snapshot()["edges"]
+    assert mr.heap_snapshot()["strings"]
+
+    assert_no_v8_objects(mr)