text: /*
   Copyright 2016-2020 Intel Corporation
   
   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 "common/global/global.hpp"
  #include "sched/queue/queue.hpp"
  
  size_t ccl_sched_queue::get_idx() const {
      return idx;
  }
  
  void ccl_sched_bin::add(ccl_sched* sched) {
      if (ccl::global_data::env().priority_mode != ccl_priority_none) {
          CCL_ASSERT(sched->coll_attr.priority == priority,
                     "unexpected sched priority ",
                     sched->coll_attr.priority,
                     " expected ",
                     priority);
      }
      CCL_ASSERT(sched);
      CCL_ASSERT(!sched->bin);
      sched->bin = this;
      sched->queue = queue;
      sched_list.add(sched);
  }
  
  size_t ccl_sched_bin::erase(size_t idx, size_t& next_idx) {
      ccl_sched* sched = nullptr;
      size_t size = 0;
      {
          std::lock_guard<sched_queue_lock_t> lock(sched_list.elem_guard);
          size = sched_list.elems.size();
          CCL_THROW_IF_NOT(size > 0, "unexpected sched_list size ", size);
          CCL_ASSERT(idx < size);
          sched = sched_list.elems[idx];
          sched->set_in_bin_status(ccl_sched_in_bin_erased);
          sched->bin = nullptr;
          size--;
          std::swap(sched_list.elems[size], sched_list.elems[idx]);
          sched_list.elems.resize(size);
          CCL_ASSERT(sched_list.elems.size() == (size));
          next_idx = idx;
      }
      return size;
  }
  
  ccl_sched_queue::ccl_sched_queue(size_t idx, std::vector<size_t> atl_eps)
          : idx(idx),
            atl_eps(atl_eps) {
      LOG_DEBUG("created sched_queue, idx ",
                idx,
                ", atl_eps count ",
                atl_eps.size(),
                ", atl_eps[0] ",
                atl_eps[0]);
  
      if (ccl::global_data::env().priority_mode != ccl_priority_none) {
          CCL_ASSERT(atl_eps.size() == CCL_PRIORITY_BUCKET_COUNT,
                     "unexpected atl_eps count ",
                     atl_eps.size(),
                     ", expected ",
                     CCL_PRIORITY_BUCKET_COUNT);
      }
      else
          CCL_ASSERT(!atl_eps.empty());
  }
  
  ccl_sched_queue::~ccl_sched_queue() {
      size_t expected_max_priority = 0;
      ccl_sched_bin* expected_cached_max_priority_bin = nullptr;
  
      if (bins.size() >= 1) {
          ccl_sched_bin* bin = &(bins.begin()->second);
          expected_max_priority = bin->priority;
          expected_cached_max_priority_bin = bin;
          if (bins.size() > 1)
              LOG_WARN("unexpected bins size ", bins.size(), ", expected <= 1");
      }
  
      if (max_priority != expected_max_priority)
          LOG_WARN("unexpected max_priority ", max_priority, ", expected ", expected_max_priority);
  
      if (cached_max_priority_bin != expected_cached_max_priority_bin)
          LOG_WARN("unexpected cached_max_priority_bin");
  
      clear();
  }
  
  void ccl_sched_queue::add(ccl_sched* sched) {
      CCL_ASSERT(sched);
      CCL_ASSERT(!sched->bin);
  
      size_t priority = sched->get_priority();
      if (ccl::global_data::env().priority_mode != ccl_priority_none) {
          if (sched->coll_param.ctype == ccl_coll_barrier) {
              priority = max_priority;
              sched->coll_attr.priority = priority;
          }
      }
  
      sched->set_in_bin_status(ccl_sched_in_bin_added);
  
      LOG_DEBUG("add to bin: sched ", sched, ", priority ", priority);
  
      ccl_sched_bin* bin = nullptr;
  
      std::lock_guard<sched_queue_lock_t> lock(bins_guard);
  
      sched_bin_list_t::iterator it = bins.find(priority);
      if (it != bins.end()) {
          bin = &(it->second);
          LOG_DEBUG("found bin ", bin);
          CCL_ASSERT(bin->priority == priority);
          bin->add(sched);
      }
      else {
          size_t atl_ep = 0;
          if (ccl::global_data::env().priority_mode == ccl_priority_none)
              atl_ep = atl_eps[0];
          else {
              size_t ep_idx = (priority / CCL_PRIORITY_BUCKET_SIZE) % CCL_PRIORITY_BUCKET_COUNT;
              atl_ep = atl_eps[ep_idx];
              LOG_DEBUG("priority ", priority, ", ep_idx ", ep_idx);
          }
  
          // in-place construct priority bin with added sched
          auto emplace_result = bins.emplace(std::piecewise_construct,
                                             std::forward_as_tuple(priority),
                                             std::forward_as_tuple(this, atl_ep, priority, sched));
          CCL_ASSERT(emplace_result.second);
          bin = &(emplace_result.first->second);
  
          if (priority >= max_priority) {
              max_priority = priority;
              cached_max_priority_bin = bin;
          }
          LOG_DEBUG("didn't find bin, emplaced new one ", bin, " , max_priority ", max_priority);
      }
  
      CCL_ASSERT(bin);
  }
  
  size_t ccl_sched_queue::erase(ccl_sched_bin* bin, size_t idx) {
      CCL_ASSERT(bin);
      size_t bin_priority = bin->get_priority();
  
      LOG_DEBUG("queue ", this, ", bin ", bin);
      size_t next_idx = 0;
  
      // erase sched and check bin size after
      // no need to lock whole `bins` for single erase
      if (!bin->erase(idx, next_idx)) {
          // 'bin 'looks like empty, we can erase it from 'bins'.
          // double check on bin.empty(), before remove it from whole table
          std::lock_guard<sched_queue_lock_t> lock{ bins_guard };
          {
              // no need to lock 'bin' here, because all adding are under bins_guard protection
              if (bin->sched_list.elems.empty() /* && (bins.size() > 1)*/) {
                  bins.erase(bin_priority);
  
                  // change priority
                  if (bins.empty()) {
                      max_priority = 0;
                      cached_max_priority_bin = nullptr;
                  }
                  else if (bin_priority == max_priority) {
                      max_priority--;
                      sched_bin_list_t::iterator it;
                      while ((it = bins.find(max_priority)) == bins.end()) {
                          max_priority--;
                      }
                      cached_max_priority_bin = &(it->second);
                  }
              } // or do nothing, because somebody added new element in bin while we getting a lock
          }
      }
  
      return next_idx;
  }
  
  ccl_sched_bin* ccl_sched_queue::peek() {
      return cached_max_priority_bin;
  }
  
  std::vector<ccl_sched_bin*> ccl_sched_queue::peek_all() {
      std::lock_guard<sched_queue_lock_t> lock{ bins_guard };
      std::vector<ccl_sched_bin*> result;
      result.reserve(bins.size());
      for (auto& bin : bins) {
          result.emplace_back(&(bin.second));
      }
      return result;
  }
  
  void ccl_sched_queue::clear() {
      cached_max_priority_bin = nullptr;
      bins.clear();
      max_priority = 0;
  }
  