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