mmap_array.h

Go to the documentation of this file.

 
#ifndef GRAPHSCOPE_UTILS_MMAP_ARRAY_H_
#define GRAPHSCOPE_UTILS_MMAP_ARRAY_H_
 
#include <assert.h>
 
#include <atomic>
#include <filesystem>
#include <string>
#include <string_view>
 
#include "flex/storages/rt_mutable_graph/file_names.h"
#include "glog/logging.h"
#include "grape/util.h"
 
#ifdef __ia64__
#define ADDR (void*) (0x8000000000000000UL)
#define FLAGS (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB | MAP_FIXED)
#else
#define ADDR (void*) (0x0UL)
#define FLAGS (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB)
#endif
 
#define PROTECTION (PROT_READ | PROT_WRITE)
 
#define HUGEPAGE_SIZE (2UL * 1024 * 1024)
#define HUGEPAGE_MASK (2UL * 1024 * 1024 - 1UL)
#define ROUND_UP(size) (((size) + HUGEPAGE_MASK) & (~HUGEPAGE_MASK))
 
inline void* allocate_hugepages(size_t size) {
  return mmap(ADDR, ROUND_UP(size), PROTECTION, FLAGS, -1, 0);
}
 
inline size_t hugepage_round_up(size_t size) { return ROUND_UP(size); }
 
#undef ADDR
#undef FLAGS
#undef HUGEPAGE_SIZE
#undef HUGEPAGE_MASK
#undef ROUND_UP
 
namespace gs {
 
enum class MemoryStrategy {
  kSyncToFile,
  kMemoryOnly,
  kHugepagePrefered,
};
 
template <typename T>
class mmap_array {
 public:
  mmap_array()
      : filename_(""),
        fd_(-1),
        data_(NULL),
        size_(0),
        mmap_size_(0),
        sync_to_file_(false),
        hugepage_prefered_(false) {}
 
  mmap_array(const mmap_array<T>& rhs) : fd_(-1) {
    resize(rhs.size_);
    memcpy(data_, rhs.data_, size_ * sizeof(T));
  }
 
  mmap_array(mmap_array&& rhs) : mmap_array() { swap(rhs); }
  ~mmap_array() { reset(); }
 
  void reset() {
    if (data_ != NULL && mmap_size_ != 0) {
      if (munmap(data_, mmap_size_) != 0) {
        std::stringstream ss;
        ss << "Failed to mummap file [ " << filename_ << " ] "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
    }
    data_ = NULL;
    size_ = 0;
    mmap_size_ = 0;
    if (fd_ != -1) {
      if (close(fd_) != 0) {
        std::stringstream ss;
        ss << "Failed to close file [ " << filename_ << " ] "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      fd_ = -1;
    }
    filename_ = "";
    sync_to_file_ = false;
  }
 
  void unlink() {
    std::string old_filename = filename_;
    reset();
    if (old_filename != "" && std::filesystem::exists(old_filename)) {
      if (std::filesystem::remove(old_filename) == 0) {
        std::stringstream ss;
        ss << "Failed to remove file [ " << old_filename << " ] "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
    }
  }
 
  void set_hugepage_prefered(bool val) {
    hugepage_prefered_ = (val && !sync_to_file_);
  }
 
  void open(const std::string& filename, bool sync_to_file = false) {
    reset();
    filename_ = filename;
    sync_to_file_ = sync_to_file;
    hugepage_prefered_ = false;
    if (sync_to_file_) {
      bool creat = !std::filesystem::exists(filename_);
      fd_ = ::open(filename_.c_str(), O_RDWR | O_CREAT, 0777);
      if (fd_ == -1) {
        std::stringstream ss;
        ss << "Failed to open file [" << filename_ << "], " << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      if (creat) {
        std::filesystem::perms readWritePermission =
            std::filesystem::perms::owner_read |
            std::filesystem::perms::owner_write;
        std::error_code errorCode;
        std::filesystem::permissions(filename, readWritePermission,
                                     std::filesystem::perm_options::add,
                                     errorCode);
        if (errorCode) {
          std::stringstream ss;
          ss << "Failed to set read/write permission for file: " << filename
             << " " << errorCode.message() << std::endl;
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
      }
 
      size_t file_size = std::filesystem::file_size(filename_);
      size_ = file_size / sizeof(T);
      mmap_size_ = file_size;
      if (mmap_size_ == 0) {
        data_ = NULL;
      } else {
        data_ = reinterpret_cast<T*>(
            mmap(NULL, mmap_size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
        if (data_ == MAP_FAILED) {
          std::stringstream ss;
          ss << "Failed to mmap file [" << filename_ << "], "
             << strerror(errno);
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
        int rt = madvise(data_, mmap_size_, MADV_RANDOM | MADV_WILLNEED);
        if (rt != 0) {
          std::stringstream ss;
          ss << "Failed to madvise file [" << filename_ << "], "
             << strerror(errno);
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
      }
    } else {
      if (!filename_.empty() && std::filesystem::exists(filename_)) {
        size_t file_size = std::filesystem::file_size(filename_);
        fd_ = ::open(filename_.c_str(), O_RDWR, 0777);
        if (fd_ == -1) {
          std::stringstream ss;
          ss << "Failed to open file [" << filename_ << "], "
             << strerror(errno);
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
        size_ = file_size / sizeof(T);
        mmap_size_ = file_size;
        if (mmap_size_ == 0) {
          data_ = NULL;
        } else {
          data_ = reinterpret_cast<T*>(mmap(
              NULL, mmap_size_, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd_, 0));
          if (data_ == MAP_FAILED) {
            std::stringstream ss;
            ss << "Failed to mmap file [" << filename_ << "], "
               << strerror(errno);
            LOG(ERROR) << ss.str();
            throw std::runtime_error(ss.str());
          }
        }
      }
    }
  }
 
  void open_with_hugepages(const std::string& filename, size_t capacity = 0) {
    reset();
    hugepage_prefered_ = true;
    if (!filename.empty() && std::filesystem::exists(filename)) {
      size_t file_size = std::filesystem::file_size(filename);
      size_ = file_size / sizeof(T);
      if (size_ != 0) {
        capacity = std::max(capacity, size_);
        mmap_size_ = hugepage_round_up(capacity * sizeof(T));
        data_ = static_cast<T*>(allocate_hugepages(mmap_size_));
        if (data_ != MAP_FAILED) {
          FILE* fin = fopen(filename.c_str(), "rb");
          if (fin == NULL) {
            std::stringstream ss;
            ss << "Failed to open file [ " << filename << " ], "
               << strerror(errno);
            LOG(ERROR) << ss.str();
            throw std::runtime_error(ss.str());
          }
          if (fread(data_, sizeof(T), size_, fin) != size_) {
            std::stringstream ss;
            ss << "Failed to fread file [ " << filename << " ], "
               << strerror(errno);
            LOG(ERROR) << ss.str();
            throw std::runtime_error(ss.str());
          }
          if (fclose(fin) != 0) {
            std::stringstream ss;
            ss << "Failed to fclose file [ " << filename << " ], "
               << strerror(errno);
            LOG(ERROR) << ss.str();
            throw std::runtime_error(ss.str());
          }
        } else {
          LOG(ERROR) << "allocating hugepage failed, " << strerror(errno)
                     << ", try with normal pages";
          data_ = NULL;
          open(filename, false);
        }
      } else {
        mmap_size_ = 0;
      }
    }
  }
 
  void dump(const std::string& filename) {
    if (sync_to_file_) {
      std::string old_filename = filename_;
      reset();
      std::error_code errorCode;
      std::filesystem::rename(old_filename, filename, errorCode);
      if (errorCode) {
        std::stringstream ss;
        ss << "Failed to rename file " << old_filename << " to " << filename
           << " " << errorCode.message() << std::endl;
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
    } else {
      FILE* fout = fopen(filename.c_str(), "wb");
      if (fout == NULL) {
        std::stringstream ss;
        ss << "Failed to open file [ " << filename << " ], " << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      if (fwrite(data_, sizeof(T), size_, fout) != size_) {
        std::stringstream ss;
        ss << "Failed to fwrite file [ " << filename << " ], "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      if (fflush(fout) != 0) {
        std::stringstream ss;
        ss << "Failed to fflush file [ " << filename << " ], "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      if (fclose(fout) != 0) {
        std::stringstream ss;
        ss << "Failed to fclose file [ " << filename << " ], "
           << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      reset();
    }
 
    std::filesystem::perms readPermission = std::filesystem::perms::owner_read;
 
    std::error_code errorCode;
    std::filesystem::permissions(filename, readPermission,
                                 std::filesystem::perm_options::add, errorCode);
 
    if (errorCode) {
      std::stringstream ss;
      ss << "Failed to set read permission for file: " << filename << " "
         << errorCode.message() << std::endl;
      LOG(ERROR) << ss.str();
      throw std::runtime_error(ss.str());
    }
  }
 
  void resize(size_t size) {
    if (size == size_) {
      return;
    }
 
    if (sync_to_file_) {
      if (data_ != NULL && mmap_size_ != 0) {
        if (munmap(data_, mmap_size_) != 0) {
          std::stringstream ss;
          ss << "Failed to mummap file [ " << filename_ << " ], "
             << strerror(errno);
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
      }
      size_t new_mmap_size = size * sizeof(T);
      int rt = ftruncate(fd_, new_mmap_size);
      if (rt == -1) {
        std::stringstream ss;
        ss << "Failed to ftruncate " << rt << ", " << strerror(errno);
        LOG(ERROR) << ss.str();
        throw std::runtime_error(ss.str());
      }
      if (new_mmap_size == 0) {
        data_ = NULL;
      } else {
        data_ = reinterpret_cast<T*>(mmap(
            NULL, new_mmap_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
        if (data_ == MAP_FAILED) {
          std::stringstream ss;
          ss << "Failed to mmap, " << strerror(errno);
          LOG(ERROR) << ss.str();
          throw std::runtime_error(ss.str());
        }
      }
      size_ = size;
      mmap_size_ = new_mmap_size;
    } else {
      size_t target_mmap_size = size * sizeof(T);
      if (target_mmap_size <= mmap_size_) {
        size_ = size;
      } else {
        T* new_data = NULL;
        size_t new_mmap_size = size * sizeof(T);
        if (hugepage_prefered_) {
          new_data = reinterpret_cast<T*>(allocate_hugepages(new_mmap_size));
          if (new_data == MAP_FAILED) {
            LOG(ERROR) << "mmap with hugepage failed, " << strerror(errno)
                       << ", try with normal pages";
            new_data = NULL;
          } else {
            new_mmap_size = hugepage_round_up(new_mmap_size);
          }
        }
        if (new_data == NULL) {
          new_data = reinterpret_cast<T*>(
              mmap(NULL, new_mmap_size, PROT_READ | PROT_WRITE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
          if (new_data == MAP_FAILED) {
            std::stringstream ss;
            ss << "mmap failed " << strerror(errno);
            LOG(ERROR) << ss.str();
            throw std::runtime_error(ss.str());
          }
        }
 
        size_t copy_size = std::min(size, size_);
        if (copy_size > 0 && data_ != NULL) {
          memcpy(reinterpret_cast<void*>(new_data),
                 reinterpret_cast<void*>(data_), copy_size * sizeof(T));
        }
 
        reset();
 
        data_ = new_data;
        size_ = size;
        mmap_size_ = new_mmap_size;
      }
    }
  }
 
  void touch(const std::string& filename) {
    dump(filename);
    open(filename, true);
  }
 
  T* data() { return data_; }
  const T* data() const { return data_; }
 
  void set(size_t idx, const T& val) { data_[idx] = val; }
 
  const T& get(size_t idx) const { return data_[idx]; }
 
  const T& operator[](size_t idx) const { return data_[idx]; }
  T& operator[](size_t idx) { return data_[idx]; }
 
  size_t size() const { return size_; }
 
  void swap(mmap_array<T>& rhs) {
    std::swap(filename_, rhs.filename_);
    std::swap(fd_, rhs.fd_);
    std::swap(data_, rhs.data_);
    std::swap(size_, rhs.size_);
    std::swap(mmap_size_, rhs.mmap_size_);
    std::swap(hugepage_prefered_, rhs.hugepage_prefered_);
    std::swap(sync_to_file_, rhs.sync_to_file_);
  }
 
  const std::string& filename() const { return filename_; }
 
 private:
  std::string filename_;
  int fd_;
  T* data_;
  size_t size_;
 
  size_t mmap_size_;
 
  bool sync_to_file_;
  bool hugepage_prefered_;
};
 
struct string_item {
  uint64_t offset : 48;
  uint32_t length : 16;
};
 
template <>
class mmap_array<std::string_view> {
 public:
  mmap_array() {}
  mmap_array(mmap_array&& rhs) : mmap_array() { swap(rhs); }
  ~mmap_array() {}
 
  void reset() {
    items_.reset();
    data_.reset();
  }
 
  void set_hugepage_prefered(bool val) {
    items_.set_hugepage_prefered(val);
    data_.set_hugepage_prefered(val);
  }
 
  void open(const std::string& filename, bool sync_to_file) {
    items_.open(filename + ".items", sync_to_file);
    data_.open(filename + ".data", sync_to_file);
  }
 
  void open_with_hugepages(const std::string& filename) {
    items_.open_with_hugepages(filename + ".items");
    data_.open_with_hugepages(filename + ".data");
  }
 
  void touch(const std::string& filename) {
    items_.touch(filename + ".items");
    data_.touch(filename + ".data");
  }
 
  void dump(const std::string& filename) {
    items_.dump(filename + ".items");
    data_.dump(filename + ".data");
  }
 
  void resize(size_t size, size_t data_size) {
    items_.resize(size);
    data_.resize(data_size);
  }
 
  void set(size_t idx, size_t offset, const std::string_view& val) {
    items_.set(idx, {offset, static_cast<uint32_t>(val.size())});
    memcpy(data_.data() + offset, val.data(), val.size());
  }
 
  std::string_view get(size_t idx) const {
    const string_item& item = items_.get(idx);
    return std::string_view(data_.data() + item.offset, item.length);
  }
 
  size_t size() const { return items_.size(); }
 
  size_t data_size() const { return data_.size(); }
 
  void swap(mmap_array& rhs) {
    items_.swap(rhs.items_);
    data_.swap(rhs.data_);
  }
  void unlink() {
    items_.unlink();
    data_.unlink();
  }
 
 private:
  mmap_array<string_item> items_;
  mmap_array<char> data_;
};
 
}  // namespace gs
 
#endif  // GRAPHSCOPE_UTILS_MMAP_ARRAY_H_