From 73c58150f53d9d828c4fc8fb455cca6831eb8ddd Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Mon, 20 Apr 2020 23:04:40 +1000
Subject: [PATCH] extmod/modbtree: Retain reference to underlying stream so
 it's not GC'd.

For ports that have a system malloc which is not garbage collected (eg
unix, esp32), the stream object for the DB must be retained separately to
prevent it from being reclaimed by the MicroPython GC (because the
berkeley-db library uses malloc to allocate the DB structure which stores
the only reference to the stream).

Although in some cases the user code will explicitly retain a reference to
the underlying stream because it needs to call close() on it, this is not
always the case, eg in cases where the DB is intended to live forever.

Fixes issue #5940.
---
 examples/natmod/btree/btree_c.c |  2 +-
 extmod/modbtree.c               |  6 ++-
 tests/extmod/btree_gc.py        | 23 ++++++++++
 tests/extmod/btree_gc.py.exp    | 80 +++++++++++++++++++++++++++++++++
 4 files changed, 108 insertions(+), 3 deletions(-)
 create mode 100644 tests/extmod/btree_gc.py
 create mode 100644 tests/extmod/btree_gc.py.exp

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'