Fix leaks in test/correctness/memoize.cpp

test/correctness/memoize.cpp

@@ -59,18 +59,22 @@ extern "C" HALIDE_EXPORT_SYMBOL int computed_eviction_key(int a) {
 }
 HalideExtern_1(int, computed_eviction_key, int);
 
-void simple_free(JITUserContext *user_context, void *ptr) {
-    free(ptr);
-}
+void *(*default_malloc)(JITUserContext *, size_t);
+void (*default_free)(JITUserContext *, void *);
 
-void *flakey_malloc(JITUserContext * /* user_context */, size_t x) {
+// A flaky allocator that wraps the built-in runtime one.
+void *flaky_malloc(JITUserContext *user_context, size_t x) {
     if ((rand() % 4) == 0) {
         return nullptr;
     } else {
-        return malloc(x);
+        return default_malloc(user_context, x);
     }
 }
 
+void simple_free(JITUserContext *user_context, void *ptr) {
+    return default_free(user_context, ptr);
+}
+
 bool error_occured = false;
 void record_error(JITUserContext *user_context, const char *msg) {
     error_occured = true;
@@ -584,6 +588,15 @@ int main(int argc, char **argv) {
         return 0;
     } else {
         // Test out of memory handling.
+
+        // Get the runtime's malloc and free. We need to use the ones
+        // in the runtime to ensure the matching free is called when
+        // we release all the runtimes at the end.
+        JITUserContext ctx;
+        Internal::JITSharedRuntime::populate_jit_handlers(&ctx, JITHandlers{});
+        default_malloc = ctx.handlers.custom_malloc;
+        default_free = ctx.handlers.custom_free;
+
         Param<float> val;
 
         Func count_calls;
@@ -600,7 +613,7 @@ int main(int argc, char **argv) {
 
         Pipeline pipe(g);
         pipe.jit_handlers().custom_error = record_error;
-        pipe.jit_handlers().custom_malloc = flakey_malloc;
+        pipe.jit_handlers().custom_malloc = flaky_malloc;
         pipe.jit_handlers().custom_free = simple_free;
 
         int total_errors = 0;
@@ -644,7 +657,8 @@ int main(int argc, char **argv) {
             }
         }
 
-        printf("In 100 attempts with flakey malloc, %d errors and %d full completions occured.\n", total_errors, completed);
+        printf("In 100 attempts with flaky malloc, %d errors and %d full completions occured.\n",
+               total_errors, completed);
     }
 
     {
@@ -733,6 +747,7 @@ int main(int argc, char **argv) {
 
         assert(call_count == 8);
     }
+    Internal::JITSharedRuntime::release_all();
 
     printf("Success!\n");
     return 0;

test/correctness/memoize_cloned.cpp

@@ -40,6 +40,8 @@ int main(int argc, char **argv) {
         return 1;
     }
 
+    Internal::JITSharedRuntime::release_all();
+
     printf("Success!\n");
     return 0;
 }