resource_manager.cc

/*!
 *  Copyright (c) 2020 by Contributors
 * \file resource_manager.cc
 * \brief Manage the resources.
 */

#include "resource_manager.h"

#include <dmlc/logging.h>

#include <utility>

namespace dgl {
namespace runtime {

/*
 * The runtime allocates resources during the computation. Some of the resources
 * cannot be destroyed after the process exits especially when the process
 * doesn't exits normally. We need to keep track of the resources in the system
 * and clean them up properly.
 */
class ResourceManager {
  std::unordered_map<std::string, std::shared_ptr<Resource>> resources;

 public:
  void Add(const std::string &key, std::shared_ptr<Resource> resource) {
    auto it = resources.find(key);
    CHECK(it == resources.end()) << key << " already exists";
    resources.insert(
        std::pair<std::string, std::shared_ptr<Resource>>(key, resource));
  }

  void Erase(const std::string &key) { resources.erase(key); }

  void Cleanup() {
    for (auto it = resources.begin(); it != resources.end(); it++) {
      it->second->Destroy();
    }
    resources.clear();
  }
};

static ResourceManager manager;

void AddResource(const std::string &key, std::shared_ptr<Resource> resource) {
  manager.Add(key, resource);
}

void DeleteResource(const std::string &key) { manager.Erase(key); }

void CleanupResources() { manager.Cleanup(); }

}  // namespace runtime
}  // namespace dgl