Provide support for FSBuffer for point lookups (#12266)

Summary:
Provide support for FSBuffer for point lookups

It also add support for compaction and scan reads that goes through BlockFetcher when readahead/prefetching is not enabled.

Some of the compaction/Scan reads goes through FilePrefetchBuffer and some through BlockFetcher. This PR add support to use underlying file system scratch buffer for reads that go through BlockFetcher as for FilePrefetch reads, design is complicated to support this feature.

Design - In order to use underlying FileSystem provided scratch for Reads, it uses MultiRead with 1 request instead of Read API which required API change.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/12266

Test Plan: Stress test using underlying file system  scratch buffer internally.

Reviewed By: anand1976

Differential Revision: D53019089

Pulled By: akankshamahajan15

fbshipit-source-id: 4fe3d090d77363320e4b67186fd4d51c005c0961

file/random_access_file_reader.cc

@@ -420,7 +420,8 @@ IOStatus RandomAccessFileReader::MultiRead(const IOOptions& opts,
           remaining_bytes -= request_bytes;
         }
       }
-      io_s = file_->MultiRead(fs_reqs, num_fs_reqs, opts, nullptr);
+      io_s = file_->MultiRead(fs_reqs, num_fs_reqs, opts,
+                              /*IODebugContext*=*/nullptr);
       RecordInHistogram(stats_, MULTIGET_IO_BATCH_SIZE, num_fs_reqs);
     }
 

table/block_based/block_based_table_reader_sync_and_async.h

@@ -192,10 +192,7 @@ DEFINE_SYNC_AND_ASYNC(void, BlockBasedTable::RetrieveMultipleBlocks)
 
     BlockContents serialized_block;
     if (s.ok()) {
-      if (use_fs_scratch) {
-        serialized_block =
-            BlockContents(Slice(req.result.data() + req_offset, handle.size()));
-      } else if (!use_shared_buffer) {
+      if (!use_fs_scratch && !use_shared_buffer) {
         // We allocated a buffer for this block. Give ownership of it to
         // BlockContents so it can free the memory
         assert(req.result.data() == req.scratch);
@@ -206,7 +203,8 @@ DEFINE_SYNC_AND_ASYNC(void, BlockBasedTable::RetrieveMultipleBlocks)
       } else {
         // We used the scratch buffer or direct io buffer
         // which are shared by the blocks.
-        // serialized_block does not have the ownership.
+        // In case of use_fs_scratch, underlying file system provided buffer is
+        // used. serialized_block does not have the ownership.
         serialized_block =
             BlockContents(Slice(req.result.data() + req_offset, handle.size()));
       }

table/block_fetcher.cc

@@ -191,16 +191,21 @@ inline void BlockFetcher::CopyBufferToCompressedBuf() {
 #endif
 }
 
-// Entering this method means the block is not compressed or do not need to be
-// uncompressed. The block can be in one of the following buffers:
+// Before - Entering this method means the block is uncompressed or do not need
+// to be uncompressed.
+//
+// The block can be in one of the following buffers:
 // 1. prefetch buffer if prefetch is enabled and the block is prefetched before
 // 2. stack_buf_ if block size is smaller than the stack_buf_ size and block
 //    is not compressed
 // 3. heap_buf_ if the block is not compressed
 // 4. compressed_buf_ if the block is compressed
-// 5. direct_io_buf_ if direct IO is enabled
-// After this method, if the block is compressed, it should be in
-// compressed_buf_, otherwise should be in heap_buf_.
+// 5. direct_io_buf_ if direct IO is enabled or
+// 6. underlying file_system scratch is used (FSReadRequest.fs_scratch).
+//
+// After - After this method, if the block is compressed, it should be in
+// compressed_buf_ and heap_buf_ points to compressed_buf_, otherwise should be
+// in heap_buf_.
 inline void BlockFetcher::GetBlockContents() {
   if (slice_.data() != used_buf_) {
     // the slice content is not the buffer provided
@@ -217,7 +222,7 @@ inline void BlockFetcher::GetBlockContents() {
       } else {
         heap_buf_ = std::move(compressed_buf_);
       }
-    } else if (direct_io_buf_.get() != nullptr) {
+    } else if (direct_io_buf_.get() != nullptr || use_fs_scratch_) {
       if (compression_type_ == kNoCompression) {
         CopyBufferToHeapBuf();
       } else {
@@ -233,6 +238,8 @@ inline void BlockFetcher::GetBlockContents() {
 }
 
 IOStatus BlockFetcher::ReadBlockContents() {
+  FSReadRequest read_req;
+  read_req.status.PermitUncheckedError();
   if (TryGetUncompressBlockFromPersistentCache()) {
     compression_type_ = kNoCompression;
 #ifndef NDEBUG
@@ -256,18 +263,35 @@ IOStatus BlockFetcher::ReadBlockContents() {
             ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);
         io_status_ =
             file_->Read(opts, handle_.offset(), block_size_with_trailer_,
-                        &slice_, nullptr, &direct_io_buf_);
+                        &slice_, /*scratch=*/nullptr, &direct_io_buf_);
         PERF_COUNTER_ADD(block_read_count, 1);
         used_buf_ = const_cast<char*>(slice_.data());
-      } else {
-        PrepareBufferForBlockFromFile();
+      } else if (use_fs_scratch_) {
         PERF_TIMER_GUARD(block_read_time);
         PERF_CPU_TIMER_GUARD(
             block_read_cpu_time,
             ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);
-        io_status_ =
-            file_->Read(opts, handle_.offset(), block_size_with_trailer_,
-                        &slice_, used_buf_, nullptr);
+        read_req.offset = handle_.offset();
+        read_req.len = block_size_with_trailer_;
+        read_req.scratch = nullptr;
+        io_status_ = file_->MultiRead(opts, &read_req, /*num_reqs=*/1,
+                                      /*AlignedBuf* =*/nullptr);
+        PERF_COUNTER_ADD(block_read_count, 1);
+
+        slice_ = Slice(read_req.result.data(), read_req.result.size());
+        used_buf_ = const_cast<char*>(slice_.data());
+      } else {
+        // It allocates/assign used_buf_
+        PrepareBufferForBlockFromFile();
+
+        PERF_TIMER_GUARD(block_read_time);
+        PERF_CPU_TIMER_GUARD(
+            block_read_cpu_time,
+            ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);
+
+        io_status_ = file_->Read(
+            opts, handle_.offset(), /*size*/ block_size_with_trailer_,
+            /*result*/ &slice_, /*scratch*/ used_buf_, /*aligned_buf=*/nullptr);
         PERF_COUNTER_ADD(block_read_count, 1);
 #ifndef NDEBUG
         if (slice_.data() == &stack_buf_[0]) {
@@ -303,10 +327,17 @@ IOStatus BlockFetcher::ReadBlockContents() {
 
     PERF_COUNTER_ADD(block_read_byte, block_size_with_trailer_);
     if (!io_status_.ok()) {
+      ReleaseFileSystemProvidedBuffer(&read_req);
       return io_status_;
     }
 
+    if (use_fs_scratch_ && !read_req.status.ok()) {
+      ReleaseFileSystemProvidedBuffer(&read_req);
+      return read_req.status;
+    }
+
     if (slice_.size() != block_size_with_trailer_) {
+      ReleaseFileSystemProvidedBuffer(&read_req);
       return IOStatus::Corruption(
           "truncated block read from " + file_->file_name() + " offset " +
           std::to_string(handle_.offset()) + ", expected " +
@@ -318,6 +349,7 @@ IOStatus BlockFetcher::ReadBlockContents() {
     if (io_status_.ok()) {
       InsertCompressedBlockToPersistentCacheIfNeeded();
     } else {
+      ReleaseFileSystemProvidedBuffer(&read_req);
       return io_status_;
     }
   }
@@ -341,6 +373,7 @@ IOStatus BlockFetcher::ReadBlockContents() {
   }
 
   InsertUncompressedBlockToPersistentCacheIfNeeded();
+  ReleaseFileSystemProvidedBuffer(&read_req);
 
   return io_status_;
 }

table/block_fetcher.h

@@ -8,6 +8,7 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 
 #pragma once
+#include "file/file_util.h"
 #include "memory/memory_allocator_impl.h"
 #include "table/block_based/block.h"
 #include "table/block_based/block_type.h"
@@ -68,6 +69,9 @@ class BlockFetcher {
         memory_allocator_compressed_(memory_allocator_compressed),
         for_compaction_(for_compaction) {
     io_status_.PermitUncheckedError();  // TODO(AR) can we improve on this?
+    if (CheckFSFeatureSupport(ioptions_.fs.get(), FSSupportedOps::kFSBuffer)) {
+      use_fs_scratch_ = true;
+    }
   }
 
   IOStatus ReadBlockContents();
@@ -127,6 +131,7 @@ class BlockFetcher {
   bool got_from_prefetch_buffer_ = false;
   CompressionType compression_type_;
   bool for_compaction_ = false;
+  bool use_fs_scratch_ = false;
 
   // return true if found
   bool TryGetUncompressBlockFromPersistentCache();
@@ -142,5 +147,15 @@ class BlockFetcher {
   void InsertCompressedBlockToPersistentCacheIfNeeded();
   void InsertUncompressedBlockToPersistentCacheIfNeeded();
   void ProcessTrailerIfPresent();
+
+  void ReleaseFileSystemProvidedBuffer(FSReadRequest* read_req) {
+    if (use_fs_scratch_) {
+      // Free the scratch buffer allocated by FileSystem.
+      if (read_req->fs_scratch != nullptr) {
+        read_req->fs_scratch.reset();
+        read_req->fs_scratch = nullptr;
+      }
+    }
+  }
 };
 }  // namespace ROCKSDB_NAMESPACE

unreleased_history/new_features/fsbuffer.md

@@ -0,0 +1 @@
+Provide support for FSBuffer for point lookups. Also added support for scans and compactions that don't go through prefetching.

utilities/fault_injection_fs.cc

@@ -457,8 +457,8 @@ IOStatus TestFSRandomAccessFile::MultiRead(FSReadRequest* reqs, size_t num_reqs,
     }
     bool this_injected_error;
     reqs[i].status = fs_->InjectThreadSpecificReadError(
-        FaultInjectionTestFS::ErrorOperation::kMultiReadSingleReq,
-        &(reqs[i].result), use_direct_io(), reqs[i].scratch,
+        FaultInjectionTestFS::ErrorOperation::kRead, &(reqs[i].result),
+        use_direct_io(), reqs[i].scratch,
         /*need_count_increase=*/true,
         /*fault_injected=*/&this_injected_error);
     injected_error |= this_injected_error;