Update vdex in place if input_vdex_fd == output_vdex_fd

Test: test-art-host
bug: 30937355
Change-Id: Ib8180d67996faec518d9092725b5de00d4dba9f6

5 files changed, 80 insertions, 36 deletions

diff --git a/compiler/image_test.cc b/compiler/image_test.cc
index 1290379569..377eddcf74 100644
--- a/compiler/image_test.cc
+++ b/compiler/image_test.cc
@@ -244,6 +244,7 @@ void CompilationHelper::Compile(CompilerDriver* driver,
             driver->GetInstructionSetFeatures(),
             &key_value_store,
             /* verify */ false,           // Dex files may be dex-to-dex-ed, don't verify.
+            /* update_input_vdex */ false,
             &cur_opened_dex_files_map,
             &cur_opened_dex_files);
         ASSERT_TRUE(dex_files_ok);
diff --git a/compiler/oat_test.cc b/compiler/oat_test.cc
index edc93ab0c4..86d92ff0b5 100644
--- a/compiler/oat_test.cc
+++ b/compiler/oat_test.cc
@@ -198,6 +198,7 @@ class OatTest : public CommonCompilerTest {
                                          compiler_driver_->GetInstructionSetFeatures(),
                                          &key_value_store,
                                          verify,
+                                         /* update_input_vdex */ false,
                                          &opened_dex_files_map,
                                          &opened_dex_files)) {
       return false;
diff --git a/compiler/oat_writer.cc b/compiler/oat_writer.cc
index bebd5f5ae2..6de08c837e 100644
--- a/compiler/oat_writer.cc
+++ b/compiler/oat_writer.cc
@@ -483,6 +483,7 @@ bool OatWriter::WriteAndOpenDexFiles(
     const InstructionSetFeatures* instruction_set_features,
     SafeMap<std::string, std::string>* key_value_store,
     bool verify,
+    bool update_input_vdex,
     /*out*/ std::unique_ptr<MemMap>* opened_dex_files_map,
     /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files) {
   CHECK(write_state_ == WriteState::kAddingDexFileSources);
@@ -511,15 +512,15 @@ bool OatWriter::WriteAndOpenDexFiles(
   if (kIsVdexEnabled) {
     std::unique_ptr<BufferedOutputStream> vdex_out(
         MakeUnique<BufferedOutputStream>(MakeUnique<FileOutputStream>(vdex_file)));
-
     // Write DEX files into VDEX, mmap and open them.
-    if (!WriteDexFiles(vdex_out.get(), vdex_file) ||
+    if (!WriteDexFiles(vdex_out.get(), vdex_file, update_input_vdex) ||
         !OpenDexFiles(vdex_file, verify, &dex_files_map, &dex_files)) {
       return false;
     }
   } else {
+    DCHECK(!update_input_vdex);
     // Write DEX files into OAT, mmap and open them.
-    if (!WriteDexFiles(oat_rodata, vdex_file) ||
+    if (!WriteDexFiles(oat_rodata, vdex_file, update_input_vdex) ||
         !OpenDexFiles(vdex_file, verify, &dex_files_map, &dex_files)) {
       return false;
     }
@@ -2096,47 +2097,56 @@ bool OatWriter::ValidateDexFileHeader(const uint8_t* raw_header, const char* loc
   return true;
 }
 
-bool OatWriter::WriteDexFiles(OutputStream* out, File* file) {
+bool OatWriter::WriteDexFiles(OutputStream* out, File* file, bool update_input_vdex) {
   TimingLogger::ScopedTiming split("Write Dex files", timings_);
 
   vdex_dex_files_offset_ = vdex_size_;
 
   // Write dex files.
   for (OatDexFile& oat_dex_file : oat_dex_files_) {
-    if (!WriteDexFile(out, file, &oat_dex_file)) {
+    if (!WriteDexFile(out, file, &oat_dex_file, update_input_vdex)) {
       return false;
     }
   }
 
-  // Close sources.
+  CloseSources();
+  return true;
+}
+
+void OatWriter::CloseSources() {
   for (OatDexFile& oat_dex_file : oat_dex_files_) {
     oat_dex_file.source_.Clear();  // Get rid of the reference, it's about to be invalidated.
   }
   zipped_dex_files_.clear();
   zip_archives_.clear();
   raw_dex_files_.clear();
-  return true;
 }
 
-bool OatWriter::WriteDexFile(OutputStream* out, File* file, OatDexFile* oat_dex_file) {
+bool OatWriter::WriteDexFile(OutputStream* out,
+                             File* file,
+                             OatDexFile* oat_dex_file,
+                             bool update_input_vdex) {
   if (!SeekToDexFile(out, file, oat_dex_file)) {
     return false;
   }
   if (profile_compilation_info_ != nullptr) {
+    DCHECK(!update_input_vdex);
     if (!LayoutAndWriteDexFile(out, oat_dex_file)) {
       return false;
     }
   } else if (oat_dex_file->source_.IsZipEntry()) {
+    DCHECK(!update_input_vdex);
     if (!WriteDexFile(out, file, oat_dex_file, oat_dex_file->source_.GetZipEntry())) {
       return false;
     }
   } else if (oat_dex_file->source_.IsRawFile()) {
+    DCHECK(!update_input_vdex);
     if (!WriteDexFile(out, file, oat_dex_file, oat_dex_file->source_.GetRawFile())) {
       return false;
     }
   } else {
     DCHECK(oat_dex_file->source_.IsRawData());
-    if (!WriteDexFile(out, oat_dex_file, oat_dex_file->source_.GetRawData())) {
+    if (!WriteDexFile(out, oat_dex_file, oat_dex_file->source_.GetRawData(), update_input_vdex)) {
       return false;
     }
   }
@@ -2146,6 +2156,7 @@ bool OatWriter::WriteDexFile(OutputStream* out, File* file, OatDexFile* oat_dex_
     DCHECK_EQ(vdex_size_, oat_dex_file->dex_file_offset_);
     vdex_size_ += oat_dex_file->dex_file_size_;
   } else {
+    DCHECK(!update_input_vdex);
     DCHECK_EQ(oat_size_, oat_dex_file->dex_file_offset_);
     oat_size_ += oat_dex_file->dex_file_size_;
   }
@@ -2216,7 +2227,7 @@ bool OatWriter::LayoutAndWriteDexFile(OutputStream* out, OatDexFile* oat_dex_fil
   DexLayout dex_layout(options, profile_compilation_info_, nullptr);
   dex_layout.ProcessDexFile(location.c_str(), dex_file.get(), 0);
   std::unique_ptr<MemMap> mem_map(dex_layout.GetAndReleaseMemMap());
-  if (!WriteDexFile(out, oat_dex_file, mem_map->Begin())) {
+  if (!WriteDexFile(out, oat_dex_file, mem_map->Begin(), /* update_input_vdex */ false)) {
     return false;
   }
   // Set the checksum of the new oat dex file to be the original file's checksum.
@@ -2373,22 +2384,27 @@ bool OatWriter::WriteDexFile(OutputStream* out,
 
 bool OatWriter::WriteDexFile(OutputStream* out,
                              OatDexFile* oat_dex_file,
-                             const uint8_t* dex_file) {
+                             const uint8_t* dex_file,
+                             bool update_input_vdex) {
   // Note: The raw data has already been checked to contain the header
   // and all the data that the header specifies as the file size.
   DCHECK(dex_file != nullptr);
   DCHECK(ValidateDexFileHeader(dex_file, oat_dex_file->GetLocation()));
   const UnalignedDexFileHeader* header = AsUnalignedDexFileHeader(dex_file);
 
-  if (!out->WriteFully(dex_file, header->file_size_)) {
-    PLOG(ERROR) << "Failed to write dex file " << oat_dex_file->GetLocation()
-                << " to " << out->GetLocation();
-    return false;
-  }
-  if (!out->Flush()) {
-    PLOG(ERROR) << "Failed to flush stream after writing dex file."
-                << " File: " << oat_dex_file->GetLocation();
-    return false;
+  if (update_input_vdex) {
+    // The vdex already contains the dex code, no need to write it again.
+  } else {
+    if (!out->WriteFully(dex_file, header->file_size_)) {
+      PLOG(ERROR) << "Failed to write dex file " << oat_dex_file->GetLocation()
+                  << " to " << out->GetLocation();
+      return false;
+    }
+    if (!out->Flush()) {
+      PLOG(ERROR) << "Failed to flush stream after writing dex file."
+                  << " File: " << oat_dex_file->GetLocation();
+      return false;
+    }
   }
 
   // Update dex file size and resize class offsets in the OatDexFile.
diff --git a/compiler/oat_writer.h b/compiler/oat_writer.h
index da221d6029..8d087f4f91 100644
--- a/compiler/oat_writer.h
+++ b/compiler/oat_writer.h
@@ -158,12 +158,15 @@ class OatWriter {
   // Supporting data structures are written into the .rodata section of the oat file.
   // The `verify` setting dictates whether the dex file verifier should check the dex files.
   // This is generally the case, and should only be false for tests.
+  // If `update_input_vdex` is true, then this method won't actually write the dex files,
+  // and the compiler will just re-use the existing vdex file.
   bool WriteAndOpenDexFiles(File* vdex_file,
                             OutputStream* oat_rodata,
                             InstructionSet instruction_set,
                             const InstructionSetFeatures* instruction_set_features,
                             SafeMap<std::string, std::string>* key_value_store,
                             bool verify,
+                            bool update_input_vdex,
                             /*out*/ std::unique_ptr<MemMap>* opened_dex_files_map,
                             /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files);
   bool WriteQuickeningInfo(OutputStream* vdex_out);
@@ -263,8 +266,13 @@ class OatWriter {
   // with a given DexMethodVisitor.
   bool VisitDexMethods(DexMethodVisitor* visitor);
 
-  bool WriteDexFiles(OutputStream* out, File* file);
-  bool WriteDexFile(OutputStream* out, File* file, OatDexFile* oat_dex_file);
+  // If `update_input_vdex` is true, then this method won't actually write the dex files,
+  // and the compiler will just re-use the existing vdex file.
+  bool WriteDexFiles(OutputStream* out, File* file, bool update_input_vdex);
+  bool WriteDexFile(OutputStream* out,
+                    File* file,
+                    OatDexFile* oat_dex_file,
+                    bool update_input_vdex);
   bool SeekToDexFile(OutputStream* out, File* file, OatDexFile* oat_dex_file);
   bool LayoutAndWriteDexFile(OutputStream* out, OatDexFile* oat_dex_file);
   bool WriteDexFile(OutputStream* out,
@@ -275,7 +283,10 @@ class OatWriter {
                     File* file,
                     OatDexFile* oat_dex_file,
                     File* dex_file);
-  bool WriteDexFile(OutputStream* out, OatDexFile* oat_dex_file, const uint8_t* dex_file);
+  bool WriteDexFile(OutputStream* out,
+                    OatDexFile* oat_dex_file,
+                    const uint8_t* dex_file,
+                    bool update_input_vdex);
   bool OpenDexFiles(File* file,
                     bool verify,
                     /*out*/ std::unique_ptr<MemMap>* opened_dex_files_map,
@@ -306,6 +317,7 @@ class OatWriter {
                              const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files);
   bool WriteCodeAlignment(OutputStream* out, uint32_t aligned_code_delta);
   void SetMultiOatRelativePatcherAdjustment();
+  void CloseSources();
 
   enum class WriteState {
     kAddingDexFileSources,
diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 2346635198..ece81e3919 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1295,17 +1295,23 @@ class Dex2Oat FINAL {
 
         DCHECK_EQ(output_vdex_fd_, -1);
         std::string vdex_filename = ReplaceFileExtension(oat_filename, "vdex");
-        std::unique_ptr<File> vdex_file(OS::CreateEmptyFile(vdex_filename.c_str()));
-        if (vdex_file.get() == nullptr) {
-          PLOG(ERROR) << "Failed to open vdex file: " << vdex_filename;
-          return false;
-        }
-        if (fchmod(vdex_file->Fd(), 0644) != 0) {
-          PLOG(ERROR) << "Failed to make vdex file world readable: " << vdex_filename;
-          vdex_file->Erase();
-          return false;
+        if (vdex_filename == input_vdex_) {
+          update_input_vdex_ = true;
+          std::unique_ptr<File> vdex_file(OS::OpenFileReadWrite(vdex_filename.c_str()));
+          vdex_files_.push_back(std::move(vdex_file));
+        } else {
+          std::unique_ptr<File> vdex_file(OS::CreateEmptyFile(vdex_filename.c_str()));
+          if (vdex_file.get() == nullptr) {
+            PLOG(ERROR) << "Failed to open vdex file: " << vdex_filename;
+            return false;
+          }
+          if (fchmod(vdex_file->Fd(), 0644) != 0) {
+            PLOG(ERROR) << "Failed to make vdex file world readable: " << vdex_filename;
+            vdex_file->Erase();
+            return false;
+          }
+          vdex_files_.push_back(std::move(vdex_file));
         }
-        vdex_files_.push_back(std::move(vdex_file));
       }
     } else {
       std::unique_ptr<File> oat_file(new File(oat_fd_, oat_location_, /* check_usage */ true));
@@ -1319,7 +1325,6 @@ class Dex2Oat FINAL {
       }
       oat_files_.push_back(std::move(oat_file));
 
-      DCHECK_NE(input_vdex_fd_, output_vdex_fd_);
       if (input_vdex_fd_ != -1) {
         struct stat s;
         int rc = TEMP_FAILURE_RETRY(fstat(input_vdex_fd_, &s));
@@ -1352,8 +1357,13 @@ class Dex2Oat FINAL {
         return false;
       }
       vdex_file->DisableAutoClose();
-      if (vdex_file->SetLength(0) != 0) {
-        PLOG(WARNING) << "Truncating vdex file " << vdex_location << " failed.";
+      if (input_vdex_file_ != nullptr && output_vdex_fd_ == input_vdex_fd_) {
+        update_input_vdex_ = true;
+      } else {
+        if (vdex_file->SetLength(0) != 0) {
+          PLOG(ERROR) << "Truncating vdex file " << vdex_location << " failed.";
+          return false;
+        }
       }
       vdex_files_.push_back(std::move(vdex_file));
 
@@ -1542,6 +1552,7 @@ class Dex2Oat FINAL {
             instruction_set_features_.get(),
             key_value_store_.get(),
             verify,
+            update_input_vdex_,
             &opened_dex_files_map,
             &opened_dex_files)) {
           return false;
@@ -2732,6 +2743,9 @@ class Dex2Oat FINAL {
   // See CompilerOptions.force_determinism_.
   bool force_determinism_;
 
+  // Whether the given input vdex is also the output.
+  bool update_input_vdex_ = false;
+
   DISALLOW_IMPLICIT_CONSTRUCTORS(Dex2Oat);
 };