Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef ADX_COMMON_UTILS_BOUND_QUEUE_H
      12              : #define ADX_COMMON_UTILS_BOUND_QUEUE_H
      13              : #include <condition_variable>
      14              : #include <queue>
      15              : #include <mutex>
      16              : namespace Adx {
      17              : template <typename T>
      18              : class BoundQueue {
      19              : public:
      20           29 :     explicit BoundQueue(uint32_t capacity) : quit_(false), capacity_(capacity) {}
      21           29 :     virtual ~BoundQueue() {}
      22              :     bool TryPush(T& value)
      23              :     {
      24              :         std::lock_guard<std::mutex> lk(mtx_);
      25              :         if (this->IsFull()) {
      26              :             return false;
      27              :         }
      28              : 
      29              :         dataQueue_.push(value);
      30              :         cvPush_.notify_all();
      31              :         return true;
      32              :     }
      33              : 
      34            9 :     bool Push(T& value)
      35              :     {
      36            9 :         std::unique_lock<std::mutex> lk(mtx_);
      37           18 :         cvPop_.wait(lk, [=] { return !this->IsFull() || quit_; });
      38            9 :         dataQueue_.push(value);
      39            9 :         cvPush_.notify_all();
      40            9 :         return true;
      41            9 :     }
      42              : 
      43              :     bool TryPop(T& value)
      44              :     {
      45              :         std::lock_guard<std::mutex> lk(mtx_);
      46              :         if (dataQueue_.empty()) {
      47              :             return false;
      48              :         }
      49              : 
      50              :         value = dataQueue_.front();
      51              :         dataQueue_.pop();
      52              :         cvPop_.notify_all();
      53              :         return true;
      54              :     }
      55              : 
      56            9 :     bool Pop(T& value)
      57              :     {
      58            9 :         std::unique_lock<std::mutex> lk(mtx_);
      59           18 :         cvPush_.wait(lk, [=] { return !this->IsEmpty() || quit_; });
      60            9 :         if (!this->IsEmpty()) {
      61            9 :             value = this->dataQueue_.front();
      62            9 :             this->dataQueue_.pop();
      63            9 :             cvPop_.notify_all();
      64            9 :             return true;
      65              :         }
      66              : 
      67            0 :         return false;
      68            9 :     }
      69              : 
      70           18 :     bool IsEmpty() const { return dataQueue_.empty(); }
      71              : 
      72            9 :     bool IsFull() const { return dataQueue_.size() == capacity_; }
      73              : 
      74              :     void Quit()
      75              :     {
      76              :         std::lock_guard<std::mutex> lk(mtx_);
      77              :         if (!quit_) {
      78              :             quit_ = true;
      79              :             cvPush_.notify_all();
      80              :             cvPop_.notify_all();
      81              :         }
      82              :     }
      83              : 
      84              :     uint32_t Size() { return dataQueue_.size(); }
      85              : 
      86              : private:
      87              :     mutable bool quit_;
      88              :     mutable std::mutex mtx_;
      89              :     std::queue<T> dataQueue_;
      90              :     std::condition_variable cvPop_;
      91              :     std::condition_variable cvPush_;
      92              :     uint32_t capacity_;
      93              : };
      94              : } // namespace Adx
      95              : #endif