Buffer.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 <mutex>

#include "DataMgr/AbstractBuffer.h"
#include "DataMgr/BufferMgr/BufferSeg.h"

using namespace Data_Namespace;

namespace Buffer_Namespace {

class BufferMgr;

class Buffer : public AbstractBuffer {
  friend class BufferMgr;
  friend class FileMgr;

 public:
  /*
  Buffer(const int8_t * mem, const size_t numPages, const size_t pageSize, const int
  epoch);
  */

  Buffer(BufferMgr* bm,
         BufferList::iterator seg_it,
         const int device_id,
         const size_t page_size = 512,
         const size_t num_bytes = 0);

  ~Buffer() override;

  void read(int8_t* const dst,
            const size_t num_bytes,
            const size_t offset = 0,
            const MemoryLevel dst_buffer_type = CPU_LEVEL,
            const int device_id = -1) override;

  void reserve(const size_t num_bytes) override;
  void write(int8_t* src,
             const size_t num_bytes,
             const size_t offset = 0,
             const MemoryLevel src_buffer_type = CPU_LEVEL,
             const int device_id = -1) override;

  void append(int8_t* src,
              const size_t num_bytes,
              const MemoryLevel src_buffer_type = CPU_LEVEL,
              const int deviceId = -1) override;

  int8_t* getMemoryPtr() override;

  void setMemoryPtr(int8_t* new_ptr) override;
  inline size_t reservedSize() const override { return page_size_ * num_pages_; }

  inline size_t pageCount() const override { return num_pages_; }

  inline size_t pageSize() const override { return page_size_; }

  inline int pin() override {
    std::lock_guard<std::mutex> pin_lock(pin_mutex_);
    return (++pin_count_);
  }

  inline int unPin() override {
    std::lock_guard<std::mutex> pin_lock(pin_mutex_);
    CHECK(pin_count_ > 0);
    return (--pin_count_);
  }

  inline int getPinCount() override {
    std::lock_guard<std::mutex> pin_lock(pin_mutex_);
    return (pin_count_);
  }

  // Added for testing.
  int32_t getSlabNum() const { return seg_it_->slab_num; }

 protected:
  int8_t* mem_;  

 private:
  Buffer(const Buffer&);             // private copy constructor
  Buffer& operator=(const Buffer&);  // private overloaded assignment operator
  virtual void readData(int8_t* const dst,
                        const size_t num_bytes,
                        const size_t offset = 0,
                        const MemoryLevel dst_buffer_type = CPU_LEVEL,
                        const int dst_device_id = -1) = 0;
  virtual void writeData(int8_t* const src,
                         const size_t num_bytes,
                         const size_t offset = 0,
                         const MemoryLevel src_buffer_type = CPU_LEVEL,
                         const int src_device_id = -1) = 0;

  BufferMgr* bm_;
  BufferList::iterator seg_it_;
  size_t page_size_;  
  size_t num_pages_;
  int epoch_;  
  std::vector<bool> page_dirty_flags_;
  int pin_count_;
  std::mutex pin_mutex_;
};

}  // namespace Buffer_Namespace