Skip to content

Commit b726f10

Browse files
ZipFileZipFile
committed
Properly handle signals while sleeping
1 parent 9d86ce5 commit b726f10

File tree

3 files changed

+92
-77
lines changed

3 files changed

+92
-77
lines changed

‎src/sonyflake_turbo/sonyflake.c‎

Lines changed: 57 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,10 @@ bool incr_combined_sequence(struct sonyflake_state *self) {
3737
return self->combined_sequence == 0;
3838
}
3939

40-
static inline void get_relative_current_time(struct sonyflake_state *self, struct timespec *now) {
40+
static inline void get_relative_current_time(struct sonyflake_state *self, struct timespec *now, struct timespec *sf_now) {
4141
timespec_get(now, TIME_UTC);
42-
sub_diff(now, &self->start_time);
42+
*sf_now = *now;
43+
sub_diff(sf_now, &self->start_time);
4344
}
4445

4546
static PyObject *sonyflake_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
@@ -180,80 +181,80 @@ static void sonyflake_dealloc(struct sonyflake_state *self) {
180181
Py_DECREF(tp);
181182
}
182183

183-
static PyObject *sonyflake_sleep(PyObject *obj, struct timespec *to_nanosleep) {
184-
if (!obj) {
184+
static PyObject *sonyflake_sleep(PyObject *obj, struct sonyflake_next_sleep_info *sleep_info) {
185+
if (!(obj && sleep_info)) {
185186
return NULL;
186187
}
187188

188-
if (!to_nanosleep) {
189-
return obj;
190-
}
189+
struct timespec duration, now = sleep_info->now;
190+
int ret = 0;
191191

192-
if (to_nanosleep->tv_sec == 0 && to_nanosleep->tv_nsec == 0) {
193-
return obj;
194-
}
192+
for (;;) {
193+
duration = sleep_info->future;
194+
sub_diff(&duration, &now);
195195

196-
int ret = 0;
196+
if (duration.tv_sec == 0 && duration.tv_nsec == 0) {
197+
return obj;
198+
}
197199

198-
Py_BEGIN_ALLOW_THREADS;
200+
Py_BEGIN_ALLOW_THREADS;
199201
#if __STDC_NO_THREADS__
200-
ret = nanosleep(to_nanosleep, NULL);
202+
ret = nanosleep(&duration, NULL);
201203
#else
202-
ret = thrd_sleep(to_nanosleep, NULL);
204+
ret = thrd_sleep(&duration, NULL);
203205
#endif
204-
Py_END_ALLOW_THREADS;
206+
Py_END_ALLOW_THREADS;
205207

206-
if (ret == 0) {
207-
return obj;
208-
}
208+
if (ret == 0) {
209+
return obj;
210+
}
209211

210-
#if __STDC_NO_THREADS__
211-
if (ret < 0 && errno == EINTR && PyErr_CheckSignals()) {
212-
goto err;
213-
}
212+
#if __STDC_NO_THREADS__ // nanosleep
213+
if (ret < 0 && errno == EINTR) {
214214
#else
215-
if (ret == -1) {
216-
if (PyErr_CheckSignals()) {
217-
goto err;
215+
if (ret == -1) {
216+
#endif
217+
if (PyErr_CheckSignals()) {
218+
goto err;
219+
} else {
220+
timespec_get(&now, TIME_UTC);
221+
continue;
222+
}
218223
}
219224

220-
return obj;
225+
break;
221226
}
222-
#endif
223227

224228
PyErr_SetFromErrno(PyExc_OSError);
225229
err:
226230
Py_DECREF(obj);
227231
return NULL;
228232
}
229233

230-
PyObject *sonyflake_next(struct sonyflake_state *self, struct timespec *to_nanosleep) {
231-
struct timespec now, future;
234+
PyObject *sonyflake_next(struct sonyflake_state *self, struct sonyflake_next_sleep_info *sleep_info) {
235+
struct timespec now, sf_now;
232236
sonyflake_time current;
233237
uint64_t sonyflake_id;
234238

235239
PyThread_acquire_lock(self->lock, 1);
236240

237-
get_relative_current_time(self, &now);
241+
get_relative_current_time(self, &now, &sf_now);
238242

239-
current = to_sonyflake_time(&now);
240-
241-
if (to_nanosleep) {
242-
to_nanosleep->tv_sec = 0;
243-
to_nanosleep->tv_nsec = 0;
243+
if (sleep_info) {
244+
sleep_info->now = now;
245+
sleep_info->future = now;
244246
}
245247

248+
current = to_sonyflake_time(&sf_now);
249+
246250
if (self->elapsed_time < current) {
247251
self->elapsed_time = current;
248252
self->combined_sequence = 0;
249253
} else if (incr_combined_sequence(self)) {
250254
self->elapsed_time++;
251255

252-
if (to_nanosleep) {
253-
from_sonyflake_time(self->elapsed_time, &future);
254-
sub_diff(&future, &now);
255-
256-
*to_nanosleep = future;
256+
if (sleep_info) {
257+
from_sonyflake_time(&self->start_time, self->elapsed_time, &sleep_info->future);
257258
}
258259
}
259260

@@ -264,7 +265,7 @@ PyObject *sonyflake_next(struct sonyflake_state *self, struct timespec *to_nanos
264265
return PyLong_FromUnsignedLongLong(sonyflake_id);
265266
}
266267

267-
PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct timespec *to_nanosleep) {
268+
PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct sonyflake_next_sleep_info *sleep_info) {
268269
assert(n > 0);
269270

270271
PyObject *out = PyList_New(n);
@@ -273,14 +274,14 @@ PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct ti
273274
return NULL;
274275
}
275276

276-
struct timespec now, future;
277-
sonyflake_time current, diff;
277+
struct timespec now, sf_now;
278+
sonyflake_time current;
278279

279280
PyThread_acquire_lock(self->lock, 1);
280281

281-
get_relative_current_time(self, &now);
282+
get_relative_current_time(self, &now, &sf_now);
282283

283-
current = to_sonyflake_time(&now);
284+
current = to_sonyflake_time(&sf_now);
284285

285286
if (self->elapsed_time < current) {
286287
self->elapsed_time = current;
@@ -299,28 +300,13 @@ PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct ti
299300
PyList_SetItem(out, i, PyLong_FromUnsignedLongLong(compose(self)));
300301
}
301302

302-
if (!to_nanosleep) {
303-
goto end;
304-
}
305-
306-
to_nanosleep->tv_sec = 0;
307-
to_nanosleep->tv_nsec = 0;
308-
diff = self->elapsed_time - current;
303+
PyThread_release_lock(self->lock);
309304

310-
if (diff <= 0) {
311-
goto end;
312-
} else if (diff > 1) {
313-
get_relative_current_time(self, &now);
305+
if (sleep_info) {
306+
timespec_get(&sleep_info->now, TIME_UTC);
307+
from_sonyflake_time(&self->start_time, self->elapsed_time, &sleep_info->future);
314308
}
315309

316-
from_sonyflake_time(self->elapsed_time, &future);
317-
sub_diff(&future, &now);
318-
319-
*to_nanosleep = future;
320-
321-
end:
322-
PyThread_release_lock(self->lock);
323-
324310
return out;
325311
}
326312

@@ -375,10 +361,11 @@ static PyObject *sonyflake_repr(struct sonyflake_state *self) {
375361
}
376362

377363
static PyObject *sonyflake_iternext(struct sonyflake_state *self) {
378-
struct timespec to_nanosleep = { 0, 0 };
379-
PyObject *sonyflake_id = sonyflake_next(self, &to_nanosleep);
364+
struct sonyflake_next_sleep_info sleep_info;
365+
366+
PyObject *sonyflake_id = sonyflake_next(self, &sleep_info);
380367

381-
return sonyflake_sleep(sonyflake_id, &to_nanosleep);
368+
return sonyflake_sleep(sonyflake_id, &sleep_info);
382369
}
383370

384371
static PyObject *sonyflake_call(struct sonyflake_state *self, PyObject *args) {
@@ -393,10 +380,10 @@ static PyObject *sonyflake_call(struct sonyflake_state *self, PyObject *args) {
393380
return NULL;
394381
}
395382

396-
struct timespec to_nanosleep = { 0, 0 };
397-
PyObject *sonyflake_ids = sonyflake_next_n(self, n, &to_nanosleep);
383+
struct sonyflake_next_sleep_info sleep_info;
384+
PyObject *sonyflake_ids = sonyflake_next_n(self, n, &sleep_info);
398385

399-
return sonyflake_sleep(sonyflake_ids, &to_nanosleep);
386+
return sonyflake_sleep(sonyflake_ids, &sleep_info);
400387
}
401388

402389
PyDoc_STRVAR(sonyflake_doc,

‎src/sonyflake_turbo/sonyflake.h‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,41 @@ const static struct timespec default_start_time = {
2525
.tv_nsec = 0
2626
};
2727

28+
struct sonyflake_next_sleep_info {
29+
struct timespec now;
30+
struct timespec future;
31+
};
32+
33+
#define print_time(name, ts) printf("%s: %ld.%09ld\n", name, (long) ts.tv_sec, (long) ts.tv_nsec)
34+
2835
static inline sonyflake_time to_sonyflake_time(const struct timespec *ts) {
2936
return ts->tv_sec * 100 + ts->tv_nsec / 10000000;
3037
}
3138

32-
static inline void from_sonyflake_time(sonyflake_time sf_time, struct timespec *ts) {
33-
ts->tv_sec = sf_time / 100;
34-
ts->tv_nsec = (sf_time % 100) * 10000000;
39+
static inline void from_sonyflake_time(struct timespec *start_time, sonyflake_time sf_time, struct timespec *ts) {
40+
ts->tv_sec = start_time->tv_sec + sf_time / 100;
41+
ts->tv_nsec = start_time->tv_nsec + (sf_time % 100) * 10000000;
42+
43+
if (ts->tv_nsec >= 1000000000) {
44+
ts->tv_sec++;
45+
ts->tv_nsec -= 1000000000;
46+
}
3547
}
3648

3749
static inline void sub_diff(struct timespec *a, const struct timespec *b) {
50+
if (a->tv_sec < b->tv_sec) {
51+
a->tv_sec = 0;
52+
a->tv_nsec = 0;
53+
return;
54+
}
55+
3856
a->tv_sec -= b->tv_sec;
57+
58+
if (a->tv_sec == 0 && a->tv_nsec < b->tv_nsec) {
59+
a->tv_nsec = 0;
60+
return;
61+
}
62+
3963
a->tv_nsec -= b->tv_nsec;
4064

4165
if (a->tv_nsec < 0) {
@@ -44,5 +68,5 @@ static inline void sub_diff(struct timespec *a, const struct timespec *b) {
4468
}
4569
}
4670

47-
PyObject *sonyflake_next(struct sonyflake_state *self, struct timespec *to_nanosleep);
48-
PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct timespec *to_nanosleep);
71+
PyObject *sonyflake_next(struct sonyflake_state *self, struct sonyflake_next_sleep_info *sleep_info);
72+
PyObject *sonyflake_next_n(struct sonyflake_state *self, Py_ssize_t n, struct sonyflake_next_sleep_info *sleep_info);

‎tests/test_sonyflake_turbo.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import threading
44
from concurrent.futures import ThreadPoolExecutor
55
from datetime import datetime
6-
from time import sleep
6+
from time import sleep, perf_counter
77

88
from pytest import mark, raises
99

@@ -109,8 +109,12 @@ def thread_func() -> None:
109109
signal.pthread_kill(thread_id, signal.SIGUSR1)
110110

111111
threading.Thread(target=thread_func).start()
112+
a = perf_counter()
113+
ids = sf(50000)
114+
b = perf_counter()
112115

113-
assert len(sf(50000)) == 50000
116+
assert len(ids) == 50000
117+
assert (b - a) > 1 # should take roughly 1.9sec
114118

115119

116120
def test_sonyflake_repr() -> None:

0 commit comments

Comments
 (0)