Add Poll interface to CloudFileSystem.

rockset · Dec 16, 2023 · 0ccbb7e · 0ccbb7e
1 parent 4d8ec6a
commit 0ccbb7e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/include/rocksdb/cloud/cloud_file_system.h b/include/rocksdb/cloud/cloud_file_system.h
@@ -508,6 +508,13 @@ class CloudFileSystem : public FileSystem {
                   const std::shared_ptr<Logger>& logger);
 
  public:
+  // Helper class that allows us to provide a customizable Poll implementation
+  // for AsyncGet and AsyncMultiGet calls.
+  struct AsyncIoHandleWrapper {
+    IOStatus(*poll_fwd)(void*);
+    void *io_handle;
+  };
+
   mutable std::shared_ptr<Logger> info_log_;  // informational messages
 
   virtual ~CloudFileSystem();

diff --git a/include/rocksdb/cloud/cloud_file_system_impl.h b/include/rocksdb/cloud/cloud_file_system_impl.h
@@ -236,6 +236,21 @@ class CloudFileSystemImpl : public CloudFileSystem {
         "CloudFileSystemImpl::IsDirectory() not supported.");
   }
 
+
+  IOStatus Poll(std::vector<void*>& io_handles,
+                size_t /*min_completions*/) override {
+    for (auto* handle : io_handles) {
+      auto *wrapper = static_cast<CloudFileSystem::AsyncIoHandleWrapper*>(handle);
+      auto status =  wrapper->poll_fwd(wrapper->io_handle);
+      if (!status.ok()) {
+        // We do not currently support continuation on partial failures.
+        return status;
+      }
+    }
+
+    return IOStatus::OK();
+  }
+
   CloudManifest* GetCloudManifest() { return cloud_manifest_.get(); }
 
   IOStatus DeleteCloudFileFromDest(const std::string& fname) override;