3 files changed, 46 insertions, 4 deletions
diff --git a/runtime/jit/jit.h b/runtime/jit/jit.h
index 643bc23da3..e73ba82278 100644
--- a/runtime/jit/jit.h
+++ b/runtime/jit/jit.h
@@ -67,6 +67,9 @@ class Jit {
   void DumpInfo(std::ostream& os);
   // Add a timing logger to cumulative_timings_.
   void AddTimingLogger(const TimingLogger& logger);
+  JitInstrumentationCache* GetInstrumentationCache() const {
+    return instrumentation_cache_.get();
+  }
 
  private:
   Jit();
diff --git a/runtime/jit/jit_instrumentation.cc b/runtime/jit/jit_instrumentation.cc
index d437dd5d56..4f4a97f455 100644
--- a/runtime/jit/jit_instrumentation.cc
+++ b/runtime/jit/jit_instrumentation.cc
@@ -24,11 +24,21 @@
 namespace art {
 namespace jit {
 
-class JitCompileTask : public Task {
+class JitCompileTask FINAL : public Task {
  public:
-  explicit JitCompileTask(ArtMethod* method) : method_(method) {}
+  explicit JitCompileTask(ArtMethod* method) : method_(method) {
+    ScopedObjectAccess soa(Thread::Current());
+    // Add a global ref to the class to prevent class unloading until compilation is done.
+    klass_ = soa.Vm()->AddGlobalRef(soa.Self(), method_->GetDeclaringClass());
+    CHECK(klass_ != nullptr);
+  }
+
+  ~JitCompileTask() {
+    ScopedObjectAccess soa(Thread::Current());
+    soa.Vm()->DeleteGlobalRef(soa.Self(), klass_);
+  }
 
-  virtual void Run(Thread* self) OVERRIDE {
+  void Run(Thread* self) OVERRIDE {
     ScopedObjectAccess soa(self);
     VLOG(jit) << "JitCompileTask compiling method " << PrettyMethod(method_);
     if (!Runtime::Current()->GetJit()->CompileMethod(method_, self)) {
@@ -36,12 +46,13 @@ class JitCompileTask : public Task {
     }
   }
 
-  virtual void Finalize() OVERRIDE {
+  void Finalize() OVERRIDE {
     delete this;
   }
 
  private:
   ArtMethod* const method_;
+  jobject klass_;
 
   DISALLOW_IMPLICIT_CONSTRUCTORS(JitCompileTask);
 };
@@ -104,5 +115,31 @@ void JitInstrumentationListener::InvokeVirtualOrInterface(Thread* thread,
   }
 }
 
+class WaitForCompilationToFinishTask FINAL : public Task {
+ public:
+  WaitForCompilationToFinishTask() : barrier_(0) {}
+
+  void Wait(Thread* self) {
+    barrier_.Increment(self, 1);
+  }
+
+  void Run(Thread* self ATTRIBUTE_UNUSED) OVERRIDE {}
+
+  void Finalize() OVERRIDE {
+    // Do this in Finalize since Finalize is called after Run by the thread pool.
+    barrier_.Pass(Thread::Current());
+  }
+
+ private:
+  Barrier barrier_;
+  DISALLOW_COPY_AND_ASSIGN(WaitForCompilationToFinishTask);
+};
+
+void JitInstrumentationCache::WaitForCompilationToFinish(Thread* self) {
+  std::unique_ptr<WaitForCompilationToFinishTask> task(new WaitForCompilationToFinishTask);
+  thread_pool_->AddTask(self, task.get());
+  task->Wait(self);
+}
+
 }  // namespace jit
 }  // namespace art
diff --git a/runtime/jit/jit_instrumentation.h b/runtime/jit/jit_instrumentation.h
index 6fdef6585d..9eb464b841 100644
--- a/runtime/jit/jit_instrumentation.h
+++ b/runtime/jit/jit_instrumentation.h
@@ -50,6 +50,8 @@ class JitInstrumentationCache {
       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_;