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 33 : explicit BoundQueueMemory() : quit_(false), memLimit_(InitMemLimit())
31 : {
32 33 : memUsageV1_.open(MEM_USAGE_V1);
33 33 : memUsageV2_.open(MEM_USAGE_V2);
34 33 : }
35 52 : virtual ~BoundQueueMemory()
36 : {
37 27 : memUsageV1_.close();
38 27 : memUsageV2_.close();
39 79 : }
40 73 : bool Push(T& value)
41 : {
42 73 : std::unique_lock<std::mutex> lk(mtx_);
43 146 : cvPop_.wait(lk, [this] { return !this->IsFullUnlocked() || this->quit_; });
44 73 : if (this->quit_) {
45 1 : return false;
46 : }
47 72 : dataQueue_.push(value);
48 72 : cvPush_.notify_all();
49 72 : return true;
50 73 : }
51 :
52 68 : bool Pop(T& value)
53 : {
54 68 : std::unique_lock<std::mutex> lk(mtx_);
55 200 : cvPush_.wait(lk, [this] { return !this->IsEmptyUnlocked() || this->quit_; });
56 68 : if (!this->IsEmptyUnlocked()) {
57 9 : value = this->dataQueue_.front();
58 9 : this->dataQueue_.pop();
59 9 : cvPop_.notify_all();
60 9 : return true;
61 : }
62 :
63 59 : return false;
64 68 : }
65 :
66 84 : bool IsEmpty() const
67 : {
68 84 : std::lock_guard<std::mutex> lk(mtx_);
69 168 : return IsEmptyUnlocked();
70 84 : }
71 :
72 13 : bool IsFull() const
73 : {
74 13 : std::lock_guard<std::mutex> lk(mtx_);
75 26 : return IsFullUnlocked();
76 13 : }
77 :
78 88 : void Init()
79 : {
80 88 : std::lock_guard<std::mutex> lk(mtx_);
81 88 : quit_ = false;
82 88 : std::queue<T>().swap(dataQueue_);
83 88 : }
84 :
85 91 : void Quit()
86 : {
87 91 : std::lock_guard<std::mutex> lk(mtx_);
88 91 : quit_ = true;
89 91 : cvPush_.notify_all();
90 91 : cvPop_.notify_all();
91 91 : }
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 284 : bool IsEmptyUnlocked() const { return dataQueue_.empty(); }
103 :
104 86 : bool IsFullUnlocked() const
105 : {
106 : struct sysinfo info;
107 86 : const size_t queueSize = 60;
108 86 : int32_t ret = sysinfo(&info);
109 86 : 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 85 : 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 85 : return (info.freeram < (info.totalram * (1 - ADX_QUEUE_FULL_SIZE))) && dataQueue_.size() > queueSize;
123 : }
124 :
125 33 : uint64_t InitMemLimit() const
126 : {
127 33 : std::ifstream memLimitV1(MEM_LIMIT_V1);
128 33 : std::ifstream memLimitV2(MEM_LIMIT_V2);
129 33 : uint64_t ret = ReadMemory(memLimitV1, memLimitV2);
130 33 : memLimitV1.close();
131 33 : memLimitV2.close();
132 33 : return ret;
133 33 : }
134 66 : uint64_t ReadLongLong(std::ifstream& f) const
135 : {
136 66 : if (!f.is_open()) {
137 33 : return 0;
138 : }
139 : uint64_t v;
140 33 : if (f >> v) {
141 33 : f.clear();
142 33 : f.seekg(0);
143 33 : return v;
144 : }
145 0 : return 0;
146 : }
147 :
148 33 : uint64_t ReadMemory(std::ifstream& f1, std::ifstream& f2) const
149 : {
150 33 : uint64_t value = ReadLongLong(f2);
151 33 : if (value == 0) {
152 33 : value = ReadLongLong(f1);
153 : }
154 33 : 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
|