Properly delete the jit thread pool.

bug:25461989
bug:25462600
Change-Id: I273cf256285d01c085e4dea1d997955d029361b9
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 5afd28e..f691151 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -26,7 +26,6 @@
 #include "jit_instrumentation.h"
 #include "runtime.h"
 #include "runtime_options.h"
-#include "thread_list.h"
 #include "utils.h"
 
 namespace art {
@@ -145,7 +144,7 @@
 
 void Jit::DeleteThreadPool() {
   if (instrumentation_cache_.get() != nullptr) {
-    instrumentation_cache_->DeleteThreadPool();
+    instrumentation_cache_->DeleteThreadPool(Thread::Current());
   }
 }
 
@@ -164,16 +163,8 @@
 
 void Jit::CreateInstrumentationCache(size_t compile_threshold, size_t warmup_threshold) {
   CHECK_GT(compile_threshold, 0U);
-  ScopedSuspendAll ssa(__FUNCTION__);
-  // Add Jit interpreter instrumentation, tells the interpreter when to notify the jit to compile
-  // something.
   instrumentation_cache_.reset(
       new jit::JitInstrumentationCache(compile_threshold, warmup_threshold));
-  Runtime::Current()->GetInstrumentation()->AddListener(
-      new jit::JitInstrumentationListener(instrumentation_cache_.get()),
-      instrumentation::Instrumentation::kMethodEntered |
-      instrumentation::Instrumentation::kBackwardBranch |
-      instrumentation::Instrumentation::kInvokeVirtualOrInterface);
 }
 
 }  // namespace jit
diff --git a/runtime/jit/jit_instrumentation.cc b/runtime/jit/jit_instrumentation.cc
index 7931306..6531325 100644
--- a/runtime/jit/jit_instrumentation.cc
+++ b/runtime/jit/jit_instrumentation.cc
@@ -20,6 +20,7 @@
 #include "jit.h"
 #include "jit_code_cache.h"
 #include "scoped_thread_state_change.h"
+#include "thread_list.h"
 
 namespace art {
 namespace jit {
@@ -73,16 +74,48 @@
 JitInstrumentationCache::JitInstrumentationCache(size_t hot_method_threshold,
                                                  size_t warm_method_threshold)
     : hot_method_threshold_(hot_method_threshold),
-      warm_method_threshold_(warm_method_threshold) {
+      warm_method_threshold_(warm_method_threshold),
+      listener_(this) {
 }
 
 void JitInstrumentationCache::CreateThreadPool() {
+  // Create the thread pool before setting the instrumentation, so that
+  // when the threads stopped being suspended, they can use it directly.
+  // There is a DCHECK in the 'AddSamples' method to ensure the tread pool
+  // is not null when we instrument.
   thread_pool_.reset(new ThreadPool("Jit thread pool", 1));
+  thread_pool_->StartWorkers(Thread::Current());
+  {
+    // Add Jit interpreter instrumentation, tells the interpreter when
+    // to notify the jit to compile something.
+    ScopedSuspendAll ssa(__FUNCTION__);
+    Runtime::Current()->GetInstrumentation()->AddListener(
+        &listener_, JitInstrumentationListener::kJitEvents);
+  }
 }
 
-void JitInstrumentationCache::DeleteThreadPool() {
-  DCHECK(Runtime::Current()->IsShuttingDown(Thread::Current()));
-  thread_pool_.reset();
+void JitInstrumentationCache::DeleteThreadPool(Thread* self) {
+  DCHECK(Runtime::Current()->IsShuttingDown(self));
+  if (thread_pool_ != nullptr) {
+    // First remove the listener, to avoid having mutators enter
+    // 'AddSamples'.
+    ThreadPool* cache = nullptr;
+    {
+      ScopedSuspendAll ssa(__FUNCTION__);
+      Runtime::Current()->GetInstrumentation()->RemoveListener(
+          &listener_, JitInstrumentationListener::kJitEvents);
+      // Clear thread_pool_ field while the threads are suspended.
+      // A mutator in the 'AddSamples' method will check against it.
+      cache = thread_pool_.release();
+    }
+    cache->StopWorkers(self);
+    cache->RemoveAllTasks(self);
+    // We could just suspend all threads, but we know those threads
+    // will finish in a short period, so it's not worth adding a suspend logic
+    // here. Besides, this is only done for shutdown.
+    cache->Wait(self, false, false);
+    delete cache;
+  }
 }
 
 void JitInstrumentationCache::AddSamples(Thread* self, ArtMethod* method, size_t) {
@@ -91,25 +124,32 @@
   if (method->IsClassInitializer() || method->IsNative()) {
     return;
   }
-  if (thread_pool_.get() == nullptr) {
-    DCHECK(Runtime::Current()->IsShuttingDown(self));
-    return;
-  }
+  DCHECK(thread_pool_ != nullptr);
+
   uint16_t sample_count = method->IncrementCounter();
   if (sample_count == warm_method_threshold_) {
-    if (ProfilingInfo::Create(self, method, /* retry_allocation */ false)) {
+    bool success = ProfilingInfo::Create(self, method, /* retry_allocation */ false);
+    if (success) {
       VLOG(jit) << "Start profiling " << PrettyMethod(method);
-    } else {
+    }
+
+    if (thread_pool_ == nullptr) {
+      // Calling ProfilingInfo::Create might put us in a suspended state, which could
+      // lead to the thread pool being deleted when we are shutting down.
+      DCHECK(Runtime::Current()->IsShuttingDown(self));
+      return;
+    }
+
+    if (!success) {
       // We failed allocating. Instead of doing the collection on the Java thread, we push
       // an allocation to a compiler thread, that will do the collection.
       thread_pool_->AddTask(self, new JitCompileTask(method, JitCompileTask::kAllocateProfile));
-      thread_pool_->StartWorkers(self);
     }
   }
 
   if (sample_count == hot_method_threshold_) {
+    DCHECK(thread_pool_ != nullptr);
     thread_pool_->AddTask(self, new JitCompileTask(method, JitCompileTask::kCompile));
-    thread_pool_->StartWorkers(self);
   }
 }
 
@@ -118,6 +158,20 @@
   CHECK(instrumentation_cache_ != nullptr);
 }
 
+void JitInstrumentationListener::MethodEntered(Thread* thread,
+                                               mirror::Object* /*this_object*/,
+                                               ArtMethod* method,
+                                               uint32_t /*dex_pc*/) {
+  instrumentation_cache_->AddSamples(thread, method, 1);
+}
+
+void JitInstrumentationListener::BackwardBranch(Thread* thread,
+                                                ArtMethod* method,
+                                                int32_t dex_pc_offset) {
+  CHECK_LE(dex_pc_offset, 0);
+  instrumentation_cache_->AddSamples(thread, method, 1);
+}
+
 void JitInstrumentationListener::InvokeVirtualOrInterface(Thread* thread,
                                                           mirror::Object* this_object,
                                                           ArtMethod* caller,
@@ -138,7 +192,9 @@
 }
 
 void JitInstrumentationCache::WaitForCompilationToFinish(Thread* self) {
-  thread_pool_->Wait(self, false, false);
+  if (thread_pool_ != nullptr) {
+    thread_pool_->Wait(self, false, false);
+  }
 }
 
 }  // namespace jit
diff --git a/runtime/jit/jit_instrumentation.h b/runtime/jit/jit_instrumentation.h
index 9eb464b..1f96d59 100644
--- a/runtime/jit/jit_instrumentation.h
+++ b/runtime/jit/jit_instrumentation.h
@@ -31,7 +31,6 @@
 
 namespace art {
 namespace mirror {
-  class Class;
   class Object;
   class Throwable;
 }  // namespace mirror
@@ -42,24 +41,7 @@
 
 namespace jit {
 
-// Keeps track of which methods are hot.
-class JitInstrumentationCache {
- public:
-  JitInstrumentationCache(size_t hot_method_threshold, size_t warm_method_threshold);
-  void AddSamples(Thread* self, ArtMethod* method, size_t samples)
-      SHARED_REQUIRES(Locks::mutator_lock_);
-  void CreateThreadPool();
-  void DeleteThreadPool();
-  // Wait until there is no more pending compilation tasks.
-  void WaitForCompilationToFinish(Thread* self);
-
- private:
-  size_t hot_method_threshold_;
-  size_t warm_method_threshold_;
-  std::unique_ptr<ThreadPool> thread_pool_;
-
-  DISALLOW_IMPLICIT_CONSTRUCTORS(JitInstrumentationCache);
-};
+class JitInstrumentationCache;
 
 class JitInstrumentationListener : public instrumentation::InstrumentationListener {
  public:
@@ -67,9 +49,8 @@
 
   void MethodEntered(Thread* thread, mirror::Object* /*this_object*/,
                      ArtMethod* method, uint32_t /*dex_pc*/)
-      OVERRIDE SHARED_REQUIRES(Locks::mutator_lock_) {
-    instrumentation_cache_->AddSamples(thread, method, 1);
-  }
+      OVERRIDE SHARED_REQUIRES(Locks::mutator_lock_);
+
   void MethodExited(Thread* /*thread*/, mirror::Object* /*this_object*/,
                     ArtMethod* /*method*/, uint32_t /*dex_pc*/,
                     const JValue& /*return_value*/)
@@ -90,10 +71,7 @@
                   ArtMethod* /*method*/, uint32_t /*new_dex_pc*/) OVERRIDE { }
 
   void BackwardBranch(Thread* thread, ArtMethod* method, int32_t dex_pc_offset)
-      OVERRIDE SHARED_REQUIRES(Locks::mutator_lock_) {
-    CHECK_LE(dex_pc_offset, 0);
-    instrumentation_cache_->AddSamples(thread, method, 1);
-  }
+      OVERRIDE SHARED_REQUIRES(Locks::mutator_lock_);
 
   void InvokeVirtualOrInterface(Thread* thread,
                                 mirror::Object* this_object,
@@ -102,12 +80,37 @@
                                 ArtMethod* callee)
       OVERRIDE SHARED_REQUIRES(Locks::mutator_lock_);
 
+  static constexpr uint32_t kJitEvents =
+      instrumentation::Instrumentation::kMethodEntered |
+      instrumentation::Instrumentation::kBackwardBranch |
+      instrumentation::Instrumentation::kInvokeVirtualOrInterface;
+
  private:
   JitInstrumentationCache* const instrumentation_cache_;
 
   DISALLOW_IMPLICIT_CONSTRUCTORS(JitInstrumentationListener);
 };
 
+// Keeps track of which methods are hot.
+class JitInstrumentationCache {
+ public:
+  JitInstrumentationCache(size_t hot_method_threshold, size_t warm_method_threshold);
+  void AddSamples(Thread* self, ArtMethod* method, size_t samples)
+      SHARED_REQUIRES(Locks::mutator_lock_);
+  void CreateThreadPool();
+  void DeleteThreadPool(Thread* self);
+  // Wait until there is no more pending compilation tasks.
+  void WaitForCompilationToFinish(Thread* self);
+
+ private:
+  size_t hot_method_threshold_;
+  size_t warm_method_threshold_;
+  JitInstrumentationListener listener_;
+  std::unique_ptr<ThreadPool> thread_pool_;
+
+  DISALLOW_IMPLICIT_CONSTRUCTORS(JitInstrumentationCache);
+};
+
 }  // namespace jit
 }  // namespace art
 
diff --git a/runtime/thread_pool.cc b/runtime/thread_pool.cc
index 0527d3a..5a4dfb8 100644
--- a/runtime/thread_pool.cc
+++ b/runtime/thread_pool.cc
@@ -82,6 +82,11 @@
   }
 }
 
+void ThreadPool::RemoveAllTasks(Thread* self) {
+  MutexLock mu(self, task_queue_lock_);
+  tasks_.clear();
+}
+
 ThreadPool::ThreadPool(const char* name, size_t num_threads)
   : name_(name),
     task_queue_lock_("task queue lock"),
diff --git a/runtime/thread_pool.h b/runtime/thread_pool.h
index a2338d6..6cd4ad3 100644
--- a/runtime/thread_pool.h
+++ b/runtime/thread_pool.h
@@ -91,6 +91,9 @@
   // after running it, it is the caller's responsibility.
   void AddTask(Thread* self, Task* task) REQUIRES(!task_queue_lock_);
 
+  // Remove all tasks in the queue.
+  void RemoveAllTasks(Thread* self) REQUIRES(!task_queue_lock_);
+
   ThreadPool(const char* name, size_t num_threads);
   virtual ~ThreadPool();