Make Write API work for TTL databases

Summary: Added logic to make another WriteBatch with Timestamps during the Write function execution in TTL class. Also expanded the ttl_test to test for it. Have done nothing for Merge for now.

Test Plan: make ttl_test;./ttl_test

Reviewers: haobo, vamsi, dhruba

Reviewed By: haobo

CC: leveldb

Differential Revision: https://reviews.facebook.net/D10827

utilities/ttl/db_ttl.cc

@@ -226,7 +226,34 @@ Status DBWithTTL::Merge(const WriteOptions& options,
 }
 
 Status DBWithTTL::Write(const WriteOptions& opts, WriteBatch* updates) {
-  return db_->Write(opts, updates);
+  class Handler : public WriteBatch::Handler {
+   public:
+    WriteBatch updates_ttl;
+    Status batch_rewrite_status;
+    virtual void Put(const Slice& key, const Slice& value) {
+      std::string value_with_ts;
+      Status st = AppendTS(value, value_with_ts);
+      if (!st.ok()) {
+        batch_rewrite_status = st;
+      } else {
+        updates_ttl.Put(key, value_with_ts);
+      }
+    }
+    virtual void Merge(const Slice& key, const Slice& value) {
+      // TTL doesn't support merge operation
+      batch_rewrite_status = Status::NotSupported("TTL doesn't support Merge");
+    }
+    virtual void Delete(const Slice& key) {
+      updates_ttl.Delete(key);
+    }
+  };
+  Handler handler;
+  updates->Iterate(&handler);
+  if (!handler.batch_rewrite_status.ok()) {
+    return handler.batch_rewrite_status;
+  } else {
+    return db_->Write(opts, &(handler.updates_ttl));
+  }
 }
 
 Iterator* DBWithTTL::NewIterator(const ReadOptions& opts) {

utilities/ttl/ttl_test.cc

@@ -12,6 +12,10 @@ namespace leveldb {
 
 namespace {
 typedef std::map<std::string, std::string> KVMap;
+enum BatchOperation {
+  PUT = 0,
+  DELETE = 1
+};
 }
 
 class TtlTest {
@@ -69,6 +73,29 @@ class TtlTest {
     ASSERT_EQ((int)kvmap_.size(), num_entries);//check all insertions done
   }
 
+  // Makes a write-batch with key-vals from kvmap_ and 'Write''s it
+  void MakePutWriteBatch(const BatchOperation* batch_ops, int num_ops) {
+    assert(num_ops <= (int)kvmap_.size());
+    static WriteOptions wopts;
+    static FlushOptions flush_opts;
+    WriteBatch batch;
+    kv_it_ = kvmap_.begin();
+    for (int i = 0; i < num_ops && kv_it_ != kvmap_.end(); i++, kv_it_++) {
+      switch (batch_ops[i]) {
+        case PUT:
+          batch.Put(kv_it_->first, kv_it_->second);
+          break;
+        case DELETE:
+          batch.Delete(kv_it_->first);
+          break;
+        default:
+          assert(false);
+      }
+    }
+    db_ttl_->Write(wopts, &batch);
+    db_ttl_->Flush(flush_opts);
+  }
+
   // Puts num_entries starting from start_pos_map from kvmap_ into the database
   void PutValues(int start_pos_map, int num_entries) {
     assert(db_ttl_);
@@ -288,7 +315,27 @@ TEST(TtlTest, ReadOnlyPresentForever) {
   CloseTtl();
 
   OpenReadOnlyTtl(1);
-  SleepCompactCheck(2, 0, kSampleSize, true); // T=2:Set1 should still be there
+  SleepCompactCheck(2, 0, kSampleSize);       // T=2:Set1 should still be there
+  CloseTtl();
+}
+
+// Checks whether WriteBatch works well with TTL
+// Puts all kvs in kvmap_ in a batch and writes first, then deletes first half
+TEST(TtlTest, WriteBatchTest) {
+  MakeKVMap(kSampleSize);
+  BatchOperation batch_ops[kSampleSize];
+  for (int i = 0; i < kSampleSize; i++) {
+    batch_ops[i] = PUT;
+  }
+
+  OpenTtl(2);
+  MakePutWriteBatch(batch_ops, kSampleSize);
+  for (int i = 0; i < kSampleSize / 2; i++) {
+    batch_ops[i] = DELETE;
+  }
+  MakePutWriteBatch(batch_ops, kSampleSize / 2);
+  SleepCompactCheck(0, 0, kSampleSize / 2, false);
+  SleepCompactCheck(0, kSampleSize / 2, kSampleSize - kSampleSize / 2);
   CloseTtl();
 }