44#include <pyawaitable/genwrapper.h>
55#include <pyawaitable/coro.h>
66#include <stdlib.h>
7+ #define AWAITABLE_POOL_SIZE 256
78
89PyDoc_STRVAR (
910 awaitable_doc ,
1011 "Awaitable transport utility for the C API."
1112);
1213
14+ static Py_ssize_t pool_index = 0 ;
15+ static PyObject * pool [AWAITABLE_POOL_SIZE ];
16+
1317static PyObject *
1418awaitable_new_func (PyTypeObject * tp , PyObject * args , PyObject * kwds )
1519{
@@ -22,18 +26,11 @@ awaitable_new_func(PyTypeObject *tp, PyObject *args, PyObject *kwds)
2226 return NULL ;
2327 }
2428
25- PyAwaitableObject * aw = (PyAwaitableObject * )self ;
26- aw -> aw_callbacks = NULL ;
27- aw -> aw_callback_size = 0 ;
28- aw -> aw_result = Py_NewRef (Py_None );
29- aw -> aw_gen = NULL ;
30- aw -> aw_values = NULL ;
31- aw -> aw_values_size = 0 ;
32- aw -> aw_state = 0 ;
33- aw -> aw_done = false;
29+ PyAwaitableObject * aw = (PyAwaitableObject * ) self ;
3430 aw -> aw_awaited = false;
31+ aw -> aw_done = false;
3532
36- return (PyObject * )aw ;
33+ return (PyObject * ) aw ;
3734}
3835
3936PyObject *
@@ -44,7 +41,10 @@ awaitable_next(PyObject *self)
4441
4542 if (aw -> aw_done )
4643 {
47- PyErr_SetString (PyExc_RuntimeError , "cannot reuse awaitable" );
44+ PyErr_SetString (
45+ PyExc_RuntimeError ,
46+ "pyawaitable: cannot reuse awaitable"
47+ );
4848 return NULL ;
4949 }
5050
@@ -55,36 +55,35 @@ awaitable_next(PyObject *self)
5555 return NULL ;
5656 }
5757
58- aw -> aw_gen = Py_NewRef (gen );
59- aw -> aw_done = true;
60- return gen ;
58+ aw -> aw_gen = gen ;
59+ return Py_NewRef (gen );
6160}
6261
6362static void
6463awaitable_dealloc (PyObject * self )
6564{
6665 PyAwaitableObject * aw = (PyAwaitableObject * )self ;
67- if ( aw -> aw_values )
66+ for ( int i = 0 ; i < VALUE_ARRAY_SIZE ; ++ i )
6867 {
69- for ( int i = 0 ; i < aw -> aw_values_size ; i ++ )
70- Py_DECREF ( aw -> aw_values [ i ]) ;
71- PyMem_Free (aw -> aw_values );
68+ if (! aw -> aw_values [ i ] )
69+ break ;
70+ Py_DECREF (aw -> aw_values [ i ] );
7271 }
7372
7473 Py_XDECREF (aw -> aw_gen );
7574 Py_XDECREF (aw -> aw_result );
7675
77- for (int i = 0 ; i < aw -> aw_callback_size ; i ++ )
76+ for (int i = 0 ; i < CALLBACK_ARRAY_SIZE ; ++ i )
7877 {
7978 pyawaitable_callback * cb = aw -> aw_callbacks [i ];
79+ if (cb == NULL )
80+ break ;
81+
8082 if (!cb -> done )
8183 Py_DECREF (cb -> coro );
8284 PyMem_Free (cb );
8385 }
8486
85- if (aw -> aw_arb_values )
86- PyMem_Free (aw -> aw_arb_values );
87-
8887 if (!aw -> aw_done )
8988 {
9089 if (
@@ -108,17 +107,20 @@ pyawaitable_cancel_impl(PyObject *aw)
108107 assert (aw != NULL );
109108 Py_INCREF (aw );
110109
111- PyAwaitableObject * a = (PyAwaitableObject * )aw ;
110+ PyAwaitableObject * a = (PyAwaitableObject * ) aw ;
112111
113- for (int i = 0 ; i < a -> aw_callback_size ; i ++ )
112+ for (int i = 0 ; i < CALLBACK_ARRAY_SIZE ; ++ i )
114113 {
115114 pyawaitable_callback * cb = a -> aw_callbacks [i ];
115+ if (!cb )
116+ break ;
117+
116118 if (!cb -> done )
117119 Py_DECREF (cb -> coro );
120+
121+ a -> aw_callbacks [i ] = NULL ;
118122 }
119123
120- PyMem_Free (a -> aw_callbacks );
121- a -> aw_callback_size = 0 ;
122124 Py_DECREF (aw );
123125}
124126
@@ -134,47 +136,30 @@ pyawaitable_await_impl(
134136 assert (coro != NULL );
135137 Py_INCREF (coro );
136138 Py_INCREF (aw );
137- PyAwaitableObject * a = (PyAwaitableObject * )aw ;
138-
139- pyawaitable_callback * aw_c = PyMem_Malloc (sizeof (pyawaitable_callback ));
140- if (aw_c == NULL )
139+ PyAwaitableObject * a = (PyAwaitableObject * ) aw ;
140+ if (a -> aw_callback_index == CALLBACK_ARRAY_SIZE )
141141 {
142- Py_DECREF (aw );
143- Py_DECREF (coro );
144- PyErr_NoMemory ();
145- return -1 ;
146- }
147-
148- ++ a -> aw_callback_size ;
149- if (a -> aw_callbacks == NULL )
150- {
151- a -> aw_callbacks = PyMem_Calloc (
152- a -> aw_callback_size ,
153- sizeof (pyawaitable_callback * )
154- );
155- } else
156- {
157- a -> aw_callbacks = PyMem_Realloc (
158- a -> aw_callbacks ,
159- sizeof (pyawaitable_callback * ) *
160- a -> aw_callback_size
142+ PyErr_SetString (
143+ PyExc_SystemError ,
144+ "pyawaitable: awaitable object cannot store more than 128 coroutines"
161145 );
146+ return -1 ;
162147 }
163148
164- if (a -> aw_callbacks == NULL )
149+ pyawaitable_callback * aw_c = PyMem_Malloc (sizeof (pyawaitable_callback ));
150+ if (aw_c == NULL )
165151 {
166- -- a -> aw_callback_size ;
167152 Py_DECREF (aw );
168153 Py_DECREF (coro );
169- PyMem_Free (aw_c );
170154 PyErr_NoMemory ();
171155 return -1 ;
172156 }
173157
174- aw_c -> coro = coro ; // steal our own reference
158+ aw_c -> coro = coro ; // Steal our own reference
175159 aw_c -> callback = cb ;
176160 aw_c -> err_callback = err ;
177- a -> aw_callbacks [a -> aw_callback_size - 1 ] = aw_c ;
161+ aw_c -> done = false;
162+ a -> aw_callbacks [a -> aw_callback_index ++ ] = aw_c ;
178163 Py_DECREF (aw );
179164
180165 return 0 ;
@@ -185,21 +170,57 @@ pyawaitable_set_result_impl(PyObject *awaitable, PyObject *result)
185170{
186171 assert (awaitable != NULL );
187172 assert (result != NULL );
188- Py_INCREF (result );
189- Py_INCREF (awaitable );
190173
191- PyAwaitableObject * aw = (PyAwaitableObject * )awaitable ;
174+ PyAwaitableObject * aw = (PyAwaitableObject * ) awaitable ;
192175 aw -> aw_result = Py_NewRef (result );
193- Py_DECREF (awaitable );
194- Py_DECREF (result );
195176 return 0 ;
196177}
197178
198179PyObject *
199180pyawaitable_new_impl (void )
200181{
201- PyObject * aw = awaitable_new_func (& _PyAwaitableType , NULL , NULL );
202- return aw ;
182+ if (pool_index == AWAITABLE_POOL_SIZE )
183+ {
184+ PyObject * aw = awaitable_new_func (& _PyAwaitableType , NULL , NULL );
185+ return aw ;
186+ }
187+
188+ return pool [pool_index ++ ];
189+ }
190+
191+ void
192+ dealloc_awaitable_pool (void )
193+ {
194+ for (Py_ssize_t i = pool_index ; i < AWAITABLE_POOL_SIZE ; ++ i )
195+ {
196+ if (Py_REFCNT (pool [i ]) != 1 )
197+ {
198+ PyErr_Format (
199+ PyExc_SystemError ,
200+ "expected %R to have a reference count of 1" ,
201+ pool [i ]
202+ );
203+ PyErr_WriteUnraisable (NULL );
204+ }
205+ Py_DECREF (pool [i ]);
206+ }
207+ }
208+
209+ int
210+ alloc_awaitable_pool (void )
211+ {
212+ for (Py_ssize_t i = 0 ; i < AWAITABLE_POOL_SIZE ; ++ i )
213+ {
214+ pool [i ] = awaitable_new_func (& _PyAwaitableType , NULL , NULL );
215+ if (!pool [i ])
216+ {
217+ for (Py_ssize_t x = 0 ; x < i ; ++ x )
218+ Py_DECREF (pool [x ]);
219+ return -1 ;
220+ }
221+ }
222+
223+ return 0 ;
203224}
204225
205226int
0 commit comments