From c2248afe4b0c60be8aa584a47f95c24a7fcea109 Mon Sep 17 00:00:00 2001
From: AN Long <aisk1988@gmail.com>
Date: Fri, 14 Jun 2024 16:33:25 +0800
Subject: [PATCH] Using critical section to make _socket.socket.close thread
 safe

---
 Modules/clinic/socketmodule.c.h            | 29 +++++++++++++++++++++-
 Modules/socketmodule.c                     | 21 ++++++++++------
 Tools/tsan/suppressions_free_threading.txt |  1 -
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/Modules/clinic/socketmodule.c.h b/Modules/clinic/socketmodule.c.h
index 3f4056efff2fec..7b0a3f8d4b1cc6 100644
--- a/Modules/clinic/socketmodule.c.h
+++ b/Modules/clinic/socketmodule.c.h
@@ -6,8 +6,35 @@ preserve
 #  include "pycore_gc.h"          // PyGC_Head
 #  include "pycore_runtime.h"     // _Py_ID()
 #endif
+#include "pycore_critical_section.h"// Py_BEGIN_CRITICAL_SECTION()
 #include "pycore_modsupport.h"    // _PyArg_UnpackKeywords()
 
+PyDoc_STRVAR(_socket_socket_close__doc__,
+"close($self, /)\n"
+"--\n"
+"\n"
+"close()\n"
+"\n"
+"Close the socket.  It cannot be used after this call.");
+
+#define _SOCKET_SOCKET_CLOSE_METHODDEF    \
+    {"close", (PyCFunction)_socket_socket_close, METH_NOARGS, _socket_socket_close__doc__},
+
+static PyObject *
+_socket_socket_close_impl(PySocketSockObject *s);
+
+static PyObject *
+_socket_socket_close(PySocketSockObject *s, PyObject *Py_UNUSED(ignored))
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(s);
+    return_value = _socket_socket_close_impl(s);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
+}
+
 static int
 sock_initobj_impl(PySocketSockObject *self, int family, int type, int proto,
                   PyObject *fdobj);
@@ -259,4 +286,4 @@ _socket_socket_if_nametoindex(PySocketSockObject *self, PyObject *arg)
 #ifndef _SOCKET_SOCKET_IF_NAMETOINDEX_METHODDEF
     #define _SOCKET_SOCKET_IF_NAMETOINDEX_METHODDEF
 #endif /* !defined(_SOCKET_SOCKET_IF_NAMETOINDEX_METHODDEF) */
-/*[clinic end generated code: output=eb37b5d88a1e4661 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=6037e47b012911c5 input=a9049054013a1b77]*/
diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c
index 0626d7934983db..b7508bc371675f 100644
--- a/Modules/socketmodule.c
+++ b/Modules/socketmodule.c
@@ -3328,8 +3328,19 @@ sockets the address is a tuple (ifname, proto [,pkttype [,hatype [,addr]]])");
    Set the file descriptor to -1 so operations tried subsequently
    will surely fail. */
 
+/*[clinic input]
+@critical_section
+_socket.socket.close
+    self as s: self(type="PySocketSockObject *")
+
+close()
+
+Close the socket.  It cannot be used after this call.
+[clinic start generated code]*/
+
 static PyObject *
-sock_close(PySocketSockObject *s, PyObject *Py_UNUSED(ignored))
+_socket_socket_close_impl(PySocketSockObject *s)
+/*[clinic end generated code: output=038b2418e07f6f6c input=9839a261e05bcb97]*/
 {
     SOCKET_T fd;
     int res;
@@ -3354,11 +3365,6 @@ sock_close(PySocketSockObject *s, PyObject *Py_UNUSED(ignored))
     Py_RETURN_NONE;
 }
 
-PyDoc_STRVAR(sock_close_doc,
-"close()\n\
-\n\
-Close the socket.  It cannot be used after this call.");
-
 static PyObject *
 sock_detach(PySocketSockObject *s, PyObject *Py_UNUSED(ignored))
 {
@@ -5115,8 +5121,7 @@ static PyMethodDef sock_methods[] = {
     {"bind",              (PyCFunction)sock_bind, METH_O,
                       bind_doc},
 #endif
-    {"close",             (PyCFunction)sock_close, METH_NOARGS,
-                      sock_close_doc},
+    _SOCKET_SOCKET_CLOSE_METHODDEF
 #ifdef HAVE_CONNECT
     {"connect",           (PyCFunction)sock_connect, METH_O,
                       connect_doc},
diff --git a/Tools/tsan/suppressions_free_threading.txt b/Tools/tsan/suppressions_free_threading.txt
index 05ceaf438b6353..d020d1be194ddf 100644
--- a/Tools/tsan/suppressions_free_threading.txt
+++ b/Tools/tsan/suppressions_free_threading.txt
@@ -57,7 +57,6 @@ race_top:PyInterpreterState_ThreadHead
 race_top:_PyObject_TryGetInstanceAttribute
 race_top:PyThreadState_Next
 race_top:PyUnstable_InterpreterFrame_GetLine
-race_top:sock_close
 race_top:tstate_delete_common
 race_top:tstate_is_freed
 race_top:type_modified_unlocked