Skip to content

Commit c3e9b0b

Browse files
bmosconbmoscon
committed
Start work on ref counting bugs, missing error checks and some performance improvements
1 parent 0c32932 commit c3e9b0b

5 files changed

Lines changed: 164 additions & 55 deletions

File tree

‎orderbook/orderbook.c‎

Lines changed: 50 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -61,39 +61,20 @@ int Orderbook_init(Orderbook *self, PyObject *args, PyObject *kwds)
6161
}
6262

6363

64+
// Lazy allocation: Store checksum format but don't allocate buffer until needed
6465
if (checksum_str.buf && checksum_str.len) {
6566
if (strncmp(checksum_str.buf, "KRAKEN", checksum_str.len) == 0) {
6667
self->checksum = KRAKEN;
67-
self->checksum_buffer = calloc(2048, sizeof(uint8_t));
6868
self->checksum_len = 2048;
69-
if (!self->checksum_buffer) {
70-
PyErr_SetNone(PyExc_MemoryError);
71-
return -1;
72-
}
7369
} else if ((checksum_str.len > 2) && (strncmp(checksum_str.buf, "FTX", 3) == 0)) {
7470
self->checksum = FTX;
75-
self->checksum_buffer = calloc(20480, sizeof(uint8_t));
7671
self->checksum_len = 20480;
77-
if (!self->checksum_buffer) {
78-
PyErr_SetNone(PyExc_MemoryError);
79-
return -1;
80-
}
8172
} else if ((checksum_str.len > 2) && ((strncmp(checksum_str.buf, "OKX", 3) == 0) || (strncmp(checksum_str.buf, "OKCO", 4) == 0))) {
8273
self->checksum = OKX;
83-
self->checksum_buffer = calloc(4096, sizeof(uint8_t));
8474
self->checksum_len = 4096;
85-
if (!self->checksum_buffer) {
86-
PyErr_SetNone(PyExc_MemoryError);
87-
return -1;
88-
}
8975
} else if (strncmp(checksum_str.buf, "BITGET", checksum_str.len) == 0) {
9076
self->checksum = BITGET;
91-
self->checksum_buffer = calloc(4096, sizeof(uint8_t));
9277
self->checksum_len = 4096;
93-
if (!self->checksum_buffer) {
94-
PyErr_SetNone(PyExc_MemoryError);
95-
return -1;
96-
}
9778
} else {
9879
PyBuffer_Release(&checksum_str);
9980
PyErr_SetString(PyExc_TypeError, "invalid checksum format specified");
@@ -162,6 +143,17 @@ PyObject* Orderbook_checksum(const Orderbook *self, PyObject *Py_UNUSED(ignored)
162143
return NULL;
163144
}
164145

146+
// Lazy allocation: allocate buffer on first use
147+
// Cast away const since we're doing lazy initialization
148+
Orderbook *mutable_self = (Orderbook *)self;
149+
if (!mutable_self->checksum_buffer) {
150+
mutable_self->checksum_buffer = calloc(mutable_self->checksum_len, sizeof(uint8_t));
151+
if (!mutable_self->checksum_buffer) {
152+
PyErr_SetNone(PyExc_MemoryError);
153+
return NULL;
154+
}
155+
}
156+
165157
if (EXPECT(update_keys(self->bids), 0)) {
166158
return NULL;
167159
}
@@ -190,13 +182,9 @@ PyObject *Orderbook_getitem(const Orderbook *self, PyObject *key)
190182
return NULL;
191183
}
192184

193-
PyObject *str = PyUnicode_AsEncodedString(key, "UTF-8", "strict");
194-
if (EXPECT(!str, 0)) {
195-
return NULL;
196-
}
197-
198-
enum side_e key_int = check_key(PyBytes_AsString(str));
199-
Py_DECREF(str);
185+
// Use optimized key checking with cached strings
186+
OrderBookModuleState* st = get_order_book_state(NULL);
187+
enum side_e key_int = check_key_pyobject(key, st->str_bid, st->str_ask, st->str_bids, st->str_asks);
200188

201189
if (key_int == BID) {
202190
Py_INCREF(self->bids);
@@ -219,13 +207,9 @@ int Orderbook_setitem(const Orderbook *self, PyObject *key, PyObject *value)
219207
return -1;
220208
}
221209

222-
PyObject *str = PyUnicode_AsEncodedString(key, "UTF-8", "strict");
223-
if (EXPECT(!str, 0)) {
224-
return -1;
225-
}
226-
227-
enum side_e key_int = check_key(PyBytes_AsString(str));
228-
Py_DECREF(str);
210+
// Use optimized key checking with cached strings
211+
OrderBookModuleState* st = get_order_book_state(NULL);
212+
enum side_e key_int = check_key_pyobject(key, st->str_bid, st->str_ask, st->str_bids, st->str_asks);
229213

230214
if (EXPECT(key_int == INVALID_SIDE, 0)) {
231215
PyErr_SetString(PyExc_ValueError, "key must one of bid/ask");
@@ -288,6 +272,7 @@ PyMODINIT_FUNC PyInit_order_book(void)
288272

289273
Py_INCREF(&SortedDictType);
290274
if (PyModule_AddObject(m, "SortedDict", (PyObject *) &SortedDictType) < 0) {
275+
Py_DECREF(&OrderbookType);
291276
Py_DECREF(&SortedDictType);
292277
Py_DECREF(m);
293278
return NULL;
@@ -297,6 +282,7 @@ PyMODINIT_FUNC PyInit_order_book(void)
297282

298283
PyObject* builtins = PyImport_AddModule("builtins");
299284
if (builtins == NULL) {
285+
Py_DECREF(&OrderbookType);
300286
Py_DECREF(&SortedDictType);
301287
Py_DECREF(m);
302288
return NULL;
@@ -305,6 +291,7 @@ PyMODINIT_FUNC PyInit_order_book(void)
305291
st->format = PyObject_GetAttrString(builtins, "format");
306292
Py_DECREF(builtins);
307293
if (st->format == NULL) {
294+
Py_DECREF(&OrderbookType);
308295
Py_DECREF(&SortedDictType);
309296
Py_DECREF(m);
310297
return NULL;
@@ -313,6 +300,26 @@ PyMODINIT_FUNC PyInit_order_book(void)
313300
st->formatf = PyUnicode_FromString("f");
314301
if (st->formatf == NULL) {
315302
Py_DECREF(st->format);
303+
Py_DECREF(&OrderbookType);
304+
Py_DECREF(&SortedDictType);
305+
Py_DECREF(m);
306+
return NULL;
307+
}
308+
309+
// Initialize cached interned strings for fast key comparison
310+
st->str_bid = PyUnicode_InternFromString("bid");
311+
st->str_ask = PyUnicode_InternFromString("ask");
312+
st->str_bids = PyUnicode_InternFromString("bids");
313+
st->str_asks = PyUnicode_InternFromString("asks");
314+
315+
if (!st->str_bid || !st->str_ask || !st->str_bids || !st->str_asks) {
316+
Py_XDECREF(st->str_bid);
317+
Py_XDECREF(st->str_ask);
318+
Py_XDECREF(st->str_bids);
319+
Py_XDECREF(st->str_asks);
320+
Py_DECREF(st->formatf);
321+
Py_DECREF(st->format);
322+
Py_DECREF(&OrderbookType);
316323
Py_DECREF(&SortedDictType);
317324
Py_DECREF(m);
318325
return NULL;
@@ -327,6 +334,10 @@ static int order_book_traverse(PyObject *m, visitproc visit, void *arg)
327334
OrderBookModuleState* st = get_order_book_state(m);
328335
Py_VISIT(st->format);
329336
Py_VISIT(st->formatf);
337+
Py_VISIT(st->str_bid);
338+
Py_VISIT(st->str_ask);
339+
Py_VISIT(st->str_bids);
340+
Py_VISIT(st->str_asks);
330341
return 0;
331342
}
332343

@@ -336,6 +347,10 @@ static int order_book_clear(PyObject* m)
336347
OrderBookModuleState* st = get_order_book_state(m);
337348
Py_CLEAR(st->format);
338349
Py_CLEAR(st->formatf);
350+
Py_CLEAR(st->str_bid);
351+
Py_CLEAR(st->str_ask);
352+
Py_CLEAR(st->str_bids);
353+
Py_CLEAR(st->str_asks);
339354
return 0;
340355
}
341356

@@ -583,7 +598,7 @@ static PyObject* alternating_checksum(const Orderbook *ob, const uint32_t depth,
583598
PyObject *price = NULL;
584599
PyObject *size = NULL;
585600

586-
for(uint32_t i = 0; i < depth; ++i) {
601+
for(uint32_t i = 0; i < depth; ++i) {
587602
if (i < bids_size) {
588603
price = PyTuple_GET_ITEM(ob->bids->keys, i);
589604
size = PyDict_GetItem(ob->bids->data, price);

‎orderbook/orderbook.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,11 +98,16 @@ static PyTypeObject OrderbookType = {
9898
.tp_dictoffset = 0,
9999
};
100100

101-
// the module contains reusable python objects referring to the builtin format
101+
// the module contains reusable python objects referring to the builtin format
102102
// function and a fixed string 'f'
103103
typedef struct {
104104
PyObject *format;
105105
PyObject *formatf;
106+
// Cached strings for fast key comparison
107+
PyObject *str_bid;
108+
PyObject *str_ask;
109+
PyObject *str_bids;
110+
PyObject *str_asks;
106111
} OrderBookModuleState;
107112

108113
static int order_book_traverse(PyObject *m, visitproc visit, void *arg);

‎orderbook/sorteddict.c‎

Lines changed: 71 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ int SortedDict_init(SortedDict *self, PyObject *args, PyObject *kwds)
6969
}
7070

7171
PyObject *copy = PyDict_Copy(dict);
72+
if (!copy) {
73+
return -1;
74+
}
75+
7276
if (self->data) {
7377
Py_DECREF(self->data);
7478
}
@@ -77,7 +81,13 @@ int SortedDict_init(SortedDict *self, PyObject *args, PyObject *kwds)
7781

7882

7983
if (kwds && PyDict_Check(kwds) && PyDict_Size(kwds) > 0) {
80-
if (PyDict_Contains(kwds, PyUnicode_FromString("max_depth"))) {
84+
PyObject *max_depth_string = PyUnicode_FromString("max_depth");
85+
if (!max_depth_string) {
86+
return -1;
87+
}
88+
89+
if (PyDict_Contains(kwds, max_depth_string)) {
90+
Py_DECREF(max_depth_string);
8191
PyObject *max_depth = PyDict_GetItemString(kwds, "max_depth");
8292
if (PyLong_Check(max_depth)) {
8393
self->depth = PyLong_AsLong(max_depth);
@@ -94,9 +104,17 @@ int SortedDict_init(SortedDict *self, PyObject *args, PyObject *kwds)
94104
PyErr_SetString(PyExc_ValueError, "max_depth must be an integer");
95105
return -1;
96106
}
107+
} else {
108+
Py_DECREF(max_depth_string);
97109
}
98110

99-
if (PyDict_Contains(kwds, PyUnicode_FromString("truncate"))) {
111+
PyObject *truncate_string = PyUnicode_FromString("truncate");
112+
if (!truncate_string) {
113+
return -1;
114+
}
115+
116+
if (PyDict_Contains(kwds, truncate_string)) {
117+
Py_DECREF(truncate_string);
100118
PyObject *truncate = PyDict_GetItemString(kwds, "truncate");
101119

102120
if (PyBool_Check(truncate)) {
@@ -109,9 +127,17 @@ int SortedDict_init(SortedDict *self, PyObject *args, PyObject *kwds)
109127
PyErr_SetString(PyExc_ValueError, "truncate must be a boolean");
110128
return -1;
111129
}
130+
} else {
131+
Py_DECREF(truncate_string);
132+
}
133+
134+
PyObject *ordering_string = PyUnicode_FromString("ordering");
135+
if (!ordering_string) {
136+
return -1;
112137
}
113138

114-
if (PyDict_Contains(kwds, PyUnicode_FromString("ordering"))) {
139+
if (PyDict_Contains(kwds, ordering_string)) {
140+
Py_DECREF(ordering_string);
115141
ordering = PyDict_GetItemString(kwds, "ordering");
116142
if (!PyUnicode_Check(ordering)) {
117143
PyErr_SetString(PyExc_ValueError, "ordering must be a string");
@@ -138,6 +164,7 @@ int SortedDict_init(SortedDict *self, PyObject *args, PyObject *kwds)
138164
}
139165
Py_DECREF(str);
140166
} else {
167+
Py_DECREF(ordering_string);
141168
// default is ascending
142169
self->ordering = ASCENDING;
143170
}
@@ -364,20 +391,23 @@ PyObject* SortedDict_tolist(SortedDict *self, PyObject *Py_UNUSED(ignored))
364391
// new reference
365392
PyObject *key = PySequence_GetItem(self->keys, i);
366393
if (EXPECT(!key, 0)) {
394+
Py_DECREF(ret);
367395
return NULL;
368396
}
369397

370398
// borrowed reference
371399
PyObject *value = PyDict_GetItem(self->data, key);
372400
if (EXPECT(!value, 0)) {
373401
Py_DECREF(key);
402+
Py_DECREF(ret);
374403
return value;
375404
}
376405

377406
// Build tuple of (i.e., key, value)
378407
PyObject *tuple_entry = PyTuple_New(2);
379408
if (EXPECT(!tuple_entry, 0)) {
380409
Py_DECREF(key);
410+
Py_DECREF(ret);
381411
return NULL;
382412
}
383413
PyTuple_SET_ITEM(tuple_entry, 0, key);
@@ -387,7 +417,7 @@ PyObject* SortedDict_tolist(SortedDict *self, PyObject *Py_UNUSED(ignored))
387417
// Add tuple to list
388418
PyList_SET_ITEM(ret, i, tuple_entry);
389419
}
390-
420+
391421
return ret;
392422
}
393423

@@ -410,6 +440,7 @@ PyObject* SortedDict_truncate(SortedDict *self, PyObject *Py_UNUSED(ignored))
410440
return NULL;
411441
}
412442

443+
// Delete items beyond depth
413444
for (int i = 0; i < len; ++i) {
414445
if (EXPECT(PyDict_DelItem(self->data, PySequence_Fast_GET_ITEM(delete, i)) == -1, 0)) {
415446
Py_DECREF(delete);
@@ -418,12 +449,17 @@ PyObject* SortedDict_truncate(SortedDict *self, PyObject *Py_UNUSED(ignored))
418449
}
419450
Py_DECREF(delete);
420451

452+
// Optimization: Instead of marking dirty and re-sorting, just slice the keys tuple
453+
// since we already have sorted keys and only removed items from the end
421454
if (len > 0) {
422-
self->dirty = true;
423-
}
424-
425-
if (EXPECT(update_keys(self), 0)) {
426-
return NULL;
455+
PyObject *new_keys = PySequence_GetSlice(self->keys, 0, self->depth);
456+
if (EXPECT(!new_keys, 0)) {
457+
return NULL;
458+
}
459+
Py_DECREF(self->keys);
460+
self->keys = new_keys;
461+
// Keys are still sorted, so don't mark dirty
462+
self->dirty = false;
427463
}
428464
}
429465

@@ -460,16 +496,27 @@ PyObject *SortedDict_getitem(SortedDict *self, PyObject *key)
460496
int SortedDict_setitem(SortedDict *self, PyObject *key, PyObject *value)
461497
{
462498
if (value) {
463-
if (EXPECT(PyDict_Contains(self->data, key) == 0, 0)) {
499+
int key_exists = PyDict_Contains(self->data, key);
500+
if (EXPECT(key_exists == -1, 0)) {
501+
return -1; // Error occurred
502+
}
503+
504+
if (EXPECT(key_exists == 0, 0)) {
464505
self->dirty = true;
465-
}
506+
}
466507

467508
int ret = PyDict_SetItem(self->data, key, value);
468509

469510
if (EXPECT(ret == -1, 0)) {
470511
return ret;
471-
} else if (EXPECT(self->truncate && !SortedDict_truncate(self, NULL), 0)) {
472-
return -1;
512+
}
513+
514+
// Optimization: Only truncate if dict size exceeds depth
515+
// This avoids redundant truncate calls when updating existing keys
516+
if (EXPECT(self->truncate && self->depth > 0 && PyDict_Size(self->data) > self->depth, 0)) {
517+
if (!SortedDict_truncate(self, NULL)) {
518+
return -1;
519+
}
473520
}
474521

475522
return ret;
@@ -483,7 +530,7 @@ int SortedDict_setitem(SortedDict *self, PyObject *key, PyObject *value)
483530
/* Seq Functions */
484531
int SortedDict_contains(const SortedDict *self, PyObject *value)
485532
{
486-
return PySequence_Contains(self->data, value);
533+
return PyDict_Contains(self->data, value);
487534
}
488535

489536
/* iterator methods */
@@ -504,6 +551,11 @@ PyObject *SortedDict_next(SortedDict *self)
504551
}
505552

506553
Py_ssize_t size = PySequence_Fast_GET_SIZE(self->keys);
554+
// Respect depth limit for consistency with keys(), to_dict(), etc.
555+
if (self->depth > 0 && self->depth < size) {
556+
size = self->depth;
557+
}
558+
507559
if (EXPECT(size == 0, 0)){
508560
return NULL;
509561
}
@@ -514,6 +566,11 @@ PyObject *SortedDict_next(SortedDict *self)
514566
} else {
515567
self->iterator_index++;
516568
Py_ssize_t size = PySequence_Fast_GET_SIZE(self->keys);
569+
// Respect depth limit for consistency with keys(), to_dict(), etc.
570+
if (self->depth > 0 && self->depth < size) {
571+
size = self->depth;
572+
}
573+
517574
if (size <= self->iterator_index) {
518575
self->iterator_index = -1;
519576
return NULL;

0 commit comments

Comments
 (0)