bpo-38880: List interpreters associated with a channel end

Include/cpython/pystate.h

@@ -147,6 +147,8 @@ struct _ts {
    The caller must hold the GIL.*/
 PyAPI_FUNC(PyInterpreterState *) _PyInterpreterState_Get(void);
 
+PyAPI_FUNC(PyInterpreterState *) _PyInterpreterState_LookUpID(PY_INT64_T requested_id);
+
 PyAPI_FUNC(int) _PyState_AddModule(PyObject*, struct PyModuleDef*);
 PyAPI_FUNC(PyThreadState *) _PyThreadState_Prealloc(PyInterpreterState *);
 

Include/pyport.h

@@ -72,10 +72,34 @@ Used in:  Py_SAFE_DOWNCAST
 
 #define PY_UINT32_T uint32_t
 #define PY_UINT64_T uint64_t
+#ifdef HAVE_STDINT_H
+#include <stdint.h>
+#define PY_UINT32_T_MIN UINT32_MIN
+#define PY_UINT32_T_MAX UINT32_MAX
+#define PY_UINT64_T_MIN UINT64_MIN
+#define PY_UINT64_T_MAX UINT64_MAX
+#else
+#define PY_UINT32_T_MIN 0
+#define PY_UINT32_T_MAX ((uint32_t)-1)
+#define PY_UINT64_T_MIN 0
+#define PY_UINT64_T_MAX ((uint64_t)-1)
+#endif
 
 /* Signed variants of the above */
 #define PY_INT32_T int32_t
 #define PY_INT64_T int64_t
+#ifdef HAVE_STDINT_H
+#include <stdint.h>
+#define PY_INT32_T_MIN INT32_MIN
+#define PY_INT32_T_MAX INT32_MAX
+#define PY_INT64_T_MIN INT64_MIN
+#define PY_INT64_T_MAX INT64_MAX
+#else
+#define PY_INT32_T_MAX ((int32_t)(((uint32_t)-1)>>1))
+#define PY_INT32_T_MIN (-PY_INT32_T_MAX-1)
+#define PY_INT64_T_MAX ((int64_t)(((uint64_t)-1)>>1))
+#define PY_INT64_T_MIN (-PY_INT64_T_MAX-1)
+#endif
 
 /* If PYLONG_BITS_IN_DIGIT is not defined then we'll use 30-bit digits if all
    the necessary integer types are available, and we're on a 64-bit platform

Lib/test/test__xxsubinterpreters.py

@@ -1203,6 +1203,58 @@ def test_ids_global(self):
 
         self.assertEqual(cid2, int(cid1) + 1)
 
+    def _assert_interpreters_returned(self, cid, send, recv):
+        actual = set(interpreters.channel_list_interpreters(cid, send=True))
+        self.assertEqual(actual, set(send))
+        actual = set(interpreters.channel_list_interpreters(cid, recv=True))
+        self.assertEqual(actual, set(recv))
+
+    def test_channel_list_interpreters_empty(self):
+        # Test for channel with no associated interpreters.
+        cid = interpreters.channel_create()
+        self._assert_interpreters_returned(cid, [], [])
+
+    def test_channel_list_interpreters_mainline(self):
+        interp0 = interpreters.get_main()
+        cid = interpreters.channel_create()
+        interpreters.channel_send(cid, "send")
+        # Test for a channel that has one end associated to an interpreter.
+        self._assert_interpreters_returned(
+            cid, [interp0], [])
+
+        interp1 = interpreters.create()
+        _run_output(interp1, dedent(f"""
+            import _xxsubinterpreters as _interpreters
+            obj = _interpreters.channel_recv({cid})
+            """))
+        # Test for channel that has boths ends associated to an interpreter.
+        self._assert_interpreters_returned(
+            cid, [interp0], [interp1])
+
+    def test_channel_list_interpreters_multiple(self):
+        # Test for channel with both ends associated to many interpreters.
+        interp0 = interpreters.get_main()
+        interp1 = interpreters.create()
+        interp2 = interpreters.create()
+        interp3 = interpreters.create()
+        cid = interpreters.channel_create()
+
+        interpreters.channel_send(cid, "send")
+        _run_output(interp1, dedent(f"""
+            import _xxsubinterpreters as _interpreters
+            obj = _interpreters.channel_send({cid}, "send")
+            """))
+        _run_output(interp2, dedent(f"""
+            import _xxsubinterpreters as _interpreters
+            obj = _interpreters.channel_recv({cid})
+            """))
+        _run_output(interp3, dedent(f"""
+            import _xxsubinterpreters as _interpreters
+            obj = _interpreters.channel_recv({cid})
+            """))
+        self._assert_interpreters_returned(
+            cid, [interp0, interp1], [interp3, interp2])
+
     ####################
 
     def test_send_recv_main(self):
@@ -1515,6 +1567,27 @@ def test_close_used_multiple_times_by_single_user(self):
         with self.assertRaises(interpreters.ChannelClosedError):
             interpreters.channel_recv(cid)
 
+    def test_channel_list_interpreters_invalid_channel(self):
+        cid = interpreters.channel_create()
+        # Test for invalid channel ID.
+        with self.assertRaises(interpreters.ChannelNotFoundError):
+            interpreters.channel_list_interpreters(1000, send=True)
+
+        interpreters.channel_close(cid)
+        # Test for a channel that has been closed.
+        with self.assertRaises(interpreters.ChannelClosedError):
+            interpreters.channel_list_interpreters(cid, send=True)
+
+    def test_channel_list_interpreters_invalid_args(self):
+        # Tests for invalid arguments passed to the API.
+        cid = interpreters.channel_create()
+
+        with self.assertRaises(ValueError):
+            interpreters.channel_list_interpreters(cid)
+
+        with self.assertRaises(ValueError):
+            interpreters.channel_list_interpreters(cid, send=True, recv=True)
+
 
 class ChannelReleaseTests(TestBase):
 

Misc/ACKS

@@ -444,6 +444,7 @@ Rodolpho Eckhardt
 Ulrich Eckhardt
 David Edelsohn
 John Edmonds
+Benjamin Edwards
 Grant Edwards
 Zvi Effron
 John Ehresman
@@ -557,6 +558,7 @@ Lars Marius Garshol
 Jake Garver
 Dan Gass
 Andrew Gaul
+Lewis Gaul
 Matthieu Gautier
 Stephen M. Gava
 Xavier de Gaye

Modules/_xxsubinterpretersmodule.c

@@ -627,6 +627,36 @@ _channelends_associate(_channelends *ends, int64_t interp, int send)
     return 0;
 }
 
+static int64_t *
+_channelends_list_interpreters(_channelends *ends, int64_t *count, int send)
+{
+    int64_t *ids = NULL;
+    int64_t numopen;
+
+    assert(ends != NULL);
+    assert(count != NULL);
+    numopen = send ? ends->numsendopen : ends->numrecvopen;
+
+    if (numopen >= PY_INT64_T_MAX) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "too many interpreters using the channel");
+        goto done;
+    }
+    ids = PyMem_NEW(int64_t, (Py_ssize_t)numopen);
+    if (ids == NULL) {
+        PyErr_SetNone(PyExc_MemoryError);
+        goto done;
+    }
+    _channelend *ref = send ? ends->send : ends->recv;
+    for (int64_t i=0; ref != NULL; ref = ref->next, i++) {
+        ids[i] = ref->interp;
+    }
+    *count = numopen;
+
+done:
+    return ids;
+}
+
 static int
 _channelends_is_open(_channelends *ends)
 {
@@ -1178,7 +1208,7 @@ _channels_list_all(_channels *channels, int64_t *count)
     int64_t *cids = NULL;
     PyThread_acquire_lock(channels->mutex, WAIT_LOCK);
     int64_t numopen = channels->numopen;
-    if (numopen >= PY_SSIZE_T_MAX) {
+    if (numopen >= PY_INT64_T_MAX) {
         PyErr_SetString(PyExc_RuntimeError, "too many channels open");
         goto done;
     }
@@ -2323,6 +2353,74 @@ PyDoc_STRVAR(channel_list_all_doc,
 \n\
 Return the list of all IDs for active channels.");
 
+
+static PyObject *
+channel_list_interpreters(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *kwlist[] = {"cid", "send", "recv", NULL};
+    int64_t cid;
+    _PyChannelState *chan;
+    int send = 0;           /* Send end? */
+    int recv = 0;           /* Receive end? */
+    int64_t *ids = NULL;    /* Array of interpreter IDs */
+    int64_t count = 0;      /* Number of interpreters to return */
+    PyObject *ret = NULL;   /* Python list of interpreter IDs */
+
+    if (!PyArg_ParseTupleAndKeywords(
+            args, kwds, "O&|$pp:channel_list_interpreters",
+            kwlist, channel_id_converter, &cid, &send, &recv)) {
+        return NULL;
+    }
+
+    if ((send && recv) || (!send && !recv)) {
+        PyErr_SetString(PyExc_ValueError,
+                        "Specify exactly one of send or recv");
+        goto except;
+    }
+
+    chan = _channels_lookup(&_globals.channels, cid, NULL);
+    if (chan == NULL) {
+        goto except;
+    }
+
+    ids = _channelends_list_interpreters(chan->ends, &count, send);
+    if (ids == NULL) {
+        goto except;
+    }
+
+    ret = PyList_New((Py_ssize_t)count);
+    if (ret == NULL) {
+        goto except;
+    }
+
+    for (int64_t i=0; i < count; i++) {
+        PyInterpreterState *interp = _PyInterpreterState_LookUpID(ids[i]);
+        PyObject *id_obj = _PyInterpreterState_GetIDObject(interp);
+        if (id_obj == NULL) {
+            goto except;
+        }
+        PyList_SET_ITEM(ret, i, id_obj);
+    }
+
+    goto finally;
+
+except:
+    Py_XDECREF(ret);
+    ret = NULL;
+
+finally:
+    PyMem_Free(ids);
+    return ret;
+}
+
+PyDoc_STRVAR(channel_list_interpreters_doc,
+"channel_list_interpreters(cid, *, send=False, recv=False) -> [id]\n\
+\n\
+Return the list of all interpreter IDs associated with the channel\n\
+Exactly one of 'send' or 'recv' should be True, corresponding to the end of\n\
+the channel to list interpreters for.");
+
+
 static PyObject *
 channel_send(PyObject *self, PyObject *args, PyObject *kwds)
 {
@@ -2476,6 +2574,8 @@ static PyMethodDef module_functions[] = {
      METH_VARARGS | METH_KEYWORDS, channel_destroy_doc},
     {"channel_list_all",          channel_list_all,
      METH_NOARGS, channel_list_all_doc},
+    {"channel_list_interpreters", (PyCFunction)(void(*)(void))channel_list_interpreters,
+     METH_VARARGS | METH_KEYWORDS, channel_list_interpreters_doc},
     {"channel_send",              (PyCFunction)(void(*)(void))channel_send,
      METH_VARARGS | METH_KEYWORDS, channel_send_doc},
     {"channel_recv",              (PyCFunction)(void(*)(void))channel_recv,