3 files changed, 270 insertions, 1 deletions

diff --git a/openjdkjvmti/ti_class_definition.cc b/openjdkjvmti/ti_class_definition.cc
index 3d9d9e7c59..a480ac0cc6 100644
--- a/openjdkjvmti/ti_class_definition.cc
+++ b/openjdkjvmti/ti_class_definition.cc
@@ -47,6 +47,32 @@
 
 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_.IsValid());
+  CHECK(temp_mmap_.IsValid());
+  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;
+  source.swap(temp_mmap_);
+  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.
@@ -60,6 +86,27 @@ bool ArtClassDefinition::IsModified() const {
     return false;
   }
 
+  // The dex_data_ was never touched by the agents.
+  if (dex_data_mmap_.IsValid() && 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 +241,44 @@ 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_ = art::MemMap::MapAnonymous(mmap_name.c_str(),
+                                               dequick_size,
+                                               PROT_NONE,
+                                               /*low_4gb=*/ false,
+                                               &error);
+    mmap_name += "-TEMP";
+    temp_mmap_ = art::MemMap::MapAnonymous(mmap_name.c_str(),
+                                           dequick_size,
+                                           PROT_READ | PROT_WRITE,
+                                           /*low_4gb=*/ false,
+                                           &error);
+    if (UNLIKELY(dex_data_mmap_.IsValid() && temp_mmap_.IsValid())) {
+      // 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();
   temp_mmap_.Reset();
   // Failed to mmap a large enough area (or on-demand dequickening was disabled). This is
diff --git a/openjdkjvmti/ti_class_definition.h b/openjdkjvmti/ti_class_definition.h
index a914ec2d2f..cb0853bdb5 100644
--- a/openjdkjvmti/ti_class_definition.h
+++ b/openjdkjvmti/ti_class_definition.h
@@ -49,6 +49,9 @@ 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),
@@ -132,6 +135,13 @@ class ArtClassDefinition {
     return name_;
   }
 
+  bool IsLazyDefinition() const {
+    DCHECK(IsInitialized());
+    return dex_data_mmap_.IsValid() &&
+        dex_data_.data() == dex_data_mmap_.Begin() &&
+        dex_data_mmap_.GetProtect() == PROT_NONE;
+  }
+
   jobject GetProtectionDomain() const {
     DCHECK(IsInitialized());
     return protection_domain_;
@@ -156,7 +166,8 @@ 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-compact-dex'd
+  // Mmap that will be filled with the original-dex-file lazily if it needs to be de-quickened or
+  // de-compact-dex'd
   mutable art::MemMap dex_data_mmap_;
   // This is a temporary mmap we will use to be able to fill the dex file data atomically.
   mutable art::MemMap temp_mmap_;
diff --git a/openjdkjvmti/transform.cc b/openjdkjvmti/transform.cc
index 753697dbe6..9c440ef367 100644
--- a/openjdkjvmti/transform.cc
+++ b/openjdkjvmti/transform.cc
@@ -70,13 +70,185 @@
 
 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();
+    uintptr_t ptr = reinterpret_cast<uintptr_t>(siginfo->si_addr);
+    if (UNLIKELY(uninitialized_class_definitions_lock_.IsExclusiveHeld(self))) {
+      // It's possible this is just some other unrelated segv that should be
+      // handled separately, continue to later handlers. This is likely due to
+      // running out of memory somewhere along the FixedUpDexFile pipeline and
+      // is likely unrecoverable. By returning false here though we will get a
+      // better, more accurate, stack-trace later that points to the actual
+      // issue.
+      LOG(WARNING) << "Recursive SEGV occurred during Transformation dequickening at 0x" << std::hex
+                   << ptr;
+      return false;
+    }
+    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); }) !=
+            initialized_class_definitions_.end();
+      }
+    }
+
+    if (LIKELY(self != nullptr)) {
+      CHECK_EQ(self->GetState(), art::ThreadState::kNative)
+          << "Transformation fault handler occurred outside of native mode";
+    }
+
+    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;
 static EventHandler* gEventHandler = nullptr;
 
 
 void Transformer::Register(EventHandler* eh) {
+  // 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);
+  }
   gEventHandler = eh;
 }
 
+// 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>(
@@ -100,6 +272,7 @@ void Transformer::TransformSingleClassDirect(EventHandler* event_handler,
   // native state early. This also avoids any problems that the FaultHandler might have in
   // determining if an access to the dex_data is from generated code or not.
   art::ScopedThreadStateChange stsc(self, art::ThreadState::kNative);
+  ScopedDefinitionHandler handler(def);
   jint new_len = -1;
   unsigned char* new_data = nullptr;
   art::ArrayRef<const unsigned char> dex_data = def->GetDexData();