OutstandingQueueRequests.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 <shared_mutex>
#include <unordered_map>
#include <vector>

#include "ExecutorResourceMgrCommon.h"
#include "Logger/Logger.h"
#include "Shared/BinarySemaphore.h"

namespace ExecutorResourceMgr_Namespace {

class OutstandingQueueRequests {
  using OutstandingQueueRequestsMap =
      std::unordered_map<RequestId, SemaphoreShim_Namespace::BinarySemaphore>;

 public:
  void queue_request_and_wait(const RequestId request_id) {
    auto& wait_semaphore = get_semaphore_for_request(request_id);
    wait_semaphore.try_acquire();
    delete_semaphore_for_request(request_id);
  }

  void queue_request_and_wait_with_timeout(const RequestId request_id,
                                           const size_t max_wait_in_ms) {
    CHECK_GT(max_wait_in_ms, size_t(0));
    auto& wait_semaphore = get_semaphore_for_request(request_id);
    // Binary semaphore returns false if it was not unblocked before the specified timeout
    const bool did_timeout = !(wait_semaphore.try_acquire_for(max_wait_in_ms));
    delete_semaphore_for_request(request_id);
    if (did_timeout) {
      throw QueryTimedOutWaitingInQueue(max_wait_in_ms);
    }
  }

  std::vector<RequestId> get_outstanding_request_ids() {
    std::vector<RequestId> outstanding_request_ids;
    std::shared_lock<std::shared_mutex> requests_read_lock(requests_map_mutex_);
    outstanding_request_ids.reserve(outstanding_requests_map_.size());
    for (const auto& request_entry : outstanding_requests_map_) {
      outstanding_request_ids.emplace_back(request_entry.first);
    }
    return outstanding_request_ids;
  }

  void wake_request_by_id(const RequestId request_id) {
    std::unique_lock<std::shared_mutex> requests_write_lock(requests_map_mutex_);
    const auto request_itr = outstanding_requests_map_.find(request_id);
    if (request_itr == outstanding_requests_map_.end()) {
      outstanding_requests_map_[request_id].set_ready();
    } else {
      request_itr->second.release();
    }
  }

 private:
  SemaphoreShim_Namespace::BinarySemaphore& get_semaphore_for_request(
      const RequestId request_id) {
    std::unique_lock<std::shared_mutex> requests_write_lock(requests_map_mutex_);
    return outstanding_requests_map_[request_id];
  }

  void delete_semaphore_for_request(const RequestId request_id) {
    std::unique_lock<std::shared_mutex> requests_write_lock(requests_map_mutex_);
    CHECK_EQ(outstanding_requests_map_.erase(request_id),
             size_t(1));  // Ensure the erase call returns 1, meaning there was actually
                          // an entry in the map matching this request_id to delete
  }

  OutstandingQueueRequestsMap outstanding_requests_map_;

  mutable std::shared_mutex requests_map_mutex_;
};

}  // namespace ExecutorResourceMgr_Namespace