Fix GetUniqueId for cloud files

Summary:
We found an issue when using cloud file with block cache. Block cache uses unique id as prefix when sharing the block cache, and cloud files right now don't generate unique id properly. This means if you `keep_local_sst_files=false` and share a block cache, reads from 1 DB could get results from another DB.

This diff fixes it.

Test Plan:
* Reproducible unit test
* Run in Dhruba's namespace with keep_local_sst_files = false

Reviewers: #platform, dhruba, igor

Reviewed By: #platform, dhruba

Differential Revision: https://rockset.phacility.com/D6613

Author: Duc Hieu Pham

cloud/aws/aws_env.cc

@@ -469,9 +469,9 @@ Status AwsEnv::GetPathForDbid(const std::string& bucket,
   Log(InfoLogLevel::DEBUG_LEVEL, info_log_,
       "[s3] Bucket %s GetPathForDbid dbid %s", bucket.c_str(), dbid.c_str());
 
-  std::unordered_map<std::string, std::string> metadata;
+  CloudObjectInformation info;
   Status st = cloud_env_options.storage_provider->GetCloudObjectMetadata(
-      bucket, dbidkey, &metadata);
+      bucket, dbidkey, &info);
   if (!st.ok()) {
     if (st.IsNotFound()) {
       Log(InfoLogLevel::ERROR_LEVEL, info_log_,
@@ -487,8 +487,8 @@ Status AwsEnv::GetPathForDbid(const std::string& bucket,
 
   // Find "dirname" metadata that stores the pathname of the db
   const char* kDirnameTag = "dirname";
-  auto it = metadata.find(kDirnameTag);
-  if (it != metadata.end()) {
+  auto it = info.metadata.find(kDirnameTag);
+  if (it != info.metadata.end()) {
     *dirname = it->second;
   } else {
     st = Status::NotFound("GetPathForDbid");

cloud/aws/aws_s3.cc

@@ -258,12 +258,23 @@ class S3ReadableFile : public CloudStorageReadableFileImpl {
   S3ReadableFile(const std::shared_ptr<AwsS3ClientWrapper>& s3client,
                  const std::shared_ptr<Logger>& info_log,
                  const std::string& bucket, const std::string& fname,
-                 uint64_t size)
+                 uint64_t size, std::string content_hash)
       : CloudStorageReadableFileImpl(info_log, bucket, fname, size),
-        s3client_(s3client) {}
+        s3client_(s3client),
+        content_hash_(std::move(content_hash)) {}
 
   virtual const char* Type() const { return "s3"; }
 
+  virtual size_t GetUniqueId(char* id, size_t max_size) const override {
+    if (content_hash_.empty()) {
+      return 0;
+    }
+
+    max_size = std::min(content_hash_.size(), max_size);
+    memcpy(id, content_hash_.c_str(), max_size);
+    return max_size;
+  }
+
   // random access, read data from specified offset in file
   Status DoCloudRead(uint64_t offset, size_t n, char* scratch,
                      uint64_t* bytes_read) const override {
@@ -328,6 +339,7 @@ class S3ReadableFile : public CloudStorageReadableFileImpl {
 
  private:
   std::shared_ptr<AwsS3ClientWrapper> s3client_;
+  std::string content_hash_;
 };  // End class S3ReadableFile
 
 /******************** Writablefile ******************/
@@ -371,9 +383,10 @@ class S3StorageProvider : public CloudStorageProviderImpl {
                                         uint64_t* time) override;
 
   // Get the metadata of the object in cloud storage
-  Status GetCloudObjectMetadata(
-      const std::string& bucket_name, const std::string& object_path,
-      std::unordered_map<std::string, std::string>* metadata) override;
+  Status GetCloudObjectMetadata(const std::string& bucket_name,
+                                const std::string& object_path,
+                                CloudObjectInformation* info) override;
+
   Status PutCloudObjectMetadata(
       const std::string& bucket_name, const std::string& object_path,
       const std::unordered_map<std::string, std::string>& metadata) override;
@@ -383,6 +396,7 @@ class S3StorageProvider : public CloudStorageProviderImpl {
                          const std::string& object_path_dest) override;
   Status DoNewCloudReadableFile(
       const std::string& bucket, const std::string& fname, uint64_t fsize,
+      const std::string& content_hash,
       std::unique_ptr<CloudStorageReadableFile>* result,
       const EnvOptions& options) override;
   Status NewCloudWritableFile(const std::string& local_path,
@@ -403,11 +417,13 @@ class S3StorageProvider : public CloudStorageProviderImpl {
                           uint64_t file_size) override;
 
  private:
-  // If metadata, size or modtime is non-nullptr, returns requested data
+  // If metadata, size modtime or etag is non-nullptr, returns requested data
   Status HeadObject(
       const std::string& bucket, const std::string& path,
       std::unordered_map<std::string, std::string>* metadata = nullptr,
-      uint64_t* size = nullptr, uint64_t* modtime = nullptr);
+      uint64_t* size = nullptr, uint64_t* modtime = nullptr,
+      std::string* etag = nullptr);
+
   // The S3 client
   std::shared_ptr<AwsS3ClientWrapper> s3client_;
 };
@@ -646,10 +662,12 @@ Status S3StorageProvider::GetCloudObjectModificationTime(
   return HeadObject(bucket_name, object_path, nullptr, nullptr, time);
 }
 
-Status S3StorageProvider::GetCloudObjectMetadata(
-    const std::string& bucket_name, const std::string& object_path,
-    std::unordered_map<std::string, std::string>* metadata) {
-  return HeadObject(bucket_name, object_path, metadata, nullptr, nullptr);
+Status S3StorageProvider::GetCloudObjectMetadata(const std::string& bucket_name,
+                                                 const std::string& object_path,
+                                                 CloudObjectInformation* info) {
+  assert(info != nullptr);
+  return HeadObject(bucket_name, object_path, &info->metadata, &info->size,
+                    &info->modification_time, &info->content_hash);
 }
 
 Status S3StorageProvider::PutCloudObjectMetadata(
@@ -680,10 +698,11 @@ Status S3StorageProvider::PutCloudObjectMetadata(
 
 Status S3StorageProvider::DoNewCloudReadableFile(
     const std::string& bucket, const std::string& fname, uint64_t fsize,
+    const std::string& content_hash,
     std::unique_ptr<CloudStorageReadableFile>* result,
     const EnvOptions& /*options*/) {
-  result->reset(
-      new S3ReadableFile(s3client_, env_->info_log_, bucket, fname, fsize));
+  result->reset(new S3ReadableFile(s3client_, env_->info_log_, bucket, fname,
+                                   fsize, content_hash));
   return Status::OK();
 }
 
@@ -700,7 +719,7 @@ Status S3StorageProvider::NewCloudWritableFile(
 Status S3StorageProvider::HeadObject(
     const std::string& bucket_name, const std::string& object_path,
     std::unordered_map<std::string, std::string>* metadata, uint64_t* size,
-    uint64_t* modtime) {
+    uint64_t* modtime, std::string* etag) {
   Aws::S3::Model::HeadObjectRequest request;
   request.SetBucket(ToAwsString(bucket_name));
   request.SetKey(ToAwsString(object_path));
@@ -727,6 +746,9 @@ Status S3StorageProvider::HeadObject(
   if (modtime != nullptr) {
     *modtime = res.GetLastModified().Millis();
   }
+  if (etag != nullptr) {
+    *etag = std::string(res.GetETag().data(), res.GetETag().length());
+  }
   return Status::OK();
 }
 

cloud/cloud_env_wrapper.h

@@ -51,9 +51,9 @@ class MockStorageProvider : public CloudStorageProvider {
                                         uint64_t* /*time*/) override {
     return notsup_;
   }
-  Status GetCloudObjectMetadata(
-      const std::string& /*bucket_name*/, const std::string& /*object_path*/,
-      std::unordered_map<std::string, std::string>* /*metadata*/) override {
+  Status GetCloudObjectMetadata(const std::string& /*bucket_name*/,
+                                const std::string& /*object_path*/,
+                                CloudObjectInformation* /* info */) override {
     return notsup_;
   }
   Status CopyCloudObject(const std::string& /*bucket_name_src*/,

cloud/cloud_storage_provider.cc

@@ -99,24 +99,6 @@ Status CloudStorageReadableFileImpl::Skip(uint64_t n) {
   return Status::OK();
 }
 
-size_t CloudStorageReadableFileImpl::GetUniqueId(char* id,
-                                                 size_t max_size) const {
-  // If this is an SST file name, then it can part of the persistent cache.
-  // We need to generate a unique id for the cache.
-  // If it is not a sst file, then nobody should be using this id.
-  uint64_t file_number;
-  FileType file_type;
-  WalFileType log_type;
-  ParseFileName(RemoveEpoch(basename(fname_)), &file_number, &file_type,
-                &log_type);
-  if (max_size >= kMaxVarint64Length && file_number > 0) {
-    char* rid = id;
-    rid = EncodeVarint64(rid, file_number);
-    return static_cast<size_t>(rid - id);
-  }
-  return 0;
-}
-
 /******************** Writablefile ******************/
 
 CloudStorageWritableFileImpl::CloudStorageWritableFileImpl(
@@ -306,14 +288,14 @@ Status CloudStorageProviderImpl::NewCloudReadableFile(
     const std::string& bucket, const std::string& fname,
     std::unique_ptr<CloudStorageReadableFile>* result,
     const EnvOptions& options) {
-  // First, check if the file exists and also find its size. We use size in
-  // CloudReadableFile to make sure we always read the valid ranges of the file
-  uint64_t size;
-  Status st = GetCloudObjectSize(bucket, fname, &size);
+  CloudObjectInformation info;
+  Status st = GetCloudObjectMetadata(bucket, fname, &info);
+
   if (!st.ok()) {
     return st;
   }
-  return DoNewCloudReadableFile(bucket, fname, size, result, options);
+  return DoNewCloudReadableFile(bucket, fname, info.size, info.content_hash,
+                                result, options);
 }
 
 Status CloudStorageProviderImpl::GetCloudObject(

cloud/cloud_storage_provider_impl.h

@@ -21,8 +21,6 @@ class CloudStorageReadableFileImpl : public CloudStorageReadableFile {
 
   virtual Status Skip(uint64_t n) override;
 
-  virtual size_t GetUniqueId(char* id, size_t max_size) const override;
-
  protected:
   virtual Status DoCloudRead(uint64_t offset, size_t n, char* scratch,
                              uint64_t* bytes_read) const = 0;
@@ -124,6 +122,7 @@ class CloudStorageProviderImpl : public CloudStorageProvider {
 
   virtual Status DoNewCloudReadableFile(
       const std::string& bucket, const std::string& fname, uint64_t fsize,
+      const std::string& content_hash,
       std::unique_ptr<CloudStorageReadableFile>* result,
       const EnvOptions& options) = 0;
 

cloud/db_cloud_test.cc

@@ -161,7 +161,8 @@ class CloudTest : public testing::Test {
                  const std::string& dest_bucket_name,
                  const std::string& dest_object_path,
                  std::unique_ptr<DBCloud>* cloud_db,
-                 std::unique_ptr<CloudEnv>* cloud_env) {
+                 std::unique_ptr<CloudEnv>* cloud_env,
+                 bool force_keep_local_on_invalid_dest_bucket = true) {
     // The local directory where the clone resides
     std::string cname = clone_dir_ + "/" + clone_name;
 
@@ -177,7 +178,8 @@ class CloudTest : public testing::Test {
       copt.dest_bucket.SetBucketName(dest_bucket_name);
     }
     copt.dest_bucket.SetObjectPath(dest_object_path);
-    if (!copt.dest_bucket.IsValid()) {
+    if (!copt.dest_bucket.IsValid() &&
+        force_keep_local_on_invalid_dest_bucket) {
       copt.keep_local_sst_files = true;
     }
     // Create new AWS env
@@ -1483,6 +1485,56 @@ TEST_F(CloudTest, PersistentCache) {
   CloseDB();
 }
 
+// This test create 2 DBs that shares a block cache. Ensure that reads from one
+// DB do not get the values from the other DB.
+TEST_F(CloudTest, SharedBlockCache) {
+  cloud_env_options_.keep_local_sst_files = false;
+
+  // Share the block cache.
+  rocksdb::BlockBasedTableOptions bbto;
+  bbto.block_cache = NewLRUCache(10 * 1024 * 1024);
+  bbto.format_version = 4;
+  options_.table_factory.reset(NewBlockBasedTableFactory(bbto));
+
+  OpenDB();
+
+  std::unique_ptr<CloudEnv> clone_env;
+  std::unique_ptr<DBCloud> clone_db;
+  CloneDB("newdb1", cloud_env_options_.src_bucket.GetBucketName(),
+          cloud_env_options_.src_bucket.GetObjectPath() + "-clone", &clone_db,
+          &clone_env, false /* force_keep_local_on_invalid_dest_bucket */);
+
+  // Flush the first DB.
+  db_->Put(WriteOptions(), "db", "original");
+  db_->Flush(FlushOptions());
+
+  // Flush the second DB.
+  clone_db->Put(WriteOptions(), "db", "clone");
+  clone_db->Flush(FlushOptions());
+
+  std::vector<LiveFileMetaData> file_metadatas;
+  db_->GetLiveFilesMetaData(&file_metadatas);
+  ASSERT_EQ(1, file_metadatas.size());
+
+  file_metadatas.clear();
+  clone_db->GetLiveFilesMetaData(&file_metadatas);
+  ASSERT_EQ(1, file_metadatas.size());
+
+  std::string value;
+  clone_db->Get(ReadOptions(), "db", &value);
+  ASSERT_EQ("clone", value);
+
+  db_->Get(ReadOptions(), "db", &value);
+  ASSERT_EQ("original", value);
+
+  // Cleanup
+  clone_db->Close();
+  CloseDB();
+  clone_env->GetCloudEnvOptions().storage_provider->EmptyBucket(
+      cloud_env_options_.src_bucket.GetBucketName(),
+      cloud_env_options_.src_bucket.GetObjectPath() + "-clone");
+}
+
 }  //  namespace rocksdb
 
 // A black-box test for the cloud wrapper around rocksdb

include/rocksdb/cloud/cloud_storage_provider.h

@@ -29,6 +29,17 @@ class CloudStorageWritableFile : public WritableFile {
   virtual const char* Name() const { return "cloud"; }
 };
 
+// Generic information of the object in the cloud. Some information might be
+// vendor-dependent.
+struct CloudObjectInformation {
+  uint64_t size;
+  uint64_t modification_time;
+
+  // Cloud-vendor dependent. In S3, we will provide ETag of the object.
+  std::string content_hash;
+  std::unordered_map<std::string, std::string> metadata;
+};
+
 // A CloudStorageProvider provides the interface to the cloud object
 // store.  Methods can create and empty buckets, as well as other
 // standard bucket object operations get/put/list/delete
@@ -70,9 +81,9 @@ class CloudStorageProvider {
                                                 uint64_t* time) = 0;
 
   // Get the metadata of the object in cloud storage
-  virtual Status GetCloudObjectMetadata(
-      const std::string& bucket_name, const std::string& object_path,
-      std::unordered_map<std::string, std::string>* metadata) = 0;
+  virtual Status GetCloudObjectMetadata(const std::string& bucket_name,
+                                        const std::string& object_path,
+                                        CloudObjectInformation* info) = 0;
 
   // Copy the specified cloud object from one location in the cloud
   // storage to another location in cloud storage