From 1a136c1f137b3cfae7071a3b4dc0cfda7b631dd6 Mon Sep 17 00:00:00 2001
From: Ding Ma <din@devvm24836.prn1.facebook.com>
Date: Fri, 16 Dec 2016 18:26:56 -0800
Subject: [PATCH] Expose file size

Summary:
add a new function to SstFileWriter that will tell the user how big is there file right now.
Closes https://github.com/facebook/rocksdb/pull/1686

Differential Revision: D4338868

Pulled By: mdyuki1016

fbshipit-source-id: c1ee16a
---
 db/external_sst_file_test.cc      | 10 ++++++++++
 include/rocksdb/sst_file_writer.h |  3 +++
 table/sst_file_writer.cc          |  8 +++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/db/external_sst_file_test.cc b/db/external_sst_file_test.cc
index 458f26a0c..f94ade569 100644
--- a/db/external_sst_file_test.cc
+++ b/db/external_sst_file_test.cc
@@ -141,6 +141,9 @@ TEST_F(ExternalSSTFileTest, Basic) {
 
     SstFileWriter sst_file_writer(EnvOptions(), options, options.comparator);
 
+    // Current file size should be 0 after sst_file_writer init and before open a file.
+    ASSERT_EQ(sst_file_writer.FileSize(), 0);
+
     // file1.sst (0 => 99)
     std::string file1 = sst_files_dir_ + "file1.sst";
     ASSERT_OK(sst_file_writer.Open(file1));
@@ -150,6 +153,10 @@ TEST_F(ExternalSSTFileTest, Basic) {
     ExternalSstFileInfo file1_info;
     Status s = sst_file_writer.Finish(&file1_info);
     ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Current file size should be non-zero after success write.
+    ASSERT_GT(sst_file_writer.FileSize(), 0);
+
     ASSERT_EQ(file1_info.file_path, file1);
     ASSERT_EQ(file1_info.num_entries, 100);
     ASSERT_EQ(file1_info.smallest_key, Key(0));
@@ -184,7 +191,10 @@ TEST_F(ExternalSSTFileTest, Basic) {
     }
     ExternalSstFileInfo file3_info;
     s = sst_file_writer.Finish(&file3_info);
+
     ASSERT_TRUE(s.ok()) << s.ToString();
+    // Current file size should be non-zero after success finish.
+    ASSERT_GT(sst_file_writer.FileSize(), 0);
     ASSERT_EQ(file3_info.file_path, file3);
     ASSERT_EQ(file3_info.num_entries, 105);
     ASSERT_EQ(file3_info.smallest_key, Key(195));
diff --git a/include/rocksdb/sst_file_writer.h b/include/rocksdb/sst_file_writer.h
index 8cd1ceea5..d7ca81bc4 100644
--- a/include/rocksdb/sst_file_writer.h
+++ b/include/rocksdb/sst_file_writer.h
@@ -66,6 +66,9 @@ class SstFileWriter {
   // which will be populated with information about the created sst file
   Status Finish(ExternalSstFileInfo* file_info = nullptr);
 
+  // Return the current file size.
+  uint64_t FileSize();
+
  private:
   struct Rep;
   Rep* rep_;
diff --git a/table/sst_file_writer.cc b/table/sst_file_writer.cc
index 74feb1a30..1d4f1c40a 100644
--- a/table/sst_file_writer.cc
+++ b/table/sst_file_writer.cc
@@ -44,7 +44,9 @@ SstFileWriter::SstFileWriter(const EnvOptions& env_options,
                              const Options& options,
                              const Comparator* user_comparator,
                              ColumnFamilyHandle* column_family)
-    : rep_(new Rep(env_options, options, user_comparator, column_family)) {}
+    : rep_(new Rep(env_options, options, user_comparator, column_family)) {
+      rep_->file_info.file_size = 0;
+    }
 
 SstFileWriter::~SstFileWriter() {
   if (rep_->builder) {
@@ -187,4 +189,8 @@ Status SstFileWriter::Finish(ExternalSstFileInfo* file_info) {
   r->builder.reset();
   return s;
 }
+
+uint64_t SstFileWriter::FileSize() {
+  return rep_->file_info.file_size;
+}
 }  // namespace rocksdb