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_MEMORY_H
12 : #define ADX_COMMON_UTILS_BOUND_QUEUE_MEMORY_H
13 : #include <condition_variable>
14 : #include <fstream>
15 : #include <queue>
16 : #include <mutex>
17 : #include <sys/sysinfo.h>
18 : #include "hdc_log.h"
19 :
20 : namespace Adx {
21 : constexpr float ADX_QUEUE_FULL_SIZE = 0.85f;
22 : constexpr const char* MEM_USAGE_V1 = "/sys/fs/cgroup/memory/memory.usage_in_bytes";
23 : constexpr const char* MEM_USAGE_V2 = "/sys/fs/cgroup/memory.current";
24 : constexpr const char* MEM_LIMIT_V1 = "/sys/fs/cgroup/memory/memory.limit_in_bytes";
25 : constexpr const char* MEM_LIMIT_V2 = "/sys/fs/cgroup/memory.max";
26 :
27 : template <typename T>
28 : class BoundQueueMemory {
29 : public:
30 34 : explicit BoundQueueMemory() : quit_(false), memLimit_(InitMemLimit())
31 : {
32 34 : memUsageV1_.open(MEM_USAGE_V1);
33 34 : memUsageV2_.open(MEM_USAGE_V2);
34 34 : }
35 54 : virtual ~BoundQueueMemory()
36 : {
37 28 : memUsageV1_.close();
38 28 : memUsageV2_.close();
39 82 : }
40 74 : bool Push(T& value)
41 : {
42 74 : std::unique_lock<std::mutex> lk(mtx_);
43 148 : cvPop_.wait(lk, [this] { return !this->IsFullUnlocked() || this->quit_; });
44 74 : if (this->quit_) {
45 1 : return false;
46 : }
47 73 : dataQueue_.push(value);
48 73 : cvPush_.notify_all();
49 73 : return true;
50 74 : }
51 :
52 71 : bool Pop(T& value)
53 : {
54 71 : std::unique_lock<std::mutex> lk(mtx_);
55 208 : cvPush_.wait(lk, [this] { return !this->IsEmptyUnlocked() || this->quit_; });
56 71 : if (!this->IsEmptyUnlocked()) {
57 10 : value = this->dataQueue_.front();
58 10 : this->dataQueue_.pop();
59 10 : cvPop_.notify_all();
60 10 : return true;
61 : }
62 :
63 61 : return false;
64 71 : }
65 :
66 86 : bool IsEmpty() const
67 : {
68 86 : std::lock_guard<std::mutex> lk(mtx_);
69 172 : return IsEmptyUnlocked();
70 86 : }
71 :
72 14 : bool IsFull() const
73 : {
74 14 : std::lock_guard<std::mutex> lk(mtx_);
75 28 : return IsFullUnlocked();
76 14 : }
77 :
78 89 : void Init()
79 : {
80 89 : std::lock_guard<std::mutex> lk(mtx_);
81 89 : quit_ = false;
82 89 : std::queue<T>().swap(dataQueue_);
83 89 : }
84 :
85 93 : void Quit()
86 : {
87 93 : std::lock_guard<std::mutex> lk(mtx_);
88 93 : quit_ = true;
89 93 : cvPush_.notify_all();
90 93 : cvPop_.notify_all();
91 93 : }
92 :
93 4 : uint32_t Size() const
94 : {
95 4 : std::lock_guard<std::mutex> lk(mtx_);
96 8 : return dataQueue_.size();
97 4 : }
98 :
99 62 : void SetPath(std::string path) { path_ = path; }
100 :
101 : private:
102 294 : bool IsEmptyUnlocked() const { return dataQueue_.empty(); }
103 :
104 88 : bool IsFullUnlocked() const
105 : {
106 : struct sysinfo info;
107 88 : const size_t queueSize = 60;
108 88 : int32_t ret = sysinfo(&info);
109 88 : if (ret != EN_OK) {
110 1 : IDE_LOGW("Can not get memory, sysinfo return: %d", ret);
111 1 : return dataQueue_.size() >= queueSize;
112 : }
113 : // if in memory-limited container, make sure that memory usage in container is less than 85%
114 87 : if (memLimit_ > 0 && memLimit_ < info.totalram) {
115 0 : uint64_t memUsage = ReadMemory(memUsageV1_, memUsageV2_);
116 0 : if (memUsage == 0) {
117 0 : IDE_LOGW("Can not read memory usage from cgroup.");
118 0 : return dataQueue_.size() >= queueSize;
119 : }
120 0 : return (memUsage > ADX_QUEUE_FULL_SIZE * memLimit_) && dataQueue_.size() > queueSize;
121 : }
122 87 : return (info.freeram < (info.totalram * (1 - ADX_QUEUE_FULL_SIZE))) && dataQueue_.size() > queueSize;
123 : }
124 :
125 34 : uint64_t InitMemLimit() const
126 : {
127 34 : std::ifstream memLimitV1(MEM_LIMIT_V1);
128 34 : std::ifstream memLimitV2(MEM_LIMIT_V2);
129 34 : uint64_t ret = ReadMemory(memLimitV1, memLimitV2);
130 34 : memLimitV1.close();
131 34 : memLimitV2.close();
132 34 : return ret;
133 34 : }
134 68 : uint64_t ReadLongLong(std::ifstream& f) const
135 : {
136 68 : if (!f.is_open()) {
137 34 : return 0;
138 : }
139 : uint64_t v;
140 34 : if (f >> v) {
141 34 : f.clear();
142 34 : f.seekg(0);
143 34 : return v;
144 : }
145 0 : return 0;
146 : }
147 :
148 34 : uint64_t ReadMemory(std::ifstream& f1, std::ifstream& f2) const
149 : {
150 34 : uint64_t value = ReadLongLong(f2);
151 34 : if (value == 0) {
152 34 : value = ReadLongLong(f1);
153 : }
154 34 : return value;
155 : }
156 : mutable bool quit_;
157 : mutable std::mutex mtx_;
158 : std::queue<T> dataQueue_;
159 : std::condition_variable cvPop_;
160 : std::condition_variable cvPush_;
161 : std::string path_;
162 : mutable std::ifstream memUsageV1_;
163 : mutable std::ifstream memUsageV2_;
164 : uint64_t memLimit_;
165 : };
166 : } // namespace Adx
167 : #endif
|
