diff --git a/examples/natmod/btree/btree_c.c b/examples/natmod/btree/btree_c.c index f3e4790ed2..5e8a34ac40 100644 --- a/examples/natmod/btree/btree_c.c +++ b/examples/natmod/btree/btree_c.c @@ -115,7 +115,7 @@ STATIC mp_obj_t btree_open(size_t n_args, const mp_obj_t *args) { mp_raise_OSError(native_errno); } - return MP_OBJ_FROM_PTR(btree_new(db)); + return MP_OBJ_FROM_PTR(btree_new(db, args[0])); } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_open_obj, 5, 5, btree_open); diff --git a/extmod/modbtree.c b/extmod/modbtree.c index 14dc8f348a..2dd3a89b89 100644 --- a/extmod/modbtree.c +++ b/extmod/modbtree.c @@ -39,6 +39,7 @@ typedef struct _mp_obj_btree_t { mp_obj_base_t base; + mp_obj_t stream; // retain a reference to prevent GC from reclaiming it DB *db; mp_obj_t start_key; mp_obj_t end_key; @@ -65,9 +66,10 @@ void __dbpanic(DB *db) { mp_printf(&mp_plat_print, "__dbpanic(%p)\n", db); } -STATIC mp_obj_btree_t *btree_new(DB *db) { +STATIC mp_obj_btree_t *btree_new(DB *db, mp_obj_t stream) { mp_obj_btree_t *o = m_new_obj(mp_obj_btree_t); o->base.type = &btree_type; + o->stream = stream; o->db = db; o->start_key = mp_const_none; o->end_key = mp_const_none; @@ -361,7 +363,7 @@ STATIC mp_obj_t mod_btree_open(size_t n_args, const mp_obj_t *pos_args, mp_map_t if (db == NULL) { mp_raise_OSError(errno); } - return MP_OBJ_FROM_PTR(btree_new(db)); + return MP_OBJ_FROM_PTR(btree_new(db, pos_args[0])); } STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mod_btree_open_obj, 1, mod_btree_open); diff --git a/tests/extmod/btree_gc.py b/tests/extmod/btree_gc.py new file mode 100644 index 0000000000..153f4e7d78 --- /dev/null +++ b/tests/extmod/btree_gc.py @@ -0,0 +1,23 @@ +# Test btree interaction with the garbage collector. + +try: + import btree, uio, gc +except ImportError: + print("SKIP") + raise SystemExit + +N = 80 + +# Create a BytesIO but don't keep a reference to it. +db = btree.open(uio.BytesIO(), pagesize=512) + +# Overwrite lots of the Python stack to make sure no reference to the BytesIO remains. +x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + +# Write lots of key/value pairs, which fill up the DB and also allocate temporary heap +# memory due to the string addition, and do a GC collect to verify that the BytesIO +# is not collected. +for i in range(N): + db[b"thekey" + str(i)] = b"thelongvalue" + str(i) + print(db[b"thekey" + str(i)]) + gc.collect() diff --git a/tests/extmod/btree_gc.py.exp b/tests/extmod/btree_gc.py.exp new file mode 100644 index 0000000000..e7b8d56729 --- /dev/null +++ b/tests/extmod/btree_gc.py.exp @@ -0,0 +1,80 @@ +b'thelongvalue0' +b'thelongvalue1' +b'thelongvalue2' +b'thelongvalue3' +b'thelongvalue4' +b'thelongvalue5' +b'thelongvalue6' +b'thelongvalue7' +b'thelongvalue8' +b'thelongvalue9' +b'thelongvalue10' +b'thelongvalue11' +b'thelongvalue12' +b'thelongvalue13' +b'thelongvalue14' +b'thelongvalue15' +b'thelongvalue16' +b'thelongvalue17' +b'thelongvalue18' +b'thelongvalue19' +b'thelongvalue20' +b'thelongvalue21' +b'thelongvalue22' +b'thelongvalue23' +b'thelongvalue24' +b'thelongvalue25' +b'thelongvalue26' +b'thelongvalue27' +b'thelongvalue28' +b'thelongvalue29' +b'thelongvalue30' +b'thelongvalue31' +b'thelongvalue32' +b'thelongvalue33' +b'thelongvalue34' +b'thelongvalue35' +b'thelongvalue36' +b'thelongvalue37' +b'thelongvalue38' +b'thelongvalue39' +b'thelongvalue40' +b'thelongvalue41' +b'thelongvalue42' +b'thelongvalue43' +b'thelongvalue44' +b'thelongvalue45' +b'thelongvalue46' +b'thelongvalue47' +b'thelongvalue48' +b'thelongvalue49' +b'thelongvalue50' +b'thelongvalue51' +b'thelongvalue52' +b'thelongvalue53' +b'thelongvalue54' +b'thelongvalue55' +b'thelongvalue56' +b'thelongvalue57' +b'thelongvalue58' +b'thelongvalue59' +b'thelongvalue60' +b'thelongvalue61' +b'thelongvalue62' +b'thelongvalue63' +b'thelongvalue64' +b'thelongvalue65' +b'thelongvalue66' +b'thelongvalue67' +b'thelongvalue68' +b'thelongvalue69' +b'thelongvalue70' +b'thelongvalue71' +b'thelongvalue72' +b'thelongvalue73' +b'thelongvalue74' +b'thelongvalue75' +b'thelongvalue76' +b'thelongvalue77' +b'thelongvalue78' +b'thelongvalue79'