Skip to content

Commit 92f2acc

Browse files
committed
[fix] add cgroup memory control from contributor.
1 parent 0514873 commit 92f2acc

File tree

2 files changed

+22
-0
lines changed

2 files changed

+22
-0
lines changed

Diff for: cli/database/create.py

+3
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ def memory_statistics():
3434

3535
# 获取总内存大小(以字节为单位)
3636
total_memory = memory.total
37+
pod_memory_limit = get_pod_memory_limit()
38+
if pod_memory_limit != 0:
39+
total_memory = pod_memory_limit
3740

3841
# 格式化内存大小
3942
size_units = ["B", "KB", "MB", "GB", "TB"]

Diff for: cli/extractor/extractor.py

+19
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,9 @@ def jar_extractor_cmd(extractor_path, source_root, database, options):
203203
if not jvm_opts:
204204
mem = psutil.virtual_memory()
205205
total_memory = mem.total
206+
pod_memory_limit = get_pod_memory_limit()
207+
if pod_memory_limit != 0:
208+
total_memory = pod_memory_limit
206209
total_memory_gb = round(total_memory / (1024 ** 3))
207210
total_memory_gb = min(total_memory_gb, 32) # limit to 32G
208211
xmx = max(total_memory_gb - 1, 6)
@@ -233,3 +236,19 @@ def extractor_run(language, source_root, database, timeout, options):
233236
logging.error("Failed to obtain the %s extractor", language)
234237
return -1
235238

239+
def get_pod_memory_limit():
240+
# cgroup 文件系统路径
241+
memory_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
242+
memory_limit = 0
243+
try:
244+
with open(memory_limit_path, 'r') as f:
245+
memory_limit = int(f.read().strip())
246+
except FileNotFoundError:
247+
pass
248+
except PermissionError:
249+
logging.error("Permission denied when accessing cgroup files.")
250+
except IOError as e:
251+
logging.error(f"IO error occurred when accessing cgroup files: {e}")
252+
except Exception as e:
253+
logging.error(f"An unexpected error occurred: {e}")
254+
return memory_limit

0 commit comments

Comments
 (0)