7 files changed, 308 insertions, 6 deletions

diff --git a/openjdkjvmti/OpenjdkJvmTi.cc b/openjdkjvmti/OpenjdkJvmTi.cc
index 027635bbb5..a0c7f40b6f 100644
--- a/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/openjdkjvmti/OpenjdkJvmTi.cc
@@ -1535,6 +1535,7 @@ extern "C" bool ArtPlugin_Initialize() {
   MethodUtil::Register(&gEventHandler);
   SearchUtil::Register();
   HeapUtil::Register();
+  Transformer::Setup();
 
   {
     // Make sure we can deopt anything we need to.
diff --git a/openjdkjvmti/fixed_up_dex_file.cc b/openjdkjvmti/fixed_up_dex_file.cc
index a8d2a37fa6..e9522b3984 100644
--- a/openjdkjvmti/fixed_up_dex_file.cc
+++ b/openjdkjvmti/fixed_up_dex_file.cc
@@ -45,14 +45,13 @@
 
 namespace openjdkjvmti {
 
-static void RecomputeDexChecksum(art::DexFile* dex_file)
-    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+static void RecomputeDexChecksum(art::DexFile* dex_file) {
   reinterpret_cast<art::DexFile::Header*>(const_cast<uint8_t*>(dex_file->Begin()))->checksum_ =
       dex_file->CalculateChecksum();
 }
 
-static void DoDexUnquicken(const art::DexFile& new_dex_file, const art::DexFile& original_dex_file)
-    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+static void DoDexUnquicken(const art::DexFile& new_dex_file,
+                           const art::DexFile& original_dex_file) {
   const art::OatDexFile* oat_dex = original_dex_file.GetOatDexFile();
   if (oat_dex == nullptr) {
     return;
diff --git a/openjdkjvmti/fixed_up_dex_file.h b/openjdkjvmti/fixed_up_dex_file.h
index 7f05a2930a..594e8a7358 100644
--- a/openjdkjvmti/fixed_up_dex_file.h
+++ b/openjdkjvmti/fixed_up_dex_file.h
@@ -50,8 +50,7 @@ namespace openjdkjvmti {
 class FixedUpDexFile {
  public:
   static std::unique_ptr<FixedUpDexFile> Create(const art::DexFile& original,
-                                                const char* descriptor)
-      REQUIRES_SHARED(art::Locks::mutator_lock_);
+                                                const char* descriptor);
 
   const art::DexFile& GetDexFile() {
     return *dex_file_;
diff --git a/openjdkjvmti/ti_class_definition.cc b/openjdkjvmti/ti_class_definition.cc
index 7f2f80009a..1b641cd905 100644
--- a/openjdkjvmti/ti_class_definition.cc
+++ b/openjdkjvmti/ti_class_definition.cc
@@ -45,6 +45,31 @@
 
 namespace openjdkjvmti {
 
+void ArtClassDefinition::InitializeMemory() const {
+  DCHECK(art::MemMap::kCanReplaceMapping);
+  VLOG(signals) << "Initializing de-quickened memory for dex file of " << name_;
+  CHECK(dex_data_mmap_ != nullptr);
+  CHECK(temp_mmap_ != nullptr);
+  CHECK_EQ(dex_data_mmap_->GetProtect(), PROT_NONE);
+  CHECK_EQ(temp_mmap_->GetProtect(), PROT_READ | PROT_WRITE);
+
+  std::string desc = std::string("L") + name_ + ";";
+  std::unique_ptr<FixedUpDexFile>
+      fixed_dex_file(FixedUpDexFile::Create(*initial_dex_file_unquickened_, desc.c_str()));
+  CHECK(fixed_dex_file.get() != nullptr);
+  CHECK_LE(fixed_dex_file->Size(), temp_mmap_->Size());
+  CHECK_EQ(temp_mmap_->Size(), dex_data_mmap_->Size());
+  // Copy the data to the temp mmap.
+  memcpy(temp_mmap_->Begin(), fixed_dex_file->Begin(), fixed_dex_file->Size());
+
+  // Move the mmap atomically.
+  art::MemMap* source = temp_mmap_.release();
+  std::string error;
+  CHECK(dex_data_mmap_->ReplaceWith(&source, &error)) << "Failed to replace mmap for "
+                                                      << name_ << " because " << error;
+  CHECK(dex_data_mmap_->Protect(PROT_READ));
+}
+
 bool ArtClassDefinition::IsModified() const {
   // RedefineClasses calls always are 'modified' since they need to change the current_dex_file of
   // the class.
@@ -58,6 +83,27 @@ bool ArtClassDefinition::IsModified() const {
     return false;
   }
 
+  // The dex_data_ was never touched by the agents.
+  if (dex_data_mmap_ != nullptr && dex_data_mmap_->GetProtect() == PROT_NONE) {
+    if (current_dex_file_.data() == dex_data_mmap_->Begin()) {
+      // the dex_data_ looks like it changed (not equal to current_dex_file_) but we never
+      // initialized the dex_data_mmap_. This means the new_dex_data was filled in without looking
+      // at the initial dex_data_.
+      return true;
+    } else if (dex_data_.data() == dex_data_mmap_->Begin()) {
+      // The dex file used to have modifications but they were not added again.
+      return true;
+    } else {
+      // It's not clear what happened. It's possible that the agent got the current dex file data
+      // from some other source so we need to initialize everything to see if it is the same.
+      VLOG(signals) << "Lazy dex file for " << name_ << " was never touched but the dex_data_ is "
+                    << "changed! Need to initialize the memory to see if anything changed";
+      InitializeMemory();
+    }
+  }
+
+  // We can definitely read current_dex_file_ and dex_file_ without causing page faults.
+
   // Check if the dex file we want to set is the same as the current one.
   // Unfortunately we need to do this check even if no modifications have been done since it could
   // be that agents were removed in the mean-time so we still have a different dex file. The dex
@@ -194,6 +240,53 @@ void ArtClassDefinition::InitWithDex(GetOriginalDexFile get_original,
                                      const art::DexFile* quick_dex) {
   art::Thread* self = art::Thread::Current();
   DCHECK(quick_dex != nullptr);
+  if (art::MemMap::kCanReplaceMapping && kEnableOnDemandDexDequicken) {
+    size_t dequick_size = quick_dex->GetDequickenedSize();
+    std::string mmap_name("anon-mmap-for-redefine: ");
+    mmap_name += name_;
+    std::string error;
+    dex_data_mmap_.reset(art::MemMap::MapAnonymous(mmap_name.c_str(),
+                                                   nullptr,
+                                                   dequick_size,
+                                                   PROT_NONE,
+                                                   /*low_4gb*/ false,
+                                                   /*reuse*/ false,
+                                                   &error));
+    mmap_name += "-TEMP";
+    temp_mmap_.reset(art::MemMap::MapAnonymous(mmap_name.c_str(),
+                                               nullptr,
+                                               dequick_size,
+                                               PROT_READ | PROT_WRITE,
+                                               /*low_4gb*/ false,
+                                               /*reuse*/ false,
+                                               &error));
+    if (UNLIKELY(dex_data_mmap_ != nullptr && temp_mmap_ != nullptr)) {
+      // Need to save the initial dexfile so we don't need to search for it in the fault-handler.
+      initial_dex_file_unquickened_ = quick_dex;
+      dex_data_ = art::ArrayRef<const unsigned char>(dex_data_mmap_->Begin(),
+                                                     dex_data_mmap_->Size());
+      if (from_class_ext_) {
+        // We got initial from class_ext so the current one must have undergone redefinition so no
+        // cdex or quickening stuff.
+        // We can only do this if it's not a first load.
+        DCHECK(klass_ != nullptr);
+        const art::DexFile& cur_dex = self->DecodeJObject(klass_)->AsClass()->GetDexFile();
+        current_dex_file_ = art::ArrayRef<const unsigned char>(cur_dex.Begin(), cur_dex.Size());
+      } else {
+        // This class hasn't been redefined before. The dequickened current data is the same as the
+        // dex_data_mmap_ when it's filled it. We don't need to copy anything because the mmap will
+        // not be cleared until after everything is done.
+        current_dex_file_ = art::ArrayRef<const unsigned char>(dex_data_mmap_->Begin(),
+                                                               dequick_size);
+      }
+      return;
+    }
+  }
+  dex_data_mmap_.reset(nullptr);
+  temp_mmap_.reset(nullptr);
+  // Failed to mmap a large enough area (or on-demand dequickening was disabled). This is
+  // unfortunate. Since currently the size is just a guess though we might as well try to do it
+  // manually.
   get_original(/*out*/&dex_data_memory_);
   dex_data_ = art::ArrayRef<const unsigned char>(dex_data_memory_);
   if (from_class_ext_) {
diff --git a/openjdkjvmti/ti_class_definition.h b/openjdkjvmti/ti_class_definition.h
index 00847394e7..31c3611e72 100644
--- a/openjdkjvmti/ti_class_definition.h
+++ b/openjdkjvmti/ti_class_definition.h
@@ -32,9 +32,14 @@
 #ifndef ART_OPENJDKJVMTI_TI_CLASS_DEFINITION_H_
 #define ART_OPENJDKJVMTI_TI_CLASS_DEFINITION_H_
 
+#include <stddef.h>
+#include <sys/mman.h>
+#include <sys/types.h>
+
 #include "art_jvmti.h"
 
 #include "base/array_ref.h"
+#include "mem_map.h"
 
 namespace openjdkjvmti {
 
@@ -43,13 +48,20 @@ namespace openjdkjvmti {
 // redefinition/retransformation function that created it.
 class ArtClassDefinition {
  public:
+  // If we support doing a on-demand dex-dequickening using signal handlers.
+  static constexpr bool kEnableOnDemandDexDequicken = true;
+
   ArtClassDefinition()
       : klass_(nullptr),
         loader_(nullptr),
         name_(),
         protection_domain_(nullptr),
+        dex_data_mmap_(nullptr),
+        temp_mmap_(nullptr),
         dex_data_memory_(),
+        initial_dex_file_unquickened_(nullptr),
         dex_data_(),
+        current_dex_memory_(),
         current_dex_file_(),
         redefined_(false),
         from_class_ext_(false),
@@ -87,6 +99,12 @@ class ArtClassDefinition {
     }
   }
 
+  bool ContainsAddress(uintptr_t ptr) const {
+    return dex_data_mmap_ != nullptr &&
+        reinterpret_cast<uintptr_t>(dex_data_mmap_->Begin()) <= ptr &&
+        reinterpret_cast<uintptr_t>(dex_data_mmap_->End()) > ptr;
+  }
+
   bool IsModified() const REQUIRES_SHARED(art::Locks::mutator_lock_);
 
   bool IsInitialized() const {
@@ -108,6 +126,13 @@ class ArtClassDefinition {
     return name_;
   }
 
+  bool IsLazyDefinition() const {
+    DCHECK(IsInitialized());
+    return dex_data_mmap_ != nullptr &&
+        dex_data_.data() == dex_data_mmap_->Begin() &&
+        dex_data_mmap_->GetProtect() == PROT_NONE;
+  }
+
   jobject GetProtectionDomain() const {
     DCHECK(IsInitialized());
     return protection_domain_;
@@ -118,6 +143,8 @@ class ArtClassDefinition {
     return dex_data_;
   }
 
+  void InitializeMemory() const;
+
  private:
   jvmtiError InitCommon(art::Thread* self, jclass klass);
 
@@ -130,9 +157,17 @@ class ArtClassDefinition {
   std::string name_;
   jobject protection_domain_;
 
+  // Mmap that will be filled with the original-dex-file lazily if it needs to be de-quickened or
+  // de-compact-dex'd
+  mutable std::unique_ptr<art::MemMap> dex_data_mmap_;
+  // This is a temporary mmap we will use to be able to fill the dex file data atomically.
+  mutable std::unique_ptr<art::MemMap> temp_mmap_;
+
   // A unique_ptr to the current dex_data if it needs to be cleaned up.
   std::vector<unsigned char> dex_data_memory_;
 
+  const art::DexFile* initial_dex_file_unquickened_;
+
   // A ref to the current dex data. This is either dex_data_memory_, or current_dex_file_. This is
   // what the dex file will be turned into.
   art::ArrayRef<const unsigned char> dex_data_;
diff --git a/openjdkjvmti/transform.cc b/openjdkjvmti/transform.cc
index af838d62b9..7b19a24ba4 100644
--- a/openjdkjvmti/transform.cc
+++ b/openjdkjvmti/transform.cc
@@ -29,6 +29,9 @@
  * questions.
  */
 
+#include <stddef.h>
+#include <sys/types.h>
+
 #include <unordered_map>
 #include <unordered_set>
 
@@ -41,6 +44,7 @@
 #include "dex/dex_file_types.h"
 #include "dex/utf.h"
 #include "events-inl.h"
+#include "fault_handler.h"
 #include "gc_root-inl.h"
 #include "globals.h"
 #include "jni_env_ext-inl.h"
@@ -63,6 +67,174 @@
 
 namespace openjdkjvmti {
 
+// A FaultHandler that will deal with initializing ClassDefinitions when they are actually needed.
+class TransformationFaultHandler FINAL : public art::FaultHandler {
+ public:
+  explicit TransformationFaultHandler(art::FaultManager* manager)
+      : art::FaultHandler(manager),
+        uninitialized_class_definitions_lock_("JVMTI Initialized class definitions lock",
+                                              art::LockLevel::kSignalHandlingLock),
+        class_definition_initialized_cond_("JVMTI Initialized class definitions condition",
+                                           uninitialized_class_definitions_lock_) {
+    manager->AddHandler(this, /* generated_code */ false);
+  }
+
+  ~TransformationFaultHandler() {
+    art::MutexLock mu(art::Thread::Current(), uninitialized_class_definitions_lock_);
+    uninitialized_class_definitions_.clear();
+  }
+
+  bool Action(int sig, siginfo_t* siginfo, void* context ATTRIBUTE_UNUSED) OVERRIDE {
+    DCHECK_EQ(sig, SIGSEGV);
+    art::Thread* self = art::Thread::Current();
+    if (UNLIKELY(uninitialized_class_definitions_lock_.IsExclusiveHeld(self))) {
+      if (self != nullptr) {
+        LOG(FATAL) << "Recursive call into Transformation fault handler!";
+        UNREACHABLE();
+      } else {
+        LOG(ERROR) << "Possible deadlock due to recursive signal delivery of segv.";
+      }
+    }
+    uintptr_t ptr = reinterpret_cast<uintptr_t>(siginfo->si_addr);
+    ArtClassDefinition* res = nullptr;
+
+    {
+      // NB Technically using a mutex and condition variables here is non-posix compliant but
+      // everything should be fine since both glibc and bionic implementations of mutexs and
+      // condition variables work fine so long as the thread was not interrupted during a
+      // lock/unlock (which it wasn't) on all architectures we care about.
+      art::MutexLock mu(self, uninitialized_class_definitions_lock_);
+      auto it = std::find_if(uninitialized_class_definitions_.begin(),
+                             uninitialized_class_definitions_.end(),
+                             [&](const auto op) { return op->ContainsAddress(ptr); });
+      if (it != uninitialized_class_definitions_.end()) {
+        res = *it;
+        // Remove the class definition.
+        uninitialized_class_definitions_.erase(it);
+        // Put it in the initializing list
+        initializing_class_definitions_.push_back(res);
+      } else {
+        // Wait for the ptr to be initialized (if it is currently initializing).
+        while (DefinitionIsInitializing(ptr)) {
+          WaitForClassInitializationToFinish();
+        }
+        // Return true (continue with user code) if we find that the definition has been
+        // initialized. Return false (continue on to next signal handler) if the definition is not
+        // initialized or found.
+        return std::find_if(initialized_class_definitions_.begin(),
+                            initialized_class_definitions_.end(),
+                            [&](const auto op) { return op->ContainsAddress(ptr); }) !=
+            uninitialized_class_definitions_.end();
+      }
+    }
+
+    VLOG(signals) << "Lazy initialization of dex file for transformation of " << res->GetName()
+                  << " during SEGV";
+    res->InitializeMemory();
+
+    {
+      art::MutexLock mu(self, uninitialized_class_definitions_lock_);
+      // Move to initialized state and notify waiters.
+      initializing_class_definitions_.erase(std::find(initializing_class_definitions_.begin(),
+                                                      initializing_class_definitions_.end(),
+                                                      res));
+      initialized_class_definitions_.push_back(res);
+      class_definition_initialized_cond_.Broadcast(self);
+    }
+
+    return true;
+  }
+
+  void RemoveDefinition(ArtClassDefinition* def) REQUIRES(!uninitialized_class_definitions_lock_) {
+    art::MutexLock mu(art::Thread::Current(), uninitialized_class_definitions_lock_);
+    auto it = std::find(uninitialized_class_definitions_.begin(),
+                        uninitialized_class_definitions_.end(),
+                        def);
+    if (it != uninitialized_class_definitions_.end()) {
+      uninitialized_class_definitions_.erase(it);
+      return;
+    }
+    while (std::find(initializing_class_definitions_.begin(),
+                     initializing_class_definitions_.end(),
+                     def) != initializing_class_definitions_.end()) {
+      WaitForClassInitializationToFinish();
+    }
+    it = std::find(initialized_class_definitions_.begin(),
+                   initialized_class_definitions_.end(),
+                   def);
+    CHECK(it != initialized_class_definitions_.end()) << "Could not find class definition for "
+                                                      << def->GetName();
+    initialized_class_definitions_.erase(it);
+  }
+
+  void AddArtDefinition(ArtClassDefinition* def) REQUIRES(!uninitialized_class_definitions_lock_) {
+    DCHECK(def->IsLazyDefinition());
+    art::MutexLock mu(art::Thread::Current(), uninitialized_class_definitions_lock_);
+    uninitialized_class_definitions_.push_back(def);
+  }
+
+ private:
+  bool DefinitionIsInitializing(uintptr_t ptr) REQUIRES(uninitialized_class_definitions_lock_) {
+    return std::find_if(initializing_class_definitions_.begin(),
+                        initializing_class_definitions_.end(),
+                        [&](const auto op) { return op->ContainsAddress(ptr); }) !=
+        initializing_class_definitions_.end();
+  }
+
+  void WaitForClassInitializationToFinish() REQUIRES(uninitialized_class_definitions_lock_) {
+    class_definition_initialized_cond_.Wait(art::Thread::Current());
+  }
+
+  art::Mutex uninitialized_class_definitions_lock_ ACQUIRED_BEFORE(art::Locks::abort_lock_);
+  art::ConditionVariable class_definition_initialized_cond_
+      GUARDED_BY(uninitialized_class_definitions_lock_);
+
+  // A list of the class definitions that have a non-readable map.
+  std::vector<ArtClassDefinition*> uninitialized_class_definitions_
+      GUARDED_BY(uninitialized_class_definitions_lock_);
+
+  // A list of class definitions that are currently undergoing unquickening. Threads should wait
+  // until the definition is no longer in this before returning.
+  std::vector<ArtClassDefinition*> initializing_class_definitions_
+      GUARDED_BY(uninitialized_class_definitions_lock_);
+
+  // A list of class definitions that are already unquickened. Threads should immediately return if
+  // it is here.
+  std::vector<ArtClassDefinition*> initialized_class_definitions_
+      GUARDED_BY(uninitialized_class_definitions_lock_);
+};
+
+static TransformationFaultHandler* gTransformFaultHandler = nullptr;
+
+void Transformer::Setup() {
+  // Although we create this the fault handler is actually owned by the 'art::fault_manager' which
+  // will take care of destroying it.
+  if (art::MemMap::kCanReplaceMapping && ArtClassDefinition::kEnableOnDemandDexDequicken) {
+    gTransformFaultHandler = new TransformationFaultHandler(&art::fault_manager);
+  }
+}
+
+// Simple helper to add and remove the class definition from the fault handler.
+class ScopedDefinitionHandler {
+ public:
+  explicit ScopedDefinitionHandler(ArtClassDefinition* def)
+      : def_(def), is_lazy_(def_->IsLazyDefinition()) {
+    if (is_lazy_) {
+      gTransformFaultHandler->AddArtDefinition(def_);
+    }
+  }
+
+  ~ScopedDefinitionHandler() {
+    if (is_lazy_) {
+      gTransformFaultHandler->RemoveDefinition(def_);
+    }
+  }
+
+ private:
+  ArtClassDefinition* def_;
+  bool is_lazy_;
+};
+
 // Initialize templates.
 template
 void Transformer::TransformSingleClassDirect<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
@@ -78,6 +250,7 @@ void Transformer::TransformSingleClassDirect(EventHandler* event_handler,
   static_assert(kEvent == ArtJvmtiEvent::kClassFileLoadHookNonRetransformable ||
                 kEvent == ArtJvmtiEvent::kClassFileLoadHookRetransformable,
                 "bad event type");
+  ScopedDefinitionHandler handler(def);
   jint new_len = -1;
   unsigned char* new_data = nullptr;
   art::ArrayRef<const unsigned char> dex_data = def->GetDexData();
diff --git a/openjdkjvmti/transform.h b/openjdkjvmti/transform.h
index f43af174f0..8bbeda4b09 100644
--- a/openjdkjvmti/transform.h
+++ b/openjdkjvmti/transform.h
@@ -48,6 +48,8 @@ jvmtiError GetClassLocation(ArtJvmTiEnv* env, jclass klass, /*out*/std::string*
 
 class Transformer {
  public:
+  static void Setup();
+
   template<ArtJvmtiEvent kEvent>
   static void TransformSingleClassDirect(
       EventHandler* event_handler,