diff --git a/spec/helpers/record_alloc.cc b/spec/helpers/record_alloc.cc
index 3e5a2961..4b057a61 100644
--- a/spec/helpers/record_alloc.cc
+++ b/spec/helpers/record_alloc.cc
@@ -42,26 +42,26 @@ size_t allocation_count() {
 
 }  // namespace record_alloc
 
+static bool can_allocate() {
+  if (!_enabled)
+    return true;
+
+  return _allocation_count < _allocation_failure_index;
+}
+
 static void *record_allocation(void *result) {
   if (!_enabled)
     return result;
 
-  if (_allocation_count > _allocation_failure_index) {
-    free(result);
-    Assert::Failure("Allocated after a previous allocation failed!");
-  }
-
-  if (_allocation_count == _allocation_failure_index) {
-    _allocation_count++;
-    free(result);
-    return nullptr;
-  }
-
   _outstanding_allocations[result] = _allocation_count;
   _allocation_count++;
   return result;
 }
 
+static void record_allocation_failure() {
+  _allocation_count++;
+}
+
 static void record_deallocation(void *pointer) {
   if (!_enabled)
     return;
@@ -75,16 +75,31 @@ static void record_deallocation(void *pointer) {
 extern "C" {
 
 void *ts_record_malloc(size_t size) {
-  return record_allocation(malloc(size));
+  if (can_allocate()) {
+    return record_allocation(malloc(size));
+  } else {
+    record_allocation_failure();
+    return NULL;
+  }
 }
 
 void *ts_record_realloc(void *pointer, size_t size) {
-  record_deallocation(pointer);
-  return record_allocation(realloc(pointer, size));
+  if (can_allocate()) {
+    record_deallocation(pointer);
+    return record_allocation(realloc(pointer, size));
+  } else {
+    record_allocation_failure();
+    return NULL;
+  }
 }
 
 void *ts_record_calloc(size_t count, size_t size) {
-  return record_allocation(calloc(count, size));
+  if (can_allocate()) {
+    return record_allocation(calloc(count, size));
+  } else {
+    record_allocation_failure();
+    return NULL;
+  }
 }
 
 void ts_record_free(void *pointer) {
@@ -92,4 +107,10 @@ void ts_record_free(void *pointer) {
   record_deallocation(pointer);
 }
 
+bool ts_record_allocations_toggle(bool value) {
+  bool previous_value = _enabled;
+  _enabled = value;
+  return previous_value;
+}
+
 }
diff --git a/src/runtime/alloc.h b/src/runtime/alloc.h
index 6aee72a1..81b5f18f 100644
--- a/src/runtime/alloc.h
+++ b/src/runtime/alloc.h
@@ -11,6 +11,7 @@ void *ts_record_malloc(size_t);
 void *ts_record_calloc(size_t, size_t);
 void *ts_record_realloc(void *, size_t);
 void ts_record_free(void *);
+bool ts_record_allocations_toggle(bool);
 
 static inline void *ts_malloc(size_t size) {
   return ts_record_malloc(size);
@@ -28,10 +29,18 @@ static inline void ts_free(void *buffer) {
   return ts_record_free(buffer);
 }
 
+static inline bool ts_toggle_allocation_recording(bool value) {
+  return ts_record_allocations_toggle(value);
+}
+
 #else
 
 #include <stdlib.h>
 
+static inline bool ts_toggle_allocation_recording(bool value) {
+  return false;
+}
+
 static inline void *ts_malloc(size_t size) {
   return malloc(size);
 }
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index c5cf561f..42cae0bf 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -494,13 +494,16 @@ void ts_stack_clear(Stack *self) {
   array_push(&self->heads, ((StackHead){ self->base_node, false }));
 }
 
-int ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
+bool ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
+  bool was_recording_allocations = ts_toggle_allocation_recording(false);
+  if (!f)
+    f = stderr;
+
   fprintf(f, "digraph stack {\n");
   fprintf(f, "rankdir=\"RL\";\n");
   fprintf(f, "edge [arrowhead=none]\n");
 
-  Array(StackNode *)visited_nodes;
-  array_init(&visited_nodes);
+  Array(StackNode *) visited_nodes = array_new();
 
   array_clear(&self->pop_paths);
   for (size_t i = 0; i < self->heads.size; i++) {
@@ -511,7 +514,8 @@ int ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
     fprintf(
       f, "node_head_%lu -> node_%p [label=%lu, fontcolor=blue, weight=10000]\n",
       i, node, i);
-    array_push(&self->pop_paths, ((PopPath){.node = node }));
+    if (!array_push(&self->pop_paths, ((PopPath){.node = node })))
+      goto error;
   }
 
   bool all_paths_done = false;
@@ -582,7 +586,12 @@ int ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
   fprintf(f, "}\n");
 
   array_delete(&visited_nodes);
+  ts_toggle_allocation_recording(was_recording_allocations);
+  return true;
 
 error:
-  return -1;
+  ts_toggle_allocation_recording(was_recording_allocations);
+  if (visited_nodes.contents)
+    array_delete(&visited_nodes);
+  return false;
 }
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index ae184274..31873e36 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -121,7 +121,7 @@ void ts_stack_remove_version(Stack *, StackVersion);
  */
 void ts_stack_clear(Stack *);
 
-int ts_stack_print_dot_graph(Stack *, const char **, FILE *);
+bool ts_stack_print_dot_graph(Stack *, const char **, FILE *);
 
 #ifdef __cplusplus
 }