Add helper function to set rlimit for tests

Add a helper function to check if the current device have sufficient
MEMLOCK rlimit for bpf related unit tests. If not, try to set it to a
proper size.

Bug: 119279144
Bug: 129246448
Test: libbpf_android_test
Change-Id: I5390f4d3b21436abff69a661d1c6e6a6749542ed
Merged-In: I5390f4d3b21436abff69a661d1c6e6a6749542ed
(Cherry-picked from commit 0a1a9a11d96d937e7563659f92edf04fafa766c9)

3 files changed, 19 insertions, 0 deletions

diff --git a/libbpf_android/BpfMapTest.cpp b/libbpf_android/BpfMapTest.cpp
index cfa9d88..895dda6 100644
--- a/libbpf_android/BpfMapTest.cpp
+++ b/libbpf_android/BpfMapTest.cpp
@@ -57,11 +57,13 @@ class BpfMapTest : public testing::Test {
     void SetUp() {
         SKIP_IF_BPF_NOT_SUPPORTED;
 
+        EXPECT_EQ(0, setrlimitForTest());
         if (!access(PINNED_MAP_PATH, R_OK)) {
             EXPECT_EQ(0, remove(PINNED_MAP_PATH));
         }
         mMapFd = createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE,
                            BPF_F_NO_PREALLOC);
+        EXPECT_LE(0, mMapFd);
     }
 
     void TearDown() {
diff --git a/libbpf_android/BpfUtils.cpp b/libbpf_android/BpfUtils.cpp
index 109715b..081317b 100644
--- a/libbpf_android/BpfUtils.cpp
+++ b/libbpf_android/BpfUtils.cpp
@@ -27,6 +27,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
+#include <sys/resource.h>
 #include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/utsname.h>
@@ -237,6 +238,19 @@ int synchronizeKernelRCU() {
     return 0;
 }
 
+int setrlimitForTest() {
+    // Set the memory rlimit for the test process if the default MEMLOCK rlimit is not enough.
+    struct rlimit limit = {
+            .rlim_cur = TEST_LIMIT,
+            .rlim_max = TEST_LIMIT,
+    };
+    int res = setrlimit(RLIMIT_MEMLOCK, &limit);
+    if (res) {
+        ALOGE("Failed to set the default MEMLOCK rlimit: %s", strerror(errno));
+    }
+    return res;
+}
+
 std::string BpfLevelToString(BpfLevel bpfLevel) {
     switch (bpfLevel) {
         case BpfLevel::NONE:      return "NONE_SUPPORT";
diff --git a/libbpf_android/include/bpf/BpfUtils.h b/libbpf_android/include/bpf/BpfUtils.h
index 77a18fa..93a13a4 100644
--- a/libbpf_android/include/bpf/BpfUtils.h
+++ b/libbpf_android/include/bpf/BpfUtils.h
@@ -61,6 +61,8 @@
 
 #define MAP_CMD_SIZE 16
 
+#define TEST_LIMIT 8388608
+
 namespace android {
 namespace bpf {
 
@@ -157,6 +159,7 @@ int bpfFdGet(const char* pathname, uint32_t flags);
 int attachProgram(bpf_attach_type type, uint32_t prog_fd, uint32_t cg_fd);
 int detachProgram(bpf_attach_type type, uint32_t cg_fd);
 uint64_t getSocketCookie(int sockFd);
+int setrlimitForTest();
 std::string BpfLevelToString(BpfLevel BpfLevel);
 BpfLevel getBpfSupportLevel();
 int parseProgramsFromFile(const char* path, BpfProgInfo* programs, size_t size,