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           26 :     explicit BoundQueue (uint32_t capacity)
      21           26 :         : quit_(false), capacity_(capacity) {}
      22           26 :     virtual ~BoundQueue() {}
      23              :     bool TryPush(T &value)
      24              :     {
      25              :         std::lock_guard<std::mutex> lk(mtx_);
      26              :         if (this->IsFull()) {
      27              :             return false;
      28              :         }
      29              : 
      30              :         dataQueue_.push(value);
      31              :         cvPush_.notify_all();
      32              :         return true;
      33              :     }
      34              : 
      35            9 :     bool Push(T &value)
      36              :     {
      37            9 :         std::unique_lock<std::mutex> lk(mtx_);
      38           18 :         cvPop_.wait(lk, [=] { return !this->IsFull() || quit_;});
      39            9 :         dataQueue_.push(value);
      40            9 :         cvPush_.notify_all();
      41            9 :         return true;
      42            9 :     }
      43              : 
      44              :     bool TryPop(T &value)
      45              :     {
      46              :         std::lock_guard<std::mutex> lk(mtx_);
      47              :         if (dataQueue_.empty()) {
      48              :             return false;
      49              :         }
      50              : 
      51              :         value = dataQueue_.front();
      52              :         dataQueue_.pop();
      53              :         cvPop_.notify_all();
      54              :         return true;
      55              :     }
      56              : 
      57            9 :     bool Pop(T &value)
      58              :     {
      59            9 :         std::unique_lock<std::mutex> lk(mtx_);
      60           18 :         cvPush_.wait(lk, [=] { return !this->IsEmpty() || quit_; });
      61            9 :         if (!this->IsEmpty()) {
      62            9 :             value = this->dataQueue_.front();
      63            9 :             this->dataQueue_.pop();
      64            9 :             cvPop_.notify_all();
      65            9 :             return true;
      66              :         }
      67              : 
      68            0 :         return false;
      69            9 :     }
      70              : 
      71           18 :     bool IsEmpty() const
      72              :     {
      73           18 :         return dataQueue_.empty();
      74              :     }
      75              : 
      76            9 :     bool IsFull() const
      77              :     {
      78            9 :         return dataQueue_.size() == capacity_;
      79              :     }
      80              : 
      81              :     void Quit()
      82              :     {
      83              :         std::lock_guard<std::mutex> lk(mtx_);
      84              :         if (!quit_) {
      85              :             quit_ = true;
      86              :             cvPush_.notify_all();
      87              :             cvPop_.notify_all();
      88              :         }
      89              :     }
      90              : 
      91              :     uint32_t Size()
      92              :     {
      93              :         return dataQueue_.size();
      94              :     }
      95              : private:
      96              :     mutable bool quit_;
      97              :     mutable std::mutex mtx_;
      98              :     std::queue<T> dataQueue_;
      99              :     std::condition_variable cvPop_;
     100              :     std::condition_variable cvPush_;
     101              :     uint32_t capacity_;
     102              : };
     103              : }
     104              : #endif