From 44847c98a5d75e8b013892b41e5934d94520909c Mon Sep 17 00:00:00 2001
From: Chris Cummins <chrisc.101@gmail.com>
Date: Wed, 2 Mar 2022 16:45:43 +0000
Subject: [PATCH] [llvm] Make sure that BenchmarkFactory::close() is called.

This adds an explicit call to BenchmarkFactory::close() on the global
singleton, since otherwise it may not be called.

Fixes #582.
Issue #591.
---
 compiler_gym/envs/llvm/service/BUILD          |  1 +
 compiler_gym/envs/llvm/service/CMakeLists.txt |  1 +
 compiler_gym/envs/llvm/service/RunService.cc  | 11 +++++++++++
 3 files changed, 13 insertions(+)

diff --git a/compiler_gym/envs/llvm/service/BUILD b/compiler_gym/envs/llvm/service/BUILD
index 6b1d21c6ba..294f037e9f 100644
--- a/compiler_gym/envs/llvm/service/BUILD
+++ b/compiler_gym/envs/llvm/service/BUILD
@@ -63,6 +63,7 @@ cc_binary(
     name = "compiler_gym-llvm-service-prelinked",
     srcs = ["RunService.cc"],
     deps = [
+        ":BenchmarkFactory",
         ":LlvmSession",
         "//compiler_gym/service/runtime:cc_runtime",
     ],
diff --git a/compiler_gym/envs/llvm/service/CMakeLists.txt b/compiler_gym/envs/llvm/service/CMakeLists.txt
index c267b40d7a..62f23f9f7a 100644
--- a/compiler_gym/envs/llvm/service/CMakeLists.txt
+++ b/compiler_gym/envs/llvm/service/CMakeLists.txt
@@ -21,6 +21,7 @@ cg_cc_binary(
     "RunService.cc"
   DEPS
     ::LlvmSession
+    ::BenchmarkFactory
     compiler_gym::service::runtime::cc_runtime
 )
 
diff --git a/compiler_gym/envs/llvm/service/RunService.cc b/compiler_gym/envs/llvm/service/RunService.cc
index 0025ed77ea..508ae456e6 100644
--- a/compiler_gym/envs/llvm/service/RunService.cc
+++ b/compiler_gym/envs/llvm/service/RunService.cc
@@ -2,6 +2,7 @@
 //
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
+#include "compiler_gym/envs/llvm/service/BenchmarkFactory.h"
 #include "compiler_gym/envs/llvm/service/LlvmSession.h"
 #include "compiler_gym/service/runtime/Runtime.h"
 #include "llvm/InitializePasses.h"
@@ -61,5 +62,15 @@ int main(int argc, char** argv) {
   initLlvm();
   const auto ret = createAndRunCompilerGymService<LlvmSession>(argc, argv, usage);
 
+  // NOTE(github.com/facebookresearch/CompilerGym/issues/582): We need to make
+  // sure that BenchmarkFactory::close() is called on the global singleton
+  // instance, so that the temporary scratch directories are tidied up.
+  //
+  // TODO(github.com/facebookresearch/CompilerGym/issues/591): Once the runtime
+  // has been refactored to support intra-session mutable state, this singleton
+  // can be replaced by a member variable that is closed on
+  // CompilerGymServiceContext::shutdown().
+  BenchmarkFactory::getSingleton(FLAGS_working_dir).close();
+
   return ret;
 }