OpenHarmony-v3.2-Release/s

/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * 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.
 */
#include "minddata/dataset/engine/cache/storage_manager.h"

#include <iomanip>

#include "utils/ms_utils.h"
#include "minddata/dataset/util/log_adapter.h"
#include "minddata/dataset/util/path.h"
#include "minddata/dataset/util/random.h"
#include "minddata/dataset/util/services.h"

namespace mindspore {
namespace dataset {
std::string StorageManager::GetBaseName(const std::string &prefix, int32_t file_id) {
  std::ostringstream oss;
  oss << prefix << std::setfill('0') << std::setw(5) << file_id;
  return oss.str();
}

std::string StorageManager::ConstructFileName(const std::string &prefix, int32_t file_id, const std::string &suffix) {
  std::string base_name = GetBaseName(prefix, file_id);
  return (base_name + "." + suffix);
}

Status StorageManager::AddOneContainer(int replaced_container_pos) {
  const std::string kPrefix = "IMG";
  const std::string kSuffix = "LB";
  Path container_name = root_ / ConstructFileName(kPrefix, file_id_, kSuffix);
  std::shared_ptr<StorageContainer> sc;
  RETURN_IF_NOT_OK(StorageContainer::CreateStorageContainer(&sc, container_name.ToString()));
  containers_.push_back(sc);
  file_id_++;
  if (replaced_container_pos >= 0) {
    writable_containers_pool_[replaced_container_pos] = containers_.size() - 1;
  } else {
    writable_containers_pool_.push_back(containers_.size() - 1);
  }
  return Status::OK();
}

Status StorageManager::DoServiceStart() {
  containers_.reserve(kMaxNumContainers);
  writable_containers_pool_.reserve(pool_size_);
  if (root_.IsDirectory()) {
    // create multiple containers and store their index in a pool
    CHECK_FAIL_RETURN_UNEXPECTED(pool_size_ > 0, "Expect positive pool_size_, but got:" + std::to_string(pool_size_));
    for (int i = 0; i < pool_size_; i++) {
      RETURN_IF_NOT_OK(AddOneContainer());
    }
  } else {
    RETURN_STATUS_UNEXPECTED("Not a directory");
  }
  return Status::OK();
}

Status StorageManager::Write(key_type *key, const std::vector<ReadableSlice> &buf) {
  RETURN_UNEXPECTED_IF_NULL(key);
  size_t sz = 0;
  for (auto &v : buf) {
    sz += v.GetSize();
  }
  if (sz == 0) {
    RETURN_STATUS_UNEXPECTED("Unexpected 0 length");
  }
  auto mt = GetRandomDevice();
  std::shared_ptr<StorageContainer> cont;
  key_type out_key;
  value_type out_value;
  bool create_new_container = false;
  int old_container_pos = -1;
  size_t last_num_container = -1;
  do {
    SharedLock lock_s(&rw_lock_);
    size_t num_containers = containers_.size();
    if (create_new_container && (num_containers == last_num_container) && (old_container_pos >= 0)) {
      // Upgrade to exclusive lock.
      lock_s.Upgrade();
      create_new_container = false;
      // Check again if someone has already added a
      // new container after we got the x lock
      if (containers_.size() == num_containers) {
        // Create a new container and replace the full container in the pool with the newly created one
        RETURN_IF_NOT_OK(AddOneContainer(old_container_pos));
      }
      // Refresh how many containers there are.
      num_containers = containers_.size();
      // Downgrade back to shared lock
      lock_s.Downgrade();
    }
    if (num_containers == 0) {
      RETURN_STATUS_UNEXPECTED("num_containers is zero");
    }
    // Pick a random container from the writable container pool to insert.
    std::uniform_int_distribution<int> distribution(0, pool_size_ - 1);
    int pos_in_pool = distribution(mt);
    int cont_index = writable_containers_pool_.at(pos_in_pool);
    cont = containers_.at(cont_index);
    off64_t offset;
    Status rc = cont->Insert(buf, &offset);
    if (rc.StatusCode() == StatusCode::kMDBuddySpaceFull) {
      create_new_container = true;
      old_container_pos = pos_in_pool;
      // Remember how many containers we saw. In the next iteration we will do a comparison to see
      // if someone has already created it.
      last_num_container = num_containers;
    } else if (rc.IsOk()) {
      out_value = std::make_pair(cont_index, std::make_pair(offset, sz));
      RETURN_IF_NOT_OK(index_.insert(out_value, &out_key));
      *key = out_key;
      break;
    } else {
      return rc;
    }
  } while (true);
  return Status::OK();
}

Status StorageManager::Read(StorageManager::key_type key, WritableSlice *dest, size_t *bytesRead) const {
  RETURN_UNEXPECTED_IF_NULL(dest);
  auto r = index_.Search(key);
  if (r.second) {
    auto &it = r.first;
    value_type v = *it;
    int container_inx = v.first;
    off_t offset = v.second.first;
    size_t sz = v.second.second;
    if (dest->GetSize() < sz) {
      std::string errMsg = "Destination buffer too small. Expect at least " + std::to_string(sz) +
                           " but length = " + std::to_string(dest->GetSize());
      RETURN_STATUS_UNEXPECTED(errMsg);
    }
    if (bytesRead != nullptr) {
      *bytesRead = sz;
    }
    auto cont = containers_.at(container_inx);
    RETURN_IF_NOT_OK(cont->Read(dest, offset));
  } else {
    RETURN_STATUS_UNEXPECTED("Key not found");
  }
  return Status::OK();
}

Status StorageManager::DoServiceStop() noexcept {
  Status rc;
  Status rc1;
  for (auto const &p : containers_) {
    // The destructor of StorageContainer is not called automatically until the use
    // count drops to 0. But it is not always the case. We will do it ourselves.
    rc = p.get()->Truncate();
    if (rc.IsError()) {
      rc1 = rc;
    }
  }
  containers_.clear();
  writable_containers_pool_.clear();
  file_id_ = 0;
  return rc1;
}

StorageManager::StorageManager(const Path &root) : root_(root), file_id_(0), index_(), pool_size_(1) {}

StorageManager::StorageManager(const Path &root, int pool_size)
    : root_(root), file_id_(0), index_(), pool_size_(pool_size) {}

StorageManager::~StorageManager() { (void)StorageManager::DoServiceStop(); }

std::ostream &operator<<(std::ostream &os, const StorageManager &s) {
  os << "Dumping all containers ..."
     << "\n";
  for (auto const &p : s.containers_) {
    os << *(p.get());
  }
  return os;
}
}  // namespace dataset
}  // namespace mindspore