From c4afd17ea6098a5b7b3686d233e9343ac65d632a Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Wed, 5 Jun 2024 21:56:01 -0700
Subject: [PATCH] FIX Check if future is done before trying to set a result on
 it (#4837)

Otherwise, the future raises invalid state error
---
 docs/project/changelog.md         |  7 ++++++-
 src/core/jsproxy.c                | 14 ++++++++++++--
 src/py/_pyodide/_future_helper.py | 14 ++++++++++++++
 src/tests/test_asyncio.py         | 22 ++++++++++++++++++++++
 4 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 src/py/_pyodide/_future_helper.py

diff --git a/docs/project/changelog.md b/docs/project/changelog.md
index 6cf6169a2b1..6edc4c2d223 100644
--- a/docs/project/changelog.md
+++ b/docs/project/changelog.md
@@ -24,9 +24,14 @@ myst:
   stack switching is disabled.
   {pr}`4817`
 
-- {{ Fix }} Resolved an issue where string keys in `PyProxyJsonAdaptor` were unexpectedly cast to numbers.
+- {{ Fix }} Resolved an issue where string keys in `PyProxyJsonAdaptor` were
+  unexpectedly cast to numbers.
   {pr}`4825`
 
+- {{ Fix }} When a `Future` connected to a `Promise` is cancelled, don't raise
+  `InvalidStateError`.
+  {pr}`4837`
+
 ## Version 0.26.0
 
 _May 27, 2024_
diff --git a/src/core/jsproxy.c b/src/core/jsproxy.c
index a57af6e5628..5f064e47024 100644
--- a/src/core/jsproxy.c
+++ b/src/core/jsproxy.c
@@ -2637,6 +2637,8 @@ wrap_promise(JsVal promise, JsVal done_callback, PyObject* js2py_converter)
 {
   bool success = false;
   PyObject* loop = NULL;
+  PyObject* helpers_mod = NULL;
+  PyObject* helpers = NULL;
   PyObject* set_result = NULL;
   PyObject* set_exception = NULL;
 
@@ -2648,9 +2650,15 @@ wrap_promise(JsVal promise, JsVal done_callback, PyObject* js2py_converter)
   result = _PyObject_CallMethodIdNoArgs(loop, &PyId_create_future);
   FAIL_IF_NULL(result);
 
-  set_result = _PyObject_GetAttrId(result, &PyId_set_result);
+  helpers_mod = PyImport_ImportModule("_pyodide._future_helper");
+  FAIL_IF_NULL(helpers_mod);
+  _Py_IDENTIFIER(get_future_resolvers);
+  helpers = _PyObject_CallMethodIdOneArg(
+    helpers_mod, &PyId_get_future_resolvers, result);
+  FAIL_IF_NULL(helpers);
+  set_result = Py_XNewRef(PyTuple_GetItem(helpers, 0));
   FAIL_IF_NULL(set_result);
-  set_exception = _PyObject_GetAttrId(result, &PyId_set_exception);
+  set_exception = Py_XNewRef(PyTuple_GetItem(helpers, 1));
   FAIL_IF_NULL(set_exception);
 
   promise = JsvPromise_Resolve(promise);
@@ -2663,6 +2671,8 @@ wrap_promise(JsVal promise, JsVal done_callback, PyObject* js2py_converter)
   success = true;
 finally:
   Py_CLEAR(loop);
+  Py_CLEAR(helpers_mod);
+  Py_CLEAR(helpers);
   Py_CLEAR(set_result);
   Py_CLEAR(set_exception);
   if (!success) {
diff --git a/src/py/_pyodide/_future_helper.py b/src/py/_pyodide/_future_helper.py
new file mode 100644
index 00000000000..37b072c7500
--- /dev/null
+++ b/src/py/_pyodide/_future_helper.py
@@ -0,0 +1,14 @@
+def set_result(fut, val):
+    if fut.done():
+        return
+    fut.set_result(val)
+
+
+def set_exception(fut, val):
+    if fut.done():
+        return
+    fut.set_exception(val)
+
+
+def get_future_resolvers(fut):
+    return (set_result.__get__(fut), set_exception.__get__(fut))
diff --git a/src/tests/test_asyncio.py b/src/tests/test_asyncio.py
index 8e6a42763ef..161c6c637d0 100644
--- a/src/tests/test_asyncio.py
+++ b/src/tests/test_asyncio.py
@@ -2,6 +2,7 @@
 import time
 
 import pytest
+from pytest_pyodide import run_in_pyodide
 
 from pyodide.code import eval_code_async
 
@@ -421,3 +422,24 @@ async def temp():
         return (!!packages.packages) && (!!packages.info);
         """
     )
+
+
+@run_in_pyodide
+async def inner_test_cancellation(selenium):
+    from asyncio import ensure_future, sleep
+
+    from js import fetch
+
+    async def f():
+        while True:
+            await fetch("/")
+
+    fut = ensure_future(f())
+    await sleep(0.01)
+    fut.cancel()
+    await sleep(0.1)
+
+
+def test_cancellation(selenium):
+    inner_test_cancellation(selenium)
+    assert "InvalidStateError" not in selenium.logs