1 files changed, 31 insertions, 6 deletions

diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c
index bb5edc3..392a56d 100644
--- a/Modules/socketmodule.c
+++ b/Modules/socketmodule.c
@@ -7314,20 +7314,39 @@ os_init(void)
 }
 #endif
 
+static int
+sock_capi_traverse(PyObject *capsule, visitproc visit, void *arg)
+{
+    PySocketModule_APIObject *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
+    assert(capi != NULL);
+    Py_VISIT(capi->Sock_Type);
+    return 0;
+}
+
+static int
+sock_capi_clear(PyObject *capsule)
+{
+    PySocketModule_APIObject *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
+    assert(capi != NULL);
+    Py_CLEAR(capi->Sock_Type);
+    return 0;
+}
+
 static void
-sock_free_api(PySocketModule_APIObject *capi)
+sock_capi_free(PySocketModule_APIObject *capi)
 {
-    Py_DECREF(capi->Sock_Type);
+    Py_XDECREF(capi->Sock_Type);  // sock_capi_free() can clear it
     Py_DECREF(capi->error);
     Py_DECREF(capi->timeout_error);
     PyMem_Free(capi);
 }
 
 static void
-sock_destroy_api(PyObject *capsule)
+sock_capi_destroy(PyObject *capsule)
 {
     void *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
-    sock_free_api(capi);
+    assert(capi != NULL);
+    sock_capi_free(capi);
 }
 
 static PySocketModule_APIObject *
@@ -7432,11 +7451,17 @@ socket_exec(PyObject *m)
     }
     PyObject *capsule = PyCapsule_New(capi,
                                       PySocket_CAPSULE_NAME,
-                                      sock_destroy_api);
+                                      sock_capi_destroy);
     if (capsule == NULL) {
-        sock_free_api(capi);
+        sock_capi_free(capi);
         goto error;
     }
+    if (_PyCapsule_SetTraverse(capsule,
+                               sock_capi_traverse, sock_capi_clear) < 0) {
+        sock_capi_free(capi);
+        goto error;
+    }
+
     if (PyModule_Add(m, PySocket_CAPI_NAME, capsule) < 0) {
         goto error;
     }