text: /*
   * Copyright (c) Facebook, Inc. and its affiliates.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #include "velox/common/caching/SsdFile.h"
  #include <folly/Executor.h>
  #include <folly/portability/SysUio.h>
  #include "velox/common/caching/FileIds.h"
  
  #include <fcntl.h>
  #include <sys/stat.h>
  #include <sys/types.h>
  #include <numeric>
  
  DEFINE_bool(ssd_odirect, true, "Use O_DIRECT for SSD cache IO");
  DEFINE_bool(ssd_verify_write, false, "Read back data after writing to SSD");
  
  namespace facebook::velox::cache {
  
  SsdPin::SsdPin(SsdFile& file, SsdRun run) : file_(&file), run_(run) {
    file_->checkPinned(run_.offset());
  }
  
  SsdPin::~SsdPin() {
    if (file_) {
      file_->unpinRegion(run_.offset());
    }
  }
  
  void SsdPin::operator=(SsdPin&& other) {
    if (file_) {
      file_->unpinRegion(run_.offset());
    }
    file_ = other.file_;
    other.file_ = nullptr;
    run_ = other.run_;
  }
  
  SsdFile::SsdFile(
      const std::string& filename,
      int32_t shardId,
      int32_t maxRegions)
      : shardId_(shardId), maxRegions_(maxRegions), filename_(filename) {
    int32_t oDirect = 0;
  #ifdef linux
    oDirect = FLAGS_ssd_odirect ? O_DIRECT : 0;
  #endif
    fd_ = open(filename.c_str(), O_CREAT | O_RDWR | oDirect, S_IRUSR | S_IWUSR);
    if (fd_ < 0) {
      LOG(ERROR) << "Cannot open or create " << filename << " error " << errno;
      exit(1);
    }
    readFile_ = std::make_unique<LocalReadFile>(fd_);
    uint64_t size = lseek(fd_, 0, SEEK_END);
    numRegions_ = size / kRegionSize;
    fileSize_ = numRegions_ * kRegionSize;
    if (size % kRegionSize > 0) {
      ftruncate(fd_, fileSize_);
    }
    // The existing regions in the file are writable.
    writableRegions_.resize(numRegions_);
    std::iota(writableRegions_.begin(), writableRegions_.end(), 0);
    tracker_.resize(maxRegions_);
    regionSize_.resize(maxRegions_);
    regionPins_.resize(maxRegions_);
  }
  
  void SsdFile::pinRegion(uint64_t offset) {
    std::lock_guard<std::mutex> l(mutex_);
    pinRegionLocked(offset);
  }
  
  void SsdFile::unpinRegion(uint64_t offset) {
    std::lock_guard<std::mutex> l(mutex_);
    auto count = --regionPins_[regionIndex(offset)];
    VELOX_CHECK_LE(0, count);
    if (suspended_ && count == 0) {
      growOrEvictLocked();
    }
  }
  
  namespace {
  void addEntryToIovecs(AsyncDataCacheEntry& entry, std::vector<iovec>& iovecs) {
    if (entry.tinyData()) {
      iovecs.push_back({entry.tinyData(), static_cast<size_t>(entry.size())});
      return;
    }
    auto& data = entry.data();
    iovecs.reserve(iovecs.size() + data.numRuns());
    int64_t bytesLeft = entry.size();
    for (auto i = 0; i < data.numRuns(); ++i) {
      auto run = data.runAt(i);
      iovecs.push_back(
          {run.data<char>(), std::min<size_t>(bytesLeft, run.numBytes())});
      bytesLeft -= run.numBytes();
      if (bytesLeft <= 0) {
        break;
      };
    }
  }
  } // namespace
  
  SsdPin SsdFile::find(RawFileCacheKey key) {
    FileCacheKey ssdKey{StringIdLease(fileIds(), key.fileNum), key.offset};
    SsdRun run;
    {
      std::lock_guard<std::mutex> l(mutex_);
      if (suspended_) {
        return SsdPin();
      }
      tracker_.fileTouched(entries_.size());
      auto it = entries_.find(ssdKey);
      if (it == entries_.end()) {
        return SsdPin();
      }
      run = it->second;
      pinRegionLocked(run.offset());
    }
    return SsdPin(*this, run);
  }
  
  CoalesceIoStats SsdFile::load(
      const std::vector<SsdPin>& ssdPins,
      const std::vector<CachePin>& pins) {
    VELOX_CHECK_EQ(ssdPins.size(), pins.size());
    int payloadTotal = 0;
    for (auto i = 0; i < pins.size(); ++i) {
      auto runSize = ssdPins[i].run().size();
      VELOX_CHECK_EQ(runSize, pins[i].entry()->size());
      payloadTotal += runSize;
      regionRead(regionIndex(ssdPins[i].run().offset()), runSize);
      ++stats_.entriesRead;
      stats_.bytesRead += runSize;
    }
    // Do coalesced IO for the pins. For short payloads, the break-even
    // between discrete pread calls and a single preadv that discards
    // gaps is ~25K per gap. For longer payloads this is ~50-100K.
    auto stats = readPins(
        pins,
        payloadTotal / pins.size() < 10000 ? 25000 : 50000,
        // Max ranges in one preadv call. Longest gap + longest cache
        // entry are under 12 ranges. If a system has a limit of 1K
        // ranges, coalesce limit of 1000 is safe.
        900,
        [&](int32_t index) { return ssdPins[index].run().offset(); },
        [&](const std::vector<CachePin>& /*pins*/,
            int32_t /*begin*/,
            int32_t /*end*/,
            uint64_t offset,
            const std::vector<folly::Range<char*>>& buffers) {
          read(offset, buffers);
        });
  
    for (auto i = 0; i < ssdPins.size(); ++i) {
      pins[i].checkedEntry()->setSsdFile(this, ssdPins[i].run().offset());
    }
    return stats;
  }
  
  void SsdFile::read(
      uint64_t offset,
      const std::vector<folly::Range<char*>>& buffers) {
    readFile_->preadv(offset, buffers);
  }
  
  std::optional<std::pair<uint64_t, int32_t>> SsdFile::getSpace(
      const std::vector<CachePin>& pins,
      int32_t begin) {
    std::lock_guard<std::mutex> l(mutex_);
    for (;;) {
      if (writableRegions_.empty()) {
        if (!growOrEvictLocked()) {
          return std::nullopt;
          ;
        }
      }
      auto region = writableRegions_[0];
      auto offset = regionSize_[region];
      auto available = kRegionSize - offset;
      int64_t toWrite = 0;
      for (; begin < pins.size(); ++begin) {
        auto entry = pins[begin].checkedEntry();
        if (entry->size() > available) {
          break;
        }
        available -= entry->size();
        toWrite += entry->size();
      }
      if (toWrite) {
        // At least some pins got space from this region. If the region is full
        // the next call will get space from another region.
        regionSize_[region] += toWrite;
        return std::make_pair<uint64_t, int32_t>(
            region * kRegionSize + offset, toWrite);
      }
  
      tracker_.regionFilled(region);
      writableRegions_.erase(writableRegions_.begin());
    }
  }
  
  bool SsdFile::growOrEvictLocked() {
    if (numRegions_ < maxRegions_) {
      auto newSize = (numRegions_ + 1) * kRegionSize;
      auto rc = ftruncate(fd_, newSize);
      if (rc >= 0) {
        fileSize_ = newSize;
        writableRegions_.push_back(numRegions_);
        regionSize_[numRegions_] = 0;
        ++numRegions_;
        return true;
      } else {
        LOG(ERROR) << "Failed to grow cache file " << filename_ << " to "
                   << newSize;
      }
    }
    auto candidates =
        tracker_.findEvictionCandidates(3, numRegions_, regionPins_);
    if (candidates.empty()) {
      suspended_ = true;
      return false;
    }
    clearRegionEntriesLocked(candidates);
    writableRegions_ = std::move(candidates);
    suspended_ = false;
    return true;
  }
  
  void SsdFile::clearRegionEntriesLocked(
      const std::vector<int32_t>& regionIndices) {
    // Remove all 'entries_' where the dependent points one of 'regionIndices'.
    auto it = entries_.begin();
    while (it != entries_.end()) {
      auto region = regionIndex(it->second.offset());
      if (std::find(regionIndices.begin(), regionIndices.end(), region) !=
          regionIndices.end()) {
        it = entries_.erase(it);
      } else {
        ++it;
      }
    }
    for (auto region : regionIndices) {
      // While the region is being filled it may get score from
      // hits. When it is full, it will get a score boost to be a little
      // ahead of the best.
      tracker_.regionCleared(region);
      regionSize_[region] = 0;
    }
  }
  
  void SsdFile::write(std::vector<CachePin>& pins) {
    // Sorts the pins by their file/offset. In this way what is ajacent
    // in storage is likely adjacent on SSD.
    std::sort(pins.begin(), pins.end());
    uint64_t total = 0;
    for (auto& pin : pins) {
      auto entry = pin.checkedEntry();
      VELOX_CHECK_NULL(entry->ssdFile());
      total += entry->size();
    }
    int32_t storeIndex = 0;
    while (storeIndex < pins.size()) {
      auto space = getSpace(pins, storeIndex);
  
      if (!space.has_value()) {
        // No space can be reclaimed. The pins are freed when the caller is freed.
        return;
      }
      auto [offset, available] = space.value();
      int32_t numWritten = 0;
      int32_t bytes = 0;
      std::vector<iovec> iovecs;
      for (auto i = storeIndex; i < pins.size(); ++i) {
        auto entry = pins[i].checkedEntry();
        auto entrySize = entry->size();
        if (bytes + entrySize > available) {
          break;
        }
        addEntryToIovecs(*entry, iovecs);
        bytes += entrySize;
        ++numWritten;
      }
      VELOX_CHECK_GE(fileSize_, offset + bytes);
      auto rc = folly::pwritev(fd_, iovecs.data(), iovecs.size(), offset);
      if (rc != bytes) {
        LOG(ERROR) << "Failed to write to SSD " << errno;
        // If the write fails we return without adding the pins to the cache. The
        // entries are unchanged.
        return;
      }
      {
        std::lock_guard<std::mutex> l(mutex_);
        for (auto i = storeIndex; i < storeIndex + numWritten; ++i) {
          auto entry = pins[i].checkedEntry();
          entry->setSsdFile(this, offset);
          auto size = entry->size();
          FileCacheKey key = {
              entry->key().fileNum, static_cast<uint64_t>(entry->offset())};
          entries_[std::move(key)] = SsdRun(offset, size);
          if (FLAGS_ssd_verify_write) {
            verifyWrite(*entry, SsdRun(offset, size));
          }
          offset += size;
          ++stats_.entriesWritten;
          stats_.bytesWritten += size;
        }
      }
      storeIndex += numWritten;
    }
  }
  
  namespace {
  int32_t indexOfFirstMismatch(char* x, char* y, int n) {
    for (auto i = 0; i < n; ++i) {
      if (x[i] != y[i]) {
        return i;
      }
    }
    return -1;
  }
  } // namespace
  
  void SsdFile::verifyWrite(AsyncDataCacheEntry& entry, SsdRun ssdRun) {
    auto testData = std::make_unique<char[]>(entry.size());
    auto rc = pread(fd_, testData.get(), entry.size(), ssdRun.offset());
    VELOX_CHECK_EQ(rc, entry.size());
    if (entry.tinyData()) {
      if (0 != memcmp(testData.get(), entry.tinyData(), entry.size())) {
        VELOX_FAIL("bad read back");
      }
    } else {
      auto& data = entry.data();
      int64_t bytesLeft = entry.size();
      int64_t offset = 0;
      for (auto i = 0; i < data.numRuns(); ++i) {
        auto run = data.runAt(i);
        auto compareSize = std::min<int64_t>(bytesLeft, run.numBytes());
        auto badIndex = indexOfFirstMismatch(
            run.data<char>(), testData.get() + offset, compareSize);
        if (badIndex != -1) {
          VELOX_FAIL("Bad read back");
        }
        bytesLeft -= run.numBytes();
        offset += run.numBytes();
        if (bytesLeft <= 0) {
          break;
        };
      }
    }
  }
  
  void SsdFile::updateStats(SsdCacheStats& stats) const {
    stats.entriesWritten += stats_.entriesWritten;
    stats.bytesWritten += stats_.bytesWritten;
    stats.entriesRead += stats_.entriesRead;
    stats.bytesRead += stats_.bytesRead;
    stats.entriesCached += entries_.size();
    for (auto& regionSize : regionSize_) {
      stats.bytesCached += regionSize;
    }
  }
  
  void SsdFile::clear() {
    std::lock_guard<std::mutex> l(mutex_);
    entries_.clear();
    std::fill(regionSize_.begin(), regionSize_.end(), 0);
    writableRegions_.resize(numRegions_);
    std::iota(writableRegions_.begin(), writableRegions_.end(), 0);
  }
  
  } // namespace facebook::velox::cache
  