diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index 3e8a145e..ce1e2b98 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -37,7 +37,7 @@ TSNode ts_tree_root_node(const TSTree *self) {
 void ts_tree_edit(TSTree *self, const TSInputEdit *edit) {
   SubtreePool pool = ts_subtree_pool_new(0);
   self->root = ts_subtree_edit(self->root, edit, &pool);
-  assert(pool.tree_stack.capacity == 0 && pool.free_trees.capacity == 0);
+  ts_subtree_pool_delete(&pool);
 }
 
 TSRange *ts_tree_get_changed_ranges(const TSTree *self, const TSTree *other, uint32_t *count) {
diff --git a/test/helpers/record_alloc.cc b/test/helpers/record_alloc.cc
index 43e11abe..2e2ea648 100644
--- a/test/helpers/record_alloc.cc
+++ b/test/helpers/record_alloc.cc
@@ -1,6 +1,7 @@
 #include <stdlib.h>
 #include <map>
 #include <vector>
+#include <mutex>
 
 using std::map;
 using std::vector;
@@ -8,13 +9,16 @@ using std::vector;
 static bool _enabled = false;
 static size_t _allocation_count = 0;
 static map<void *, size_t> _outstanding_allocations;
+static std::mutex _outstanding_allocations_mutex;
+static bool _multi_threaded_mode = false;
 
 namespace record_alloc {
 
-void start() {
+void start(bool multi_threaded_mode) {
   _enabled = true;
   _allocation_count = 0;
   _outstanding_allocations.clear();
+  _multi_threaded_mode = multi_threaded_mode;
 }
 
 void stop() {
@@ -30,7 +34,11 @@ vector<size_t> outstanding_allocation_indices() {
 }
 
 size_t allocation_count() {
-  return _allocation_count;
+  size_t result;
+  _outstanding_allocations_mutex.lock();
+  result = _allocation_count;
+  _outstanding_allocations_mutex.unlock();
+  return result;
 }
 
 }  // namespace record_alloc
@@ -39,16 +47,20 @@ extern "C" {
 
 static void *record_allocation(void *result) {
   if (!_enabled) return result;
+  if (_multi_threaded_mode) _outstanding_allocations_mutex.lock();
   _outstanding_allocations[result] = _allocation_count;
   _allocation_count++;
+  if (_multi_threaded_mode) _outstanding_allocations_mutex.unlock();
   return result;
 }
 
 static void record_deallocation(void *pointer) {
+  if (_multi_threaded_mode) _outstanding_allocations_mutex.lock();
   auto entry = _outstanding_allocations.find(pointer);
   if (entry != _outstanding_allocations.end()) {
     _outstanding_allocations.erase(entry);
   }
+  if (_multi_threaded_mode) _outstanding_allocations_mutex.unlock();
 }
 
 void *ts_record_malloc(size_t size) {
diff --git a/test/helpers/record_alloc.h b/test/helpers/record_alloc.h
index 1f5968ac..f21876b4 100644
--- a/test/helpers/record_alloc.h
+++ b/test/helpers/record_alloc.h
@@ -5,7 +5,7 @@
 
 namespace record_alloc {
 
-void start();
+void start(bool multi_threaded_mode = false);
 void stop();
 void fail_at_allocation_index(size_t failure_index);
 std::vector<size_t> outstanding_allocation_indices();
diff --git a/test/runtime/tree_test.cc b/test/runtime/tree_test.cc
index 6d038220..ad9d3204 100644
--- a/test/runtime/tree_test.cc
+++ b/test/runtime/tree_test.cc
@@ -28,11 +28,17 @@ describe("Tree", [&]() {
   TSTree *tree;
 
   before_each([&]() {
+    record_alloc::start(true);
     parser = ts_parser_new();
+    tree = nullptr;
+    input = nullptr;
   });
 
   after_each([&]() {
+    if (tree) ts_tree_delete(tree);
+    if (input) delete input;
     ts_parser_delete(parser);
+    AssertThat(record_alloc::outstanding_allocation_indices(), IsEmpty());
   });
 
   auto assert_root_node = [&](const string &expected) {
@@ -85,6 +91,8 @@ describe("Tree", [&]() {
         }));
       }
 
+      ts_tree_delete(original_tree);
+
       for (auto &future : new_trees) {
         future.wait();
         TSTree *new_tree = future.get();
@@ -105,11 +113,6 @@ describe("Tree", [&]() {
       );
     });
 
-    after_each([&]() {
-      ts_tree_delete(tree);
-      delete input;
-    });
-
     auto get_changed_ranges_for_edit = [&](function<TSInputEdit()> fn) -> vector<TSRange> {
       TSInputEdit edit = fn();
       ts_tree_edit(tree, &edit);