From 5d870717ae87c572f1339d56e2dd51a4c83f4205 Mon Sep 17 00:00:00 2001
From: Igor Canadi <icanadi@fb.com>
Date: Thu, 5 Jun 2014 13:19:35 -0700
Subject: [PATCH] Correctly preallocate files in universal compaction

Summary: In universal compaction, MaxFileSizeForLevel is ULLONG_MAX. We've been preallocation files to UULONG_MAX size all these time :)

Test Plan: make check

Reviewers: dhruba, haobo, ljin, sdong, yhchiang

Reviewed By: yhchiang

Subscribers: leveldb

Differential Revision: https://reviews.facebook.net/D18915
---
 db/compaction.cc | 17 +++++++++++++++++
 db/compaction.h  |  5 +++++
 db/db_impl.cc    |  7 ++-----
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/db/compaction.cc b/db/compaction.cc
index a8caa59ef..e88e822b2 100644
--- a/db/compaction.cc
+++ b/db/compaction.cc
@@ -250,4 +250,21 @@ void Compaction::Summary(char* output, int len) {
   snprintf(output + write, len - write, "]");
 }
 
+uint64_t Compaction::OutputFilePreallocationSize() {
+  uint64_t preallocation_size = 0;
+
+  if (cfd_->options()->compaction_style == kCompactionStyleLevel) {
+    preallocation_size =
+        cfd_->compaction_picker()->MaxFileSizeForLevel(output_level());
+  } else {
+    for (const auto& f : inputs_[0]) {
+      preallocation_size += f->file_size;
+    }
+  }
+  // Over-estimate slightly so we don't end up just barely crossing
+  // the threshold
+  return preallocation_size * 1.1;
+}
+
+
 }  // namespace rocksdb
diff --git a/db/compaction.h b/db/compaction.h
index aaa402303..d6f6f80b4 100644
--- a/db/compaction.h
+++ b/db/compaction.h
@@ -91,6 +91,11 @@ class Compaction {
   // Was this compaction triggered manually by the client?
   bool IsManualCompaction() { return is_manual_compaction_; }
 
+  // Returns a number of byte that the output file should be preallocated to
+  // In level compaction, that is max_file_size_. In universal compaction, that
+  // is the sum of all input file sizes
+  uint64_t OutputFilePreallocationSize();
+
  private:
   friend class CompactionPicker;
   friend class UniversalCompactionPicker;
diff --git a/db/db_impl.cc b/db/db_impl.cc
index 5301fa599..7ecc93fd5 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -2334,13 +2334,10 @@ Status DBImpl::OpenCompactionOutputFile(CompactionState* compact) {
   Status s = env_->NewWritableFile(fname, &compact->outfile, storage_options_);
 
   if (s.ok()) {
-    // Over-estimate slightly so we don't end up just barely crossing
-    // the threshold.
-    ColumnFamilyData* cfd = compact->compaction->column_family_data();
     compact->outfile->SetPreallocationBlockSize(
-        1.1 * cfd->compaction_picker()->MaxFileSizeForLevel(
-                  compact->compaction->output_level()));
+        compact->compaction->OutputFilePreallocationSize());
 
+    ColumnFamilyData* cfd = compact->compaction->column_family_data();
     CompressionType compression_type =
         GetCompressionType(*cfd->options(), compact->compaction->output_level(),
                            compact->compaction->enable_compression());