Skip to content

Commit 513c89d

Browse files
vstinnervstinner
authored
gh-108240: Add _PyCapsule_SetTraverse() internal function (#108339)
The _socket extension uses _PyCapsule_SetTraverse() to visit and clear the socket type in the garbage collector. So the _socket.socket type can be cleared in some corner cases when it wasn't possible before.
1 parent b6be188 commit 513c89d

File tree

3 files changed

+135
-72
lines changed

3 files changed

+135
-72
lines changed

‎Include/pycapsule.h

+4
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ PyAPI_FUNC(int) PyCapsule_SetName(PyObject *capsule, const char *name);
4848

4949
PyAPI_FUNC(int) PyCapsule_SetContext(PyObject *capsule, void *context);
5050

51+
#ifdef Py_BUILD_CORE
52+
PyAPI_FUNC(int) _PyCapsule_SetTraverse(PyObject *op, traverseproc traverse_func, inquiry clear_func);
53+
#endif
54+
5155
PyAPI_FUNC(void *) PyCapsule_Import(
5256
const char *name, /* UTF-8 encoded string */
5357
int no_block);

‎Modules/socketmodule.c

+31-6
Original file line numberDiff line numberDiff line change
@@ -7314,20 +7314,39 @@ os_init(void)
73147314
}
73157315
#endif
73167316

7317+
static int
7318+
sock_capi_traverse(PyObject *capsule, visitproc visit, void *arg)
7319+
{
7320+
PySocketModule_APIObject *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
7321+
assert(capi != NULL);
7322+
Py_VISIT(capi->Sock_Type);
7323+
return 0;
7324+
}
7325+
7326+
static int
7327+
sock_capi_clear(PyObject *capsule)
7328+
{
7329+
PySocketModule_APIObject *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
7330+
assert(capi != NULL);
7331+
Py_CLEAR(capi->Sock_Type);
7332+
return 0;
7333+
}
7334+
73177335
static void
7318-
sock_free_api(PySocketModule_APIObject *capi)
7336+
sock_capi_free(PySocketModule_APIObject *capi)
73197337
{
7320-
Py_DECREF(capi->Sock_Type);
7338+
Py_XDECREF(capi->Sock_Type); // sock_capi_free() can clear it
73217339
Py_DECREF(capi->error);
73227340
Py_DECREF(capi->timeout_error);
73237341
PyMem_Free(capi);
73247342
}
73257343

73267344
static void
7327-
sock_destroy_api(PyObject *capsule)
7345+
sock_capi_destroy(PyObject *capsule)
73287346
{
73297347
void *capi = PyCapsule_GetPointer(capsule, PySocket_CAPSULE_NAME);
7330-
sock_free_api(capi);
7348+
assert(capi != NULL);
7349+
sock_capi_free(capi);
73317350
}
73327351

73337352
static PySocketModule_APIObject *
@@ -7432,11 +7451,17 @@ socket_exec(PyObject *m)
74327451
}
74337452
PyObject *capsule = PyCapsule_New(capi,
74347453
PySocket_CAPSULE_NAME,
7435-
sock_destroy_api);
7454+
sock_capi_destroy);
74367455
if (capsule == NULL) {
7437-
sock_free_api(capi);
7456+
sock_capi_free(capi);
74387457
goto error;
74397458
}
7459+
if (_PyCapsule_SetTraverse(capsule,
7460+
sock_capi_traverse, sock_capi_clear) < 0) {
7461+
sock_capi_free(capi);
7462+
goto error;
7463+
}
7464+
74407465
if (PyModule_Add(m, PySocket_CAPI_NAME, capsule) < 0) {
74417466
goto error;
74427467
}

‎Objects/capsule.c

+100-66
Original file line numberDiff line numberDiff line change
@@ -9,18 +9,28 @@ typedef struct {
99
const char *name;
1010
void *context;
1111
PyCapsule_Destructor destructor;
12+
traverseproc traverse_func;
13+
inquiry clear_func;
1214
} PyCapsule;
1315

1416

1517

1618
static int
17-
_is_legal_capsule(PyCapsule *capsule, const char *invalid_capsule)
19+
_is_legal_capsule(PyObject *op, const char *invalid_capsule)
1820
{
19-
if (!capsule || !PyCapsule_CheckExact(capsule) || capsule->pointer == NULL) {
20-
PyErr_SetString(PyExc_ValueError, invalid_capsule);
21-
return 0;
21+
if (!op || !PyCapsule_CheckExact(op)) {
22+
goto error;
23+
}
24+
PyCapsule *capsule = (PyCapsule *)op;
25+
26+
if (capsule->pointer == NULL) {
27+
goto error;
2228
}
2329
return 1;
30+
31+
error:
32+
PyErr_SetString(PyExc_ValueError, invalid_capsule);
33+
return 0;
2434
}
2535

2636
#define is_legal_capsule(capsule, name) \
@@ -50,7 +60,7 @@ PyCapsule_New(void *pointer, const char *name, PyCapsule_Destructor destructor)
5060
return NULL;
5161
}
5262

53-
capsule = PyObject_New(PyCapsule, &PyCapsule_Type);
63+
capsule = PyObject_GC_New(PyCapsule, &PyCapsule_Type);
5464
if (capsule == NULL) {
5565
return NULL;
5666
}
@@ -59,15 +69,18 @@ PyCapsule_New(void *pointer, const char *name, PyCapsule_Destructor destructor)
5969
capsule->name = name;
6070
capsule->context = NULL;
6171
capsule->destructor = destructor;
72+
capsule->traverse_func = NULL;
73+
capsule->clear_func = NULL;
74+
// Only track the capsule if _PyCapsule_SetTraverse() is called
6275

6376
return (PyObject *)capsule;
6477
}
6578

6679

6780
int
68-
PyCapsule_IsValid(PyObject *o, const char *name)
81+
PyCapsule_IsValid(PyObject *op, const char *name)
6982
{
70-
PyCapsule *capsule = (PyCapsule *)o;
83+
PyCapsule *capsule = (PyCapsule *)op;
7184

7285
return (capsule != NULL &&
7386
PyCapsule_CheckExact(capsule) &&
@@ -77,13 +90,12 @@ PyCapsule_IsValid(PyObject *o, const char *name)
7790

7891

7992
void *
80-
PyCapsule_GetPointer(PyObject *o, const char *name)
93+
PyCapsule_GetPointer(PyObject *op, const char *name)
8194
{
82-
PyCapsule *capsule = (PyCapsule *)o;
83-
84-
if (!is_legal_capsule(capsule, "PyCapsule_GetPointer")) {
95+
if (!is_legal_capsule(op, "PyCapsule_GetPointer")) {
8596
return NULL;
8697
}
98+
PyCapsule *capsule = (PyCapsule *)op;
8799

88100
if (!name_matches(name, capsule->name)) {
89101
PyErr_SetString(PyExc_ValueError, "PyCapsule_GetPointer called with incorrect name");
@@ -95,52 +107,48 @@ PyCapsule_GetPointer(PyObject *o, const char *name)
95107

96108

97109
const char *
98-
PyCapsule_GetName(PyObject *o)
110+
PyCapsule_GetName(PyObject *op)
99111
{
100-
PyCapsule *capsule = (PyCapsule *)o;
101-
102-
if (!is_legal_capsule(capsule, "PyCapsule_GetName")) {
112+
if (!is_legal_capsule(op, "PyCapsule_GetName")) {
103113
return NULL;
104114
}
115+
PyCapsule *capsule = (PyCapsule *)op;
105116
return capsule->name;
106117
}
107118

108119

109120
PyCapsule_Destructor
110-
PyCapsule_GetDestructor(PyObject *o)
121+
PyCapsule_GetDestructor(PyObject *op)
111122
{
112-
PyCapsule *capsule = (PyCapsule *)o;
113-
114-
if (!is_legal_capsule(capsule, "PyCapsule_GetDestructor")) {
123+
if (!is_legal_capsule(op, "PyCapsule_GetDestructor")) {
115124
return NULL;
116125
}
126+
PyCapsule *capsule = (PyCapsule *)op;
117127
return capsule->destructor;
118128
}
119129

120130

121131
void *
122-
PyCapsule_GetContext(PyObject *o)
132+
PyCapsule_GetContext(PyObject *op)
123133
{
124-
PyCapsule *capsule = (PyCapsule *)o;
125-
126-
if (!is_legal_capsule(capsule, "PyCapsule_GetContext")) {
134+
if (!is_legal_capsule(op, "PyCapsule_GetContext")) {
127135
return NULL;
128136
}
137+
PyCapsule *capsule = (PyCapsule *)op;
129138
return capsule->context;
130139
}
131140

132141

133142
int
134-
PyCapsule_SetPointer(PyObject *o, void *pointer)
143+
PyCapsule_SetPointer(PyObject *op, void *pointer)
135144
{
136-
PyCapsule *capsule = (PyCapsule *)o;
137-
138-
if (!pointer) {
139-
PyErr_SetString(PyExc_ValueError, "PyCapsule_SetPointer called with null pointer");
145+
if (!is_legal_capsule(op, "PyCapsule_SetPointer")) {
140146
return -1;
141147
}
148+
PyCapsule *capsule = (PyCapsule *)op;
142149

143-
if (!is_legal_capsule(capsule, "PyCapsule_SetPointer")) {
150+
if (!pointer) {
151+
PyErr_SetString(PyExc_ValueError, "PyCapsule_SetPointer called with null pointer");
144152
return -1;
145153
}
146154

@@ -150,47 +158,62 @@ PyCapsule_SetPointer(PyObject *o, void *pointer)
150158

151159

152160
int
153-
PyCapsule_SetName(PyObject *o, const char *name)
161+
PyCapsule_SetName(PyObject *op, const char *name)
154162
{
155-
PyCapsule *capsule = (PyCapsule *)o;
156-
157-
if (!is_legal_capsule(capsule, "PyCapsule_SetName")) {
163+
if (!is_legal_capsule(op, "PyCapsule_SetName")) {
158164
return -1;
159165
}
166+
PyCapsule *capsule = (PyCapsule *)op;
160167

161168
capsule->name = name;
162169
return 0;
163170
}
164171

165172

166173
int
167-
PyCapsule_SetDestructor(PyObject *o, PyCapsule_Destructor destructor)
174+
PyCapsule_SetDestructor(PyObject *op, PyCapsule_Destructor destructor)
168175
{
169-
PyCapsule *capsule = (PyCapsule *)o;
170-
171-
if (!is_legal_capsule(capsule, "PyCapsule_SetDestructor")) {
176+
if (!is_legal_capsule(op, "PyCapsule_SetDestructor")) {
172177
return -1;
173178
}
179+
PyCapsule *capsule = (PyCapsule *)op;
174180

175181
capsule->destructor = destructor;
176182
return 0;
177183
}
178184

179185

180186
int
181-
PyCapsule_SetContext(PyObject *o, void *context)
187+
PyCapsule_SetContext(PyObject *op, void *context)
182188
{
183-
PyCapsule *capsule = (PyCapsule *)o;
184-
185-
if (!is_legal_capsule(capsule, "PyCapsule_SetContext")) {
189+
if (!is_legal_capsule(op, "PyCapsule_SetContext")) {
186190
return -1;
187191
}
192+
PyCapsule *capsule = (PyCapsule *)op;
188193

189194
capsule->context = context;
190195
return 0;
191196
}
192197

193198

199+
int
200+
_PyCapsule_SetTraverse(PyObject *op, traverseproc traverse_func, inquiry clear_func)
201+
{
202+
if (!is_legal_capsule(op, "_PyCapsule_SetTraverse")) {
203+
return -1;
204+
}
205+
PyCapsule *capsule = (PyCapsule *)op;
206+
207+
if (!PyObject_GC_IsTracked(op)) {
208+
PyObject_GC_Track(op);
209+
}
210+
211+
capsule->traverse_func = traverse_func;
212+
capsule->clear_func = clear_func;
213+
return 0;
214+
}
215+
216+
194217
void *
195218
PyCapsule_Import(const char *name, int no_block)
196219
{
@@ -249,13 +272,14 @@ PyCapsule_Import(const char *name, int no_block)
249272

250273

251274
static void
252-
capsule_dealloc(PyObject *o)
275+
capsule_dealloc(PyObject *op)
253276
{
254-
PyCapsule *capsule = (PyCapsule *)o;
277+
PyCapsule *capsule = (PyCapsule *)op;
278+
PyObject_GC_UnTrack(op);
255279
if (capsule->destructor) {
256-
capsule->destructor(o);
280+
capsule->destructor(op);
257281
}
258-
PyObject_Free(o);
282+
PyObject_GC_Del(op);
259283
}
260284

261285

@@ -279,6 +303,29 @@ capsule_repr(PyObject *o)
279303
}
280304

281305

306+
static int
307+
capsule_traverse(PyCapsule *capsule, visitproc visit, void *arg)
308+
{
309+
if (capsule->traverse_func) {
310+
return capsule->traverse_func((PyObject*)capsule, visit, arg);
311+
}
312+
else {
313+
return 0;
314+
}
315+
}
316+
317+
318+
static int
319+
capsule_clear(PyCapsule *capsule)
320+
{
321+
if (capsule->clear_func) {
322+
return capsule->clear_func((PyObject*)capsule);
323+
}
324+
else {
325+
return 0;
326+
}
327+
}
328+
282329

283330
PyDoc_STRVAR(PyCapsule_Type__doc__,
284331
"Capsule objects let you wrap a C \"void *\" pointer in a Python\n\
@@ -293,27 +340,14 @@ Python import mechanism to link to one another.\n\
293340

294341
PyTypeObject PyCapsule_Type = {
295342
PyVarObject_HEAD_INIT(&PyType_Type, 0)
296-
"PyCapsule", /*tp_name*/
297-
sizeof(PyCapsule), /*tp_basicsize*/
298-
0, /*tp_itemsize*/
299-
/* methods */
300-
capsule_dealloc, /*tp_dealloc*/
301-
0, /*tp_vectorcall_offset*/
302-
0, /*tp_getattr*/
303-
0, /*tp_setattr*/
304-
0, /*tp_as_async*/
305-
capsule_repr, /*tp_repr*/
306-
0, /*tp_as_number*/
307-
0, /*tp_as_sequence*/
308-
0, /*tp_as_mapping*/
309-
0, /*tp_hash*/
310-
0, /*tp_call*/
311-
0, /*tp_str*/
312-
0, /*tp_getattro*/
313-
0, /*tp_setattro*/
314-
0, /*tp_as_buffer*/
315-
0, /*tp_flags*/
316-
PyCapsule_Type__doc__ /*tp_doc*/
343+
.tp_name = "PyCapsule",
344+
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
345+
.tp_basicsize = sizeof(PyCapsule),
346+
.tp_dealloc = capsule_dealloc,
347+
.tp_repr = capsule_repr,
348+
.tp_doc = PyCapsule_Type__doc__,
349+
.tp_traverse = (traverseproc)capsule_traverse,
350+
.tp_clear = (inquiry)capsule_clear,
317351
};
318352

319353

0 commit comments

Comments
 (0)