abstract_memory_manager.h

// NIST-developed software is provided by NIST as a public service. You may use, copy and distribute copies of the
// software in any medium, provided that you keep intact this entire notice. You may improve, modify and create
// derivative works of the software or any portion of the software, and you may copy and distribute such modifications
// or works. Modified works should carry a notice stating that you changed the software and should note the date and
// nature of any such change. Please explicitly acknowledge the National Institute of Standards and Technology as the
// source of the software. NIST-developed software is expressly provided "AS IS." NIST MAKES NO WARRANTY OF ANY KIND,
// EXPRESS, IMPLIED, IN FACT OR ARISING BY OPERATION OF LAW, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTY OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT AND DATA ACCURACY. NIST NEITHER REPRESENTS NOR
// WARRANTS THAT THE OPERATION OF THE SOFTWARE WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT ANY DEFECTS WILL BE
// CORRECTED. NIST DOES NOT WARRANT OR MAKE ANY REPRESENTATIONS REGARDING THE USE OF THE SOFTWARE OR THE RESULTS
// THEREOF, INCLUDING BUT NOT LIMITED TO THE CORRECTNESS, ACCURACY, RELIABILITY, OR USEFULNESS OF THE SOFTWARE. You
// are solely responsible for determining the appropriateness of using and distributing the software and you assume
// all risks associated with its use, including but not limited to the risks and costs of program errors, compliance
// with applicable laws, damage to or loss of data, programs or equipment, and the unavailability or interruption of 
// operation. This software is not intended to be used in any situation where a failure could cause risk of injury or
// damage to property. The software developed by NIST employees is not subject to copyright protection within the
// United States.


#ifndef HEDGEHOG_ABSTRACT_MEMORY_MANAGER_H
#define HEDGEHOG_ABSTRACT_MEMORY_MANAGER_H

#include <memory>
#include <mutex>
#include <condition_variable>

#include "../../tools/traits.h"
#include "../../tools/logger.h"
#include "../../tools/nvtx_profiler.h"
#include "../../behavior/memory_manager/pool.h"

namespace hh {
#ifndef DOXYGEN_SHOULD_SKIP_THIS
template<class ManagedMemory>
class MemoryData;
#endif //DOXYGEN_SHOULD_SKIP_THIS

template<class ManagedMemory, class = void>
class AbstractMemoryManager {
 private:
  int deviceId_ = 0; 
  bool initialized_ = false; 
  std::unique_ptr<behavior::Pool<ManagedMemory>> pool_ = {}; 

  std::shared_ptr<NvtxProfiler> profiler_ = nullptr; 

 protected:
  std::mutex memoryManagerMutex_ = {}; 

 public:
  AbstractMemoryManager() = delete;

  explicit AbstractMemoryManager(size_t const &capacity)
      : deviceId_(0) {
    auto tempCapacity = capacity > 0 ? capacity : 1; // Cannot have an empty memory manager
    pool_ = std::make_unique<behavior::Pool<ManagedMemory>>(tempCapacity);
  }

  virtual ~AbstractMemoryManager() = default;

  virtual std::shared_ptr<AbstractMemoryManager<ManagedMemory>> copy()  {
    return std::make_shared<AbstractMemoryManager<ManagedMemory>>(this->capacity());
  }

  void deviceId(int deviceId) { deviceId_ = deviceId; }

  void profiler(const std::shared_ptr<NvtxProfiler> &profiler) { this->profiler_ = profiler; }

  [[nodiscard]] size_t currentSize() {
    std::lock_guard<std::mutex> lk(memoryManagerMutex_);
    return this->pool()->size();
  }

  std::shared_ptr<ManagedMemory> getManagedMemory() {
    std::shared_ptr<ManagedMemory> managedMemory;
    HLOG(4, "StaticMemoryManager memory pool size = " << this->currentSize())
    managedMemory = this->pool()->pop_front();
    HLOG(4,
         "StaticMemoryManager After waiting: received: " << managedMemory << " pSize: " << (int) (this->pool()->size()))
    return managedMemory;
  };

  void recycleMemory(std::shared_ptr<MemoryData<ManagedMemory>> managedMemory) {
    std::lock_guard<std::mutex> lk(memoryManagerMutex_);
    managedMemory->used();
    if (managedMemory->canBeRecycled()) {
      this->profiler_->addReleaseMarker();
      managedMemory->recycle();
      this->pool_->push_back(std::dynamic_pointer_cast<ManagedMemory>(managedMemory));
    }
  };

  virtual void initialize() {
    std::lock_guard<std::mutex> lk(memoryManagerMutex_);
    if (!this->isInitialized()) {
      this->initialized();
      std::for_each(
          this->pool()->begin(), this->pool()->end(),
          [this](std::shared_ptr<ManagedMemory> &emptyShared) {
            emptyShared = std::make_shared<ManagedMemory>();
            emptyShared->memoryManager(this);
          }
      );
      initializeMemoryManager();
    }
  };

  virtual void initializeMemoryManager() {}
 protected:
  [[nodiscard]] std::unique_ptr<behavior::Pool<ManagedMemory>> const &pool() const { return pool_; }

  [[nodiscard]] size_t capacity() const { return this->pool()->capacity(); };

  [[nodiscard]] bool isInitialized() const { return initialized_; }

  [[nodiscard]] std::mutex &memoryManagerMutex() { return memoryManagerMutex_; }

  void initialized() { initialized_ = true; }

  [[nodiscard]] int deviceId() const { return deviceId_; }

};

template<class ManagedMemory>
class AbstractMemoryManager<ManagedMemory,
                            typename std::enable_if_t<!traits::is_managed_memory_v<ManagedMemory>>> {
 public:
  virtual void initialize() {
    std::ostringstream oss;
    oss << "Call Memory manager method: " << __FUNCTION__ << " without managed memory data.";
    HLOG_SELF(0, oss.str())
    throw (std::runtime_error(oss.str()));
  }

  void deviceId(int) {
    std::ostringstream oss;
    oss << "Call Memory manager method: " << __FUNCTION__ << " without managed memory data.";
    HLOG_SELF(0, oss.str())
    throw (std::runtime_error(oss.str()));
  }

  void profiler(const std::shared_ptr<NvtxProfiler> &) {
    std::ostringstream oss;
    oss << "Call Memory manager method: " << __FUNCTION__ << " without managed memory data.";
    HLOG_SELF(0, oss.str())
    throw (std::runtime_error(oss.str()));
  }

  virtual std::shared_ptr<AbstractMemoryManager<ManagedMemory>> copy() {
    std::ostringstream oss;
    oss << "Call Memory manager method: " << __FUNCTION__ << " without managed memory data.";
    HLOG_SELF(0, oss.str())
    throw (std::runtime_error(oss.str()));
  };
};
}
#endif //HEDGEHOG_ABSTRACT_MEMORY_MANAGER_H