GlobalFileMgr.h

Go to the documentation of this file.

/*
 * Copyright 2022 HEAVY.AI, Inc.
 *
 * 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.
 */

#pragma once

#include <iostream>
#include <map>
#include <mutex>
#include <set>
#include "DataMgr/AbstractBuffer.h"
#include "DataMgr/AbstractBufferMgr.h"
#include "FileMgr.h"
#include "Shared/heavyai_shared_mutex.h"

class ForeignStorageInterface;

using namespace Data_Namespace;

namespace File_Namespace {

struct FileMgrParams {
  FileMgrParams() : epoch(-1), max_rollback_epochs(-1) {}
  int32_t epoch;
  int32_t max_rollback_epochs;
};

class GlobalFileMgr : public AbstractBufferMgr {  // implements

 public:
  GlobalFileMgr(const int32_t device_id,
                std::shared_ptr<ForeignStorageInterface> fsi,
                std::string base_path = ".",
                const size_t num_reader_threads = 0,
                const size_t page_size = DEFAULT_PAGE_SIZE,
                const size_t metadata_page_size = DEFAULT_METADATA_PAGE_SIZE);

  ~GlobalFileMgr() override {}

  AbstractBuffer* createBuffer(const ChunkKey& key,
                               size_t pageSize = 0,
                               const size_t numBytes = 0) override {
    return getFileMgr(key)->createBuffer(key, pageSize, numBytes);
  }

  bool isBufferOnDevice(const ChunkKey& key) override {
    return getFileMgr(key)->isBufferOnDevice(key);
  }

  // Purge == true means delete the data chunks -
  // can't undelete and revert to previous
  // state - reclaims disk space for chunk
  void deleteBuffer(const ChunkKey& key, const bool purge = true) override {
    return getFileMgr(key)->deleteBuffer(key, purge);
  }

  void deleteBuffersWithPrefix(const ChunkKey& keyPrefix,
                               const bool purge = true) override;

  AbstractBuffer* getBuffer(const ChunkKey& key, const size_t numBytes = 0) override {
    return getFileMgr(key)->getBuffer(key, numBytes);
  }

  void fetchBuffer(const ChunkKey& key,
                   AbstractBuffer* destBuffer,
                   const size_t numBytes) override {
    return getFileMgr(key)->fetchBuffer(key, destBuffer, numBytes);
  }

  AbstractBuffer* putBuffer(const ChunkKey& key,
                            AbstractBuffer* d,
                            const size_t numBytes = 0) override {
    return getFileMgr(key)->putBuffer(key, d, numBytes);
  }

  // Buffer API
  AbstractBuffer* alloc(const size_t numBytes) override {
    LOG(FATAL) << "Operation not supported";
    return nullptr;  // satisfy return-type warning
  }

  void free(AbstractBuffer* buffer) override { LOG(FATAL) << "Operation not supported"; }

  inline MgrType getMgrType() override { return GLOBAL_FILE_MGR; };
  inline std::string getStringMgrType() override { return ToString(GLOBAL_FILE_MGR); }
  inline std::string printSlabs() override { return "Not Implemented"; }
  inline size_t getMaxSize() override { return 0; }
  inline size_t getInUseSize() override { return 0; }
  inline size_t getAllocated() override { return 0; }
  inline bool isAllocationCapped() override { return false; }

  void init();

  void getChunkMetadataVecForKeyPrefix(ChunkMetadataVector& chunkMetadataVec,
                                       const ChunkKey& keyPrefix) override {
    return getFileMgr(keyPrefix)->getChunkMetadataVecForKeyPrefix(chunkMetadataVec,
                                                                  keyPrefix);
  }

  void checkpoint() override;
  void checkpoint(const int32_t db_id, const int32_t tb_id) override;

  inline size_t getNumReaderThreads() { return num_reader_threads_; }

  size_t getNumChunks() override;

  void compactDataFiles(const int32_t db_id, const int32_t tb_id);

  static constexpr int32_t db_version_{2};

 private:
  AbstractBufferMgr* findFileMgrUnlocked(const int32_t db_id, const int32_t tb_id);
  void deleteFileMgr(const int32_t db_id, const int32_t tb_id);

 public:
  AbstractBufferMgr* findFileMgr(const int32_t db_id, const int32_t tb_id) {
    heavyai::shared_lock<heavyai::shared_mutex> read_lock(fileMgrs_mutex_);
    return findFileMgrUnlocked(db_id, tb_id);
  }
  void setFileMgrParams(const int32_t db_id,
                        const int32_t tb_id,
                        const FileMgrParams& file_mgr_params);
  AbstractBufferMgr* getFileMgr(const int32_t db_id, const int32_t tb_id);
  AbstractBufferMgr* getFileMgr(const ChunkKey& key) {
    return getFileMgr(key[0], key[1]);
  }

  std::string getBasePath() const { return basePath_; }
  size_t getPageSize() const { return page_size_; }
  size_t getMetadataPageSize() const { return metadata_page_size_; }

  void writeFileMgrData(FileMgr* fileMgr = 0);

  inline bool getDBConvert() const { return dbConvert_; }
  inline void setDBConvert(bool val) { dbConvert_ = val; }

  void removeTableRelatedDS(const int32_t db_id, const int32_t tb_id) override;
  void setTableEpoch(const int32_t db_id, const int32_t tb_id, const int32_t start_epoch);
  size_t getTableEpoch(const int32_t db_id, const int32_t tb_id);
  void resetTableEpochFloor(const int32_t db_id, const int32_t tb_id);
  StorageStats getStorageStats(const int32_t db_id, const int32_t tb_id);

  // For testing purposes only
  std::shared_ptr<FileMgr> getSharedFileMgr(const int db_id, const int table_id);

  // For testing purposes only
  void setFileMgr(const int db_id, const int table_id, std::shared_ptr<FileMgr> file_mgr);
  void closeFileMgr(const int32_t db_id,
                    const int32_t tb_id);  // A locked public wrapper for deleteFileMgr,
                                           // for now for unit testing
 protected:
  std::shared_ptr<ForeignStorageInterface> fsi_;

 private:
  bool existsDiffBetweenFileMgrParamsAndFileMgr(
      FileMgr* file_mgr,
      const FileMgrParams& file_mgr_params) const;
  std::string basePath_;       
  size_t num_reader_threads_;  
  int32_t
      epoch_; /* the current epoch (time of last checkpoint) will be used for all
               * tables except of the one for which the value of the epoch has been reset
               * using --start-epoch option at start up to rollback this table's updates.
               */
  const size_t page_size_;           
  const size_t metadata_page_size_;  
  bool dbConvert_;  

  std::map<TablePair, std::shared_ptr<FileMgr>> ownedFileMgrs_;
  std::map<TablePair, AbstractBufferMgr*> allFileMgrs_;
  std::map<TablePair, int32_t> max_rollback_epochs_per_table_;
  std::map<TablePair, StorageStats> lazy_initialized_stats_;

  heavyai::shared_mutex fileMgrs_mutex_;
};

}  // namespace File_Namespace