//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.

//  This source code is licensed under both the GPLv2 (found in the

//  COPYING file in the root directory) and Apache 2.0 License

//  (found in the LICENSE.Apache file in the root directory).

//

// Copyright (c) 2011 The LevelDB Authors. All rights reserved.

// Use of this source code is governed by a BSD-style license that can be

// found in the LICENSE file. See the AUTHORS file for names of contributors.

#pragma once

#include <errno.h>

#if defined(ROCKSDB_IOURING_PRESENT)

#include <liburing.h>

#include <pthread.h>

#include <sys/uio.h>

#include <cstdio>

#include "util/string_util.h"

// Compatibility defines for io_uring flags that may not be present in older

// kernel headers. These values are fixed and won't change, so it's safe to

// define them even if the running kernel doesn't support them.

#ifndef IORING_SETUP_SINGLE_ISSUER

#define IORING_SETUP_SINGLE_ISSUER (1U << 12)

#endif

#ifndef IORING_SETUP_DEFER_TASKRUN

#define IORING_SETUP_DEFER_TASKRUN (1U << 13)

#endif

#endif

#include <unistd.h>

#include <atomic>

#include <functional>

#include <map>

#include <string>

#include "port/lang.h"

#include "port/port.h"

#include "rocksdb/env.h"

#include "rocksdb/file_system.h"

#include "rocksdb/io_status.h"

#include "test_util/sync_point.h"

#include "util/mutexlock.h"

#include "util/thread_local.h"

// For non linux platform, the following macros are used only as place

// holder.

#if !(defined OS_LINUX) && !(defined OS_FREEBSD) && !(defined CYGWIN) && \

    !(defined OS_AIX) && !(defined OS_ANDROID)

#define POSIX_FADV_NORMAL 0     /* [MC1] no further special treatment */

#define POSIX_FADV_RANDOM 1     /* [MC1] expect random page refs */

#define POSIX_FADV_SEQUENTIAL 2 /* [MC1] expect sequential page refs */

#define POSIX_FADV_WILLNEED 3   /* [MC1] will need these pages */

#define POSIX_FADV_DONTNEED 4   /* [MC1] don't need these pages */

#define POSIX_MADV_NORMAL 0     /* [MC1] no further special treatment */

#define POSIX_MADV_RANDOM 1     /* [MC1] expect random page refs */

#define POSIX_MADV_SEQUENTIAL 2 /* [MC1] expect sequential page refs */

#define POSIX_MADV_WILLNEED 3   /* [MC1] will need these pages */

#define POSIX_MADV_DONTNEED 4   /* [MC1] don't need these pages */

#endif

namespace ROCKSDB_NAMESPACE {

std::string IOErrorMsg(const std::string& context,

                       const std::string& file_name);

// file_name can be left empty if it is not unkown.

IOStatus IOError(const std::string& context, const std::string& file_name,

                 int err_number);

// SyncPoint payload used by deterministic TSAN regression tests to observe

// which virtual address range a freshly created mapping occupies.

struct TsanMappedMemoryInfo {

  const void* addr;

  size_t size;

};

#ifdef __SANITIZE_THREAD__

extern "C" void AnnotateNewMemory(const char* file, int line,

                                  const volatile void* mem, long size);

#endif  // __SANITIZE_THREAD__

inline void TsanAnnotateMappedMemory(const volatile void* mem, size_t size) {

  TsanMappedMemoryInfo info{const_cast<const void*>(mem), size};

  TEST_SYNC_POINT_CALLBACK("TsanAnnotateMappedMemory", &info);

  (void)info;

#ifdef __SANITIZE_THREAD__

  if (mem != nullptr && size != 0) {

    // TSAN does not understand that a new mmap or io_uring setup can legally

    // reuse a virtual address from an unrelated, previously unmapped region.

    // Reset the shadow state as soon as the new mapping exists so later

    // accesses are not reported against stale accesses from the old mapping.

    AnnotateNewMemory(__FILE__, __LINE__, mem, static_cast<long>(size));

  }

#endif  // __SANITIZE_THREAD__

}

class PosixHelper {

 public:

  static const std::string& GetLogicalBlockSizeFileName() {

    static const std::string kLogicalBlockSizeFileName = "logical_block_size";

    return kLogicalBlockSizeFileName;

  }

  static const std::string& GetMaxSectorsKBFileName() {

    static const std::string kMaxSectorsKBFileName = "max_sectors_kb";

    return kMaxSectorsKBFileName;

  }

  static size_t GetUniqueIdFromFile(int fd, char* id, size_t max_size);

  static size_t GetLogicalBlockSizeOfFd(int fd);

  static Status GetLogicalBlockSizeOfDirectory(const std::string& directory,

                                               size_t* size);

  static Status GetMaxSectorsKBOfDirectory(const std::string& directory,

                                           size_t* kb);

 private:

  static const size_t kDefaultMaxSectorsKB = 2 * 1024;

  static size_t GetMaxSectorsKBOfFd(int fd);

  // Return the value in the specified `file_name` under

  // `/sys/block/xxx/queue/` for the device where the file of `fd` is on.

  // If not found, then return the specified `default_return_value`

  static size_t GetQueueSysfsFileValueOfFd(int fd, const std::string& file_name,

                                           size_t default_return_value);

  /// Return the value in the specified `file_name` under

  // `/sys/block/xxx/queue/` for the device where `directory` is on.

  // If not found, then return the specified `default_return_value`

  static Status GetQueueSysfsFileValueofDirectory(const std::string& directory,

                                                  const std::string& file_name,

                                                  size_t* value);

};

/*

 * DirectIOHelper

 */

inline bool IsSectorAligned(const size_t off, size_t sector_size) {

  assert((sector_size & (sector_size - 1)) == 0);

  return (off & (sector_size - 1)) == 0;

}

#ifndef NDEBUG

inline bool IsSectorAligned(const void* ptr, size_t sector_size) {

  return uintptr_t(ptr) % sector_size == 0;

}

#endif

#if defined(ROCKSDB_IOURING_PRESENT)

struct Posix_IOHandle {

  Posix_IOHandle(struct io_uring* _iu,

                 std::function<void(FSReadRequest&, void*)> _cb, void* _cb_arg,

                 uint64_t _offset, size_t _len, char* _scratch,

                 bool _use_direct_io, size_t _alignment)

      : iu(_iu),

        cb(_cb),

        cb_arg(_cb_arg),

        offset(_offset),

        len(_len),

        scratch(_scratch),

        use_direct_io(_use_direct_io),

        alignment(_alignment),

        is_finished(false),

        is_being_aborted(false),

        req_count(0) {}

  struct iovec iov;

  struct io_uring* iu;

  std::function<void(FSReadRequest&, void*)> cb;

  void* cb_arg;

  uint64_t offset;

  size_t len;

  char* scratch;

  bool use_direct_io;

  size_t alignment;

  bool is_finished;

  // is_being_aborted is set by AbortIO when a cancel request is submitted.

  // Used to distinguish between aborted handles (expect 2 completions) and

  // non-aborted handles (expect 1 completion) when processing completions.

  bool is_being_aborted;

  // req_count is used by AbortIO API to keep track of number of requests.

  uint32_t req_count;

};

inline void UpdateResult(struct io_uring_cqe* cqe, const std::string& file_name,

                         size_t len, size_t iov_len, bool async_read,

                         bool use_direct_io, size_t alignment,

                         size_t& finished_len, FSReadRequest* req,

                         size_t& bytes_read, bool& read_again) {

  read_again = false;

  if (cqe->res < 0) {

    req->result = Slice(req->scratch, 0);

    req->status = IOError("Req failed", file_name, cqe->res);

  } else {

    bytes_read = static_cast<size_t>(cqe->res);

    TEST_SYNC_POINT_CALLBACK("UpdateResults::io_uring_result", &bytes_read);

    if (bytes_read == iov_len) {

      req->result = Slice(req->scratch, req->len);

      req->status = IOStatus::OK();

    } else if (bytes_read == 0) {

      /// cqe->res == 0 can means EOF, or can mean partial results. See

      // comment

      // https://github.com/facebook/rocksdb/pull/6441#issuecomment-589843435

      // Fall back to pread in this case.

      if (use_direct_io && !IsSectorAligned(finished_len, alignment)) {

        // Bytes reads don't fill sectors. Should only happen at the end

        // of the file.

        req->result = Slice(req->scratch, finished_len);

        req->status = IOStatus::OK();

      } else {

        if (async_read) {

          // No  bytes read. It can means EOF. In case of partial results, it's

          // caller responsibility to call read/readasync again.

          req->result = Slice(req->scratch, 0);

          req->status = IOStatus::OK();

        } else {

          read_again = true;

        }

      }

    } else if (bytes_read < iov_len) {

      assert(bytes_read > 0);

      if (async_read) {

        req->result = Slice(req->scratch, bytes_read);

        req->status = IOStatus::OK();

      } else {

        assert(bytes_read + finished_len < len);

        finished_len += bytes_read;

      }

    } else {

      req->result = Slice(req->scratch, 0);

      req->status = IOError("Req returned more bytes than requested", file_name,

                            cqe->res);

    }

  }

#ifdef NDEBUG

  (void)len;

#endif

}

// Finalize a completed async read request.

// Processes the CQE result, marks the handle as finished, and invokes the

// callback. This is shared between Poll and AbortIO (for non-aborted handles).

inline void FinalizeAsyncRead(struct io_uring* iu, struct io_uring_cqe* cqe,

                              Posix_IOHandle* posix_handle) {

  FSReadRequest req;

  req.scratch = posix_handle->scratch;

  req.offset = posix_handle->offset;

  req.len = posix_handle->len;

  size_t finished_len = 0;

  size_t bytes_read = 0;

  bool read_again = false;

  UpdateResult(cqe, "", req.len, posix_handle->iov.iov_len, true /*async_read*/,

               posix_handle->use_direct_io, posix_handle->alignment,

               finished_len, &req, bytes_read, read_again);

  posix_handle->is_finished = true;

  io_uring_cqe_seen(iu, cqe);

  posix_handle->cb(req, posix_handle->cb_arg);

}

#endif

#ifdef OS_LINUX

// Files under a specific directory have the same logical block size.

// This class caches the logical block size for the specified directories to

// save the CPU cost of computing the size.

// Safe for concurrent access from multiple threads without any external

// synchronization.

class LogicalBlockSizeCache {

 public:

  LogicalBlockSizeCache(

      std::function<size_t(int)> get_logical_block_size_of_fd =

          PosixHelper::GetLogicalBlockSizeOfFd,

      std::function<Status(const std::string&, size_t*)>

          get_logical_block_size_of_directory =

              PosixHelper::GetLogicalBlockSizeOfDirectory)

      : get_logical_block_size_of_fd_(get_logical_block_size_of_fd),

        get_logical_block_size_of_directory_(

            get_logical_block_size_of_directory) {}

  // Takes the following actions:

  // 1. Increases reference count of the directories;

  // 2. If the directory's logical block size is not cached,

  //    compute the buffer size and cache the result.

  Status RefAndCacheLogicalBlockSize(

      const std::vector<std::string>& directories);

  // Takes the following actions:

  // 1. Decreases reference count of the directories;

  // 2. If the reference count of a directory reaches 0, remove the directory

  //    from the cache.

  void UnrefAndTryRemoveCachedLogicalBlockSize(

      const std::vector<std::string>& directories);

  // Returns the logical block size for the file.

  //

  // If the file is under a cached directory, return the cached size.

  // Otherwise, the size is computed.

  size_t GetLogicalBlockSize(const std::string& fname, int fd);

  int GetRefCount(const std::string& dir) {

    ReadLock lock(&cache_mutex_);

    auto it = cache_.find(dir);

    if (it == cache_.end()) {

      return 0;

    }

    return it->second.ref;

  }

  size_t Size() const { return cache_.size(); }

  bool Contains(const std::string& dir) {

    ReadLock lock(&cache_mutex_);

    return cache_.find(dir) != cache_.end();

  }

 private:

  struct CacheValue {

    CacheValue() : size(0), ref(0) {}

    // Logical block size of the directory.

    size_t size;

    // Reference count of the directory.

    int ref;

  };

  std::function<size_t(int)> get_logical_block_size_of_fd_;

  std::function<Status(const std::string&, size_t*)>

      get_logical_block_size_of_directory_;

  std::map<std::string, CacheValue> cache_;

  port::RWMutex cache_mutex_;

};

#endif

class PosixSequentialFile : public FSSequentialFile {

 private:

  std::string filename_;

  FILE* file_;

  int fd_;

  bool use_direct_io_;

  size_t logical_sector_size_;

 public:

  PosixSequentialFile(const std::string& fname, FILE* file, int fd,

                      size_t logical_block_size, const EnvOptions& options);

  virtual ~PosixSequentialFile();

  IOStatus Read(size_t n, const IOOptions& opts, Slice* result, char* scratch,

                IODebugContext* dbg) override;

  IOStatus PositionedRead(uint64_t offset, size_t n, const IOOptions& opts,

                          Slice* result, char* scratch,

                          IODebugContext* dbg) override;

  IOStatus Skip(uint64_t n) override;

  IOStatus InvalidateCache(size_t offset, size_t length) override;

  bool use_direct_io() const override { return use_direct_io_; }

  size_t GetRequiredBufferAlignment() const override {

    return logical_sector_size_;

  }

};

#if defined(ROCKSDB_IOURING_PRESENT)

// io_uring instance queue depth

const unsigned int kIoUringDepth = 256;

inline void TsanAnnotateIOUringMemory(struct io_uring* iu) {

  TsanAnnotateMappedMemory(iu->sq.ring_ptr, iu->sq.ring_sz);

  // CQ and SQ can share the same mmap region.

  if (iu->cq.ring_ptr != iu->sq.ring_ptr) {

    TsanAnnotateMappedMemory(iu->cq.ring_ptr, iu->cq.ring_sz);

  }

  TsanAnnotateMappedMemory(

      iu->sq.sqes, static_cast<size_t>(kIoUringDepth) * sizeof(io_uring_sqe));

}

inline void DeleteIOUring(void* p) {

  struct io_uring* iu = static_cast<struct io_uring*>(p);

  io_uring_queue_exit(iu);

  delete iu;

}

inline struct io_uring* CreateIOUring() {

  struct io_uring* new_io_uring = new struct io_uring;

  unsigned int flags = 0;

  flags |= IORING_SETUP_SINGLE_ISSUER;

  flags |= IORING_SETUP_DEFER_TASKRUN;

  int ret = io_uring_queue_init(kIoUringDepth, new_io_uring, flags);

  if (ret) {

    fprintf(stdout, "CreateIOUring failed: %s (errno=%d), thread=%lu\n",

            errnoStr(-ret).c_str(), -ret,

            static_cast<unsigned long>(pthread_self()));

    delete new_io_uring;

    new_io_uring = nullptr;

  } else {

    TsanAnnotateIOUringMemory(new_io_uring);

  }

  return new_io_uring;

}

#endif  // defined(ROCKSDB_IOURING_PRESENT)

class PosixRandomAccessFile : public FSRandomAccessFile {

 protected:

  std::string filename_;

  int fd_;

  bool use_direct_io_;

  size_t logical_sector_size_;

#if defined(ROCKSDB_IOURING_PRESENT)

  ThreadLocalPtr* thread_local_async_read_io_urings_;

  ThreadLocalPtr* thread_local_multi_read_io_urings_;

#endif

 public:

  PosixRandomAccessFile(const std::string& fname, int fd,

                        size_t logical_block_size, const EnvOptions& options

#if defined(ROCKSDB_IOURING_PRESENT)

                        ,

                        ThreadLocalPtr* thread_local_async_read_io_urings,

                        ThreadLocalPtr* thread_local_multi_read_io_urings

#endif

  );

  virtual ~PosixRandomAccessFile();

  IOStatus Read(uint64_t offset, size_t n, const IOOptions& opts, Slice* result,

                char* scratch, IODebugContext* dbg) const override;

  IOStatus MultiRead(FSReadRequest* reqs, size_t num_reqs,

                     const IOOptions& options, IODebugContext* dbg) override;

  IOStatus Prefetch(uint64_t offset, size_t n, const IOOptions& opts,

                    IODebugContext* dbg) override;

#if defined(OS_LINUX) || defined(OS_MACOSX) || defined(OS_AIX)

  size_t GetUniqueId(char* id, size_t max_size) const override;

#endif

  void Hint(AccessPattern pattern) override;

  IOStatus InvalidateCache(size_t offset, size_t length) override;

  bool use_direct_io() const override { return use_direct_io_; }

  size_t GetRequiredBufferAlignment() const override {

    return logical_sector_size_;

  }

  virtual IOStatus ReadAsync(FSReadRequest& req, const IOOptions& opts,

                             std::function<void(FSReadRequest&, void*)> cb,

                             void* cb_arg, void** io_handle,

                             IOHandleDeleter* del_fn,

                             IODebugContext* dbg) override;

  virtual IOStatus GetFileSize(uint64_t* result) override;

};

class PosixWritableFile : public FSWritableFile {

 protected:

  const std::string filename_;

  const bool use_direct_io_;

  int fd_;

  uint64_t filesize_;

  size_t logical_sector_size_;

#ifdef ROCKSDB_FALLOCATE_PRESENT

  bool allow_fallocate_;

  bool fallocate_with_keep_size_;

#endif

#ifdef ROCKSDB_RANGESYNC_PRESENT

  // Even if the syscall is present, the filesystem may still not properly

  // support it, so we need to do a dynamic check too.

  bool sync_file_range_supported_;

#endif  // ROCKSDB_RANGESYNC_PRESENT

 public:

  explicit PosixWritableFile(const std::string& fname, int fd,

                             size_t logical_block_size,

                             const EnvOptions& options,

                             uint64_t initial_file_size);

  virtual ~PosixWritableFile();

  // Need to implement this so the file is truncated correctly

  // with direct I/O

  IOStatus Truncate(uint64_t size, const IOOptions& opts,

                    IODebugContext* dbg) override;

  IOStatus Close(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Append(const Slice& data, const IOOptions& opts,

                  IODebugContext* dbg) override;

  IOStatus Append(const Slice& data, const IOOptions& opts,

                  const DataVerificationInfo& /* verification_info */,

                  IODebugContext* dbg) override {

    return Append(data, opts, dbg);

  }

  IOStatus PositionedAppend(const Slice& data, uint64_t offset,

                            const IOOptions& opts,

                            IODebugContext* dbg) override;

  IOStatus PositionedAppend(const Slice& data, uint64_t offset,

                            const IOOptions& opts,

                            const DataVerificationInfo& /* verification_info */,

                            IODebugContext* dbg) override {

    return PositionedAppend(data, offset, opts, dbg);

  }

  IOStatus Flush(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Sync(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Fsync(const IOOptions& opts, IODebugContext* dbg) override;

  bool IsSyncThreadSafe() const override;

  bool use_direct_io() const override { return use_direct_io_; }

  void SetWriteLifeTimeHint(Env::WriteLifeTimeHint hint) override;

  uint64_t GetFileSize(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus InvalidateCache(size_t offset, size_t length) override;

  size_t GetRequiredBufferAlignment() const override {

    return logical_sector_size_;

  }

#ifdef ROCKSDB_FALLOCATE_PRESENT

  IOStatus Allocate(uint64_t offset, uint64_t len, const IOOptions& opts,

                    IODebugContext* dbg) override;

#endif

  IOStatus RangeSync(uint64_t offset, uint64_t nbytes, const IOOptions& opts,

                     IODebugContext* dbg) override;

#ifdef OS_LINUX

  size_t GetUniqueId(char* id, size_t max_size) const override;

#endif

};

// mmap() based random-access

class PosixMmapReadableFile : public FSRandomAccessFile {

 private:

  int fd_;

  std::string filename_;

  void* mmapped_region_;

  size_t length_;

 public:

  PosixMmapReadableFile(const int fd, const std::string& fname, void* base,

                        size_t length, const EnvOptions& options);

  virtual ~PosixMmapReadableFile();

  IOStatus Read(uint64_t offset, size_t n, const IOOptions& opts, Slice* result,

                char* scratch, IODebugContext* dbg) const override;

  void Hint(AccessPattern pattern) override;

  IOStatus InvalidateCache(size_t offset, size_t length) override;

  virtual IOStatus GetFileSize(uint64_t* result) override;

};

class PosixMmapFile : public FSWritableFile {

 private:

  std::string filename_;

  int fd_;

  size_t page_size_;

  size_t map_size_;       // How much extra memory to map at a time

  char* base_;            // The mapped region

  char* limit_;           // Limit of the mapped region

  char* dst_;             // Where to write next  (in range [base_,limit_])

  char* last_sync_;       // Where have we synced up to

  uint64_t file_offset_;  // Offset of base_ in file

#ifdef ROCKSDB_FALLOCATE_PRESENT

  bool allow_fallocate_;  // If false, fallocate calls are bypassed

  bool fallocate_with_keep_size_;

#endif

  // Roundup x to a multiple of y

  static size_t Roundup(size_t x, size_t y) { return ((x + y - 1) / y) * y; }

  size_t TruncateToPageBoundary(size_t s) {

    s -= (s & (page_size_ - 1));

    assert((s % page_size_) == 0);

    return s;

  }

  IOStatus MapNewRegion();

  IOStatus UnmapCurrentRegion();

  IOStatus Msync();

 public:

  PosixMmapFile(const std::string& fname, int fd, size_t page_size,

                const EnvOptions& options, uint64_t initial_file_size);

  ~PosixMmapFile();

  // Means Close() will properly take care of truncate

  // and it does not need any additional information

  IOStatus Truncate(uint64_t /*size*/, const IOOptions& /*opts*/,

                    IODebugContext* /*dbg*/) override {

    return IOStatus::OK();

  }

  IOStatus Close(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Append(const Slice& data, const IOOptions& opts,

                  IODebugContext* dbg) override;

  IOStatus Append(const Slice& data, const IOOptions& opts,

                  const DataVerificationInfo& /* verification_info */,

                  IODebugContext* dbg) override {

    return Append(data, opts, dbg);

  }

  IOStatus Flush(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Sync(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Fsync(const IOOptions& opts, IODebugContext* dbg) override;

  uint64_t GetFileSize(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus InvalidateCache(size_t offset, size_t length) override;

#ifdef ROCKSDB_FALLOCATE_PRESENT

  IOStatus Allocate(uint64_t offset, uint64_t len, const IOOptions& opts,

                    IODebugContext* dbg) override;

#endif

};

class PosixRandomRWFile : public FSRandomRWFile {

 public:

  explicit PosixRandomRWFile(const std::string& fname, int fd,

                             const EnvOptions& options);

  virtual ~PosixRandomRWFile();

  IOStatus Write(uint64_t offset, const Slice& data, const IOOptions& opts,

                 IODebugContext* dbg) override;

  IOStatus Read(uint64_t offset, size_t n, const IOOptions& opts, Slice* result,

                char* scratch, IODebugContext* dbg) const override;

  IOStatus Flush(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Sync(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Fsync(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Close(const IOOptions& opts, IODebugContext* dbg) override;

 private:

  const std::string filename_;

  int fd_;

};

struct PosixMemoryMappedFileBuffer : public MemoryMappedFileBuffer {

  PosixMemoryMappedFileBuffer(void* _base, size_t _length)

      : MemoryMappedFileBuffer(_base, _length) {}

  virtual ~PosixMemoryMappedFileBuffer();

};

class PosixDirectory : public FSDirectory {

 public:

  explicit PosixDirectory(int fd, const std::string& directory_name);

  ~PosixDirectory();

  IOStatus Fsync(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus Close(const IOOptions& opts, IODebugContext* dbg) override;

  IOStatus FsyncWithDirOptions(

      const IOOptions&, IODebugContext*,

      const DirFsyncOptions& dir_fsync_options) override;

 private:

  int fd_;

  bool is_btrfs_;

  const std::string directory_name_;

};

}  // namespace ROCKSDB_NAMESPACE