xref: /aosp_15_r20/external/bcc/tests/python/test_rlimit.py (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1#!/usr/bin/env python3
2#
3# USAGE: test_usdt.py
4#
5# Copyright 2018 Facebook, Inc
6# Licensed under the Apache License, Version 2.0 (the "License")
7
8from __future__ import print_function
9from bcc import BPF
10from unittest import main, skipUnless, TestCase
11from utils import kernel_version_ge
12import os, resource
13
14@skipUnless(not kernel_version_ge(5, 11), "Since d5299b67dd59 \"bpf: Memcg-based memory accounting for bpf maps\""\
15                                          ",map mem has been counted against memcg, not rlimit")
16class TestRlimitMemlock(TestCase):
17    def testRlimitMemlock(self):
18        text = b"""
19BPF_HASH(unused, u64, u64, 65536);
20int test() { return 0; }
21"""
22        # save the original memlock limits
23        memlock_limit = resource.getrlimit(resource.RLIMIT_MEMLOCK)
24
25        # set a small RLIMIT_MEMLOCK limit
26        resource.setrlimit(resource.RLIMIT_MEMLOCK, (4096, 4096))
27
28        # below will fail
29        failed = 0
30        try:
31            b = BPF(text=text, allow_rlimit=False)
32        except:
33            failed = 1
34        self.assertEqual(failed, 1)
35
36        # below should succeed
37        b = BPF(text=text, allow_rlimit=True)
38
39        # reset to the original memlock limits
40        resource.setrlimit(resource.RLIMIT_MEMLOCK, memlock_limit)
41
42if __name__ == "__main__":
43    main()
44