Support SuperSnapshots in one more variant of MultiGet

This PR contains a couple of improvements to SuperSnapshots
implementation, the major one being support for one more variant of
MultiGet (there are three total, we now support two).

cloud/replication_test.cc

@@ -1315,13 +1315,15 @@ TEST_F(ReplicationTest, SuperSnapshot) {
  ASSERT_OK(follower->Get(ro, "k1", &val));
  EXPECT_EQ(val, "v1");
 
- auto iter = follower->NewIterator(ro, follower->DefaultColumnFamily());
+ auto iter = std::unique_ptr<rocksdb::Iterator>(
+ follower->NewIterator(ro, follower->DefaultColumnFamily()));
  iter->SeekToFirst();
  EXPECT_TRUE(iter->Valid());
  EXPECT_EQ(iter->key(), "k1");
  EXPECT_EQ(iter->value(), "v1");
  iter->Next();
  EXPECT_FALSE(iter->Valid());
+ iter.reset();
 
  ro.snapshot = nullptr;
  ASSERT_OK(follower->Get(ro, followerCF("cf1"), "cf1k1", &val));
@@ -1339,19 +1341,33 @@ TEST_F(ReplicationTest, SuperSnapshot) {
  cfs.push_back(followerCF("cf1"));
  std::vector<std::string> vals;
 
+ // Multiget on cf1, snapshot
+ ro.snapshot = snapshots[1];
  auto statuses = follower->MultiGet(ro, cfs, keys, &vals);
  ASSERT_OK(statuses[0]);
  ASSERT_TRUE(statuses[1].IsNotFound());
  ASSERT_EQ(vals[0], "cf1v1");
 
+ // Multiget on cf1, no snapshot
  ro.snapshot = nullptr;
  statuses = follower->MultiGet(ro, cfs, keys, &vals);
  ASSERT_OK(statuses[0]);
  ASSERT_TRUE(statuses[1].IsNotFound());
  ASSERT_EQ(vals[0], "cf1v2");
 
- // Column family <-> snapshot mismatch
+ // Multiget on default column family, snapshot, pinnable values
  ro.snapshot = snapshots[0];
+ keys[0] = "k1";
+ std::vector<PinnableSlice> pinnableValues;
+ pinnableValues.resize(2);
+ follower->MultiGet(ro, follower->DefaultColumnFamily(), 2,
+ keys.data(), pinnableValues.data(), statuses.data(), true);
+ ASSERT_OK(statuses[0]);
+ ASSERT_TRUE(statuses[1].IsNotFound());
+ ASSERT_EQ(pinnableValues[0], "v1");
+ pinnableValues.clear();
+
+ // Column family <-> snapshot mismatch
  ASSERT_FALSE(follower->Get(ro, followerCF("cf1"), "cf1k1", &val).ok());
 
  follower->ReleaseSnapshot(snapshots[0]);

db/db_impl/db_impl.cc

@@ -3139,8 +3139,9 @@ void DBImpl::MultiGet(const ReadOptions& read_options, const size_t num_keys,
  if (dynamic_cast<const SuperSnapshotImpl*>(read_options.snapshot)) {
  for (size_t i = 0; i < num_keys; ++i) {
  statuses[i] = Status::NotSupported(
- "MultiGet with timestamps does not support super snapshot");
+ "This variant of MultiGet does not yet support super snapshots");
  }
+ return;
  }
  // RocksDB-Cloud contribution end
 
@@ -3400,8 +3401,12 @@ void DBImpl::MultiGetWithCallback(
  multiget_cf_data[0].super_version, consistent_seqnum,
  read_callback);
  assert(s.ok() || s.IsTimedOut() || s.IsAborted());
- ReturnAndCleanupSuperVersion(multiget_cf_data[0].cfd,
- multiget_cf_data[0].super_version);
+ // RocksDB-Cloud contribution begin
+ if (!dynamic_cast<const SuperSnapshotImpl*>(read_options.snapshot)) {
+ ReturnAndCleanupSuperVersion(multiget_cf_data[0].cfd,
+ multiget_cf_data[0].super_version);
+ }
+ // RocksDB-Cloud contribution end
 }
 
 // The actual implementation of batched MultiGet. Parameters -
@@ -3844,10 +3849,12 @@ Iterator* DBImpl::NewIterator(const ReadOptions& read_options,
  result = nullptr;
 
 #else
+ // RocksDB-Cloud contribution begin
  if (dynamic_cast<const SuperSnapshotImpl*>(read_options.snapshot)) {
  return NewErrorIterator(Status::NotSupported(
  "Tailing iterator not supported with super snapshot"));
  }
+ // RocksDB-Cloud contribution end
 
  SuperVersion* sv = cfd->GetReferencedSuperVersion(this);
  auto iter = new ForwardIterator(this, read_options, cfd, sv,
@@ -4077,6 +4084,8 @@ Status DBImpl::GetSuperSnapshots(
  return Status::InvalidArgument(
  "GetSuperSnapshots only supported in RocksDB compiled with USE_RTTI=1");
 #endif
+ InstrumentedMutexLock l(&mutex_);
+
  if (!is_snapshot_supported_) {
  return Status::InvalidArgument("Snapshot not supported");
  }
@@ -4092,7 +4101,8 @@ Status DBImpl::GetSuperSnapshots(
  for (auto& cf : column_families) {
  auto cfh = static_cast_with_check<ColumnFamilyHandleImpl>(cf);
  auto cfd = cfh->cfd();
- auto sv = cfd->GetReferencedSuperVersion(this);
+ auto sv = cfd->GetSuperVersion();
+ sv->Ref();
  auto ss = new SuperSnapshotImpl(cfd, sv);
  snapshots_.New(ss, snapshot_seq, unix_time,
  /*is_write_conflict_boundary=*/false);