4 files changed, 108 insertions, 3 deletions

diff --git a/examples/natmod/btree/btree_c.c b/examples/natmod/btree/btree_c.c
index f3e4790ed..5e8a34ac4 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 14dc8f348..2dd3a89b8 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 000000000..153f4e7d7
--- /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 000000000..e7b8d5672
--- /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'