feat(kmod): reuse mempool

.github/workflows/cppcheck-differential.yaml

@@ -77,7 +77,7 @@ jobs:
       - name: Add kmod to paths
         id: add-kmod
         run: |
-          echo "filtered-full-paths=${{ steps.add-header-paths.outputs.target-paths }} $(realpath kmod/agnocast.c)" >> $GITHUB_OUTPUT
+          echo "filtered-full-paths=${{ steps.add-header-paths.outputs.target-paths }} $(realpath kmod/agnocast.c) $(realpath kmod/memory_allocator.c)" >> $GITHUB_OUTPUT
 
       - name: Run Cppcheck on modified packages
         if: ${{ steps.add-kmod.outputs.filtered-full-paths != '' }}

kmod/Makefile

@@ -1,5 +1,7 @@
 obj-m := agnocast.o
+agnocast-objs := memory_allocator.o
 CFLAGS_agnocast.o := -Wall -Werror
+CFLAGS_memory_allocator.o := -Wall -Werror
 CFLAGS_agnocast_kunit.o := -Wall -Werror -DKUNIT_BUILD
 
 KERNEL_VERSION := $(shell uname -r)
@@ -33,7 +35,7 @@ endif
 
 ifeq ($(CONFIG_KUNIT),y)
   obj-m += agnocast_kunit.o
-  agnocast_kunit-m := agnocast_kunit_main.o agnocast_tmp_copied.o \
+  agnocast_kunit-m := agnocast_kunit_main.o agnocast_tmp_copied.o memory_allocator.o \
 		agnocast_kunit/agnocast_kunit_subscriber_add.o \
 		agnocast_kunit/agnocast_kunit_publisher_add.o \
 		agnocast_kunit/agnocast_kunit_increment_rc.o \
@@ -49,6 +51,7 @@ endif
 
 ifeq ($(CONFIG_GCOV_KERNEL),y)
   CFLAGS_agnocast.o += -fprofile-arcs -ftest-coverage
+  CFLAGS_memory_allocator.o += -fprofile-arcs -ftest-coverage
   CFLAGS_agnocast_tmp_copied.o += -fprofile-arcs -ftest-coverage
   CFLAGS_agnocast_kunit.o += -fprofile-arcs -ftest-coverage
 endif

kmod/agnocast.c

@@ -1,5 +1,7 @@
 #include "agnocast.h"
 
+#include "memory_allocator.h"
+
 #include <linux/device.h>
 #include <linux/hashtable.h>
 #include <linux/kprobes.h>
@@ -1246,9 +1248,6 @@ int take_msg(
 
 int new_shm_addr(const pid_t pid, uint64_t shm_size, union ioctl_new_shm_args * ioctl_ret)
 {
-  // TODO: assume 0x40000000000~ (4398046511104) is allocatable
-  static uint64_t allocatable_addr = 0x40000000000;
-
   if (shm_size % PAGE_SIZE != 0) {
     dev_warn(
       agnocast_device, "shm_size=%llu is not aligned to PAGE_SIZE=%lu. (new_shm_addr)\n", shm_size,
@@ -1261,13 +1260,22 @@ int new_shm_addr(const pid_t pid, uint64_t shm_size, union ioctl_new_shm_args *
     return -EINVAL;
   }
 
+  uint64_t shm_addr = agnocast_assign_memory(pid, shm_size);
+  if (shm_addr == 0) {
+    dev_warn(
+      agnocast_device,
+      "Process (pid=%d) failed to allocate memory (shm_size=%llu). (new_shm_addr)\n", pid,
+      shm_size);
+    return -ENOMEM;
+  }
+
   struct process_info * new_proc_info = kmalloc(sizeof(struct process_info), GFP_KERNEL);
   if (!new_proc_info) {
     dev_warn(agnocast_device, "kmalloc failed. (new_shm_addr)\n");
     return -ENOMEM;
   }
   new_proc_info->pid = pid;
-  new_proc_info->shm_addr = allocatable_addr;
+  new_proc_info->shm_addr = shm_addr;
   new_proc_info->shm_size = shm_size;
   new_proc_info->mapped_num = 0;
   for (int i = 0; i < MAX_MAP_NUM; i++) {
@@ -1278,9 +1286,7 @@ int new_shm_addr(const pid_t pid, uint64_t shm_size, union ioctl_new_shm_args *
   uint32_t hash_val = hash_min(new_proc_info->pid, PROC_INFO_HASH_BITS);
   hash_add(proc_info_htable, &new_proc_info->node, hash_val);
 
-  allocatable_addr += shm_size;
-
-  ioctl_ret->ret_addr = new_proc_info->shm_addr;
+  ioctl_ret->ret_addr = shm_addr;
   return 0;
 }
 
@@ -1741,6 +1747,12 @@ void process_exit_cleanup(const pid_t pid)
 
   if (!agnocast_related) return;
 
+  if (agnocast_free_memory(pid) < 0) {
+    dev_warn(
+      agnocast_device, "Unreachable: pid=%d doesn't have shared memory (process_exit_cleanup)\n",
+      pid);
+  }
+
   struct topic_wrapper * wrapper;
   struct hlist_node * node;
   int bkt;
@@ -1906,6 +1918,8 @@ static int agnocast_init(void)
   ret = agnocast_init_kprobe();
   if (ret < 0) return ret;
 
+  agnocast_init_memory_allocator();
+
   dev_info(agnocast_device, "Agnocast installed!\n");
   return 0;
 }

kmod/agnocast_kunit/agnocast_kunit_new_shm.c

@@ -6,42 +6,44 @@
 
 static pid_t pid = 1000;
 
-static uint64_t get_previously_allocated_shm_addr(struct kunit * test)
-{
-  uint64_t local_pid = pid++;
-  union ioctl_new_shm_args new_shm_args;
-  int ret = new_shm_addr(local_pid, PAGE_SIZE, &new_shm_args);
-  KUNIT_ASSERT_EQ(test, ret, 0);
-  process_exit_cleanup(local_pid);
-  return new_shm_args.ret_addr;
-}
-
 void test_case_new_shm_normal(struct kunit * test)
 {
-  uint64_t previously_allocated_shm_addr = get_previously_allocated_shm_addr(test);
   KUNIT_ASSERT_EQ(test, get_proc_info_htable_size(), 0);
 
   uint64_t local_pid = pid++;
   union ioctl_new_shm_args args;
   int ret = new_shm_addr(local_pid, PAGE_SIZE, &args);
 
   KUNIT_EXPECT_EQ(test, ret, 0);
-  KUNIT_EXPECT_EQ(test, args.ret_addr, previously_allocated_shm_addr + PAGE_SIZE);
   KUNIT_EXPECT_EQ(test, get_proc_info_htable_size(), 1);
   KUNIT_EXPECT_TRUE(test, is_in_proc_info_htable(local_pid));
+
+  process_exit_cleanup(local_pid);
+}
+
+void test_case_new_shm_too_big(struct kunit * test)
+{
+  KUNIT_ASSERT_EQ(test, get_proc_info_htable_size(), 0);
+
+  uint64_t local_pid = pid++;
+  uint64_t shm_size = 8589934592 /* 8GB */ + PAGE_SIZE;
+  union ioctl_new_shm_args args;
+  int ret = new_shm_addr(local_pid, shm_size, &args);
+
+  KUNIT_EXPECT_EQ(test, ret, -ENOMEM);
+  KUNIT_EXPECT_EQ(test, get_proc_info_htable_size(), 0);
+  KUNIT_EXPECT_FALSE(test, is_in_proc_info_htable(local_pid));
 }
 
-void test_case_new_shm_many(struct kunit * test)
+void test_case_new_shm_too_many(struct kunit * test)
 {
-  const int process_num = 1000;
-  uint64_t previously_allocated_shm_addr = get_previously_allocated_shm_addr(test);
+  const int process_num = 10000;
   KUNIT_ASSERT_EQ(test, get_proc_info_htable_size(), 0);
 
   // ================================================
   // Act
 
-  pid_t local_pid_start = pid;
-  for (int i = 0; i < process_num - 1; i++) {
+  for (int i = 0; i < process_num; i++) {
     uint64_t local_pid = pid++;
     union ioctl_new_shm_args args;
     new_shm_addr(local_pid, PAGE_SIZE, &args);
@@ -54,11 +56,13 @@ void test_case_new_shm_many(struct kunit * test)
   // ================================================
   // Assert
 
-  KUNIT_EXPECT_EQ(test, ret, 0);
-  KUNIT_EXPECT_EQ(test, args.ret_addr, previously_allocated_shm_addr + PAGE_SIZE * process_num);
-  KUNIT_EXPECT_EQ(test, get_proc_info_htable_size(), process_num);
+  KUNIT_EXPECT_EQ(test, ret, -ENOMEM);
+  KUNIT_EXPECT_TRUE(test, get_proc_info_htable_size() > 0);
+  KUNIT_EXPECT_TRUE(test, get_proc_info_htable_size() < process_num);
+  KUNIT_EXPECT_FALSE(test, is_in_proc_info_htable(local_pid));
+
   for (int i = 0; i < process_num; i++) {
-    KUNIT_EXPECT_TRUE(test, is_in_proc_info_htable(local_pid_start + i));
+    process_exit_cleanup(pid - process_num + i);
   }
 }
 
@@ -84,4 +88,6 @@ void test_case_new_shm_twice(struct kunit * test)
   KUNIT_EXPECT_EQ(test, ret2, -EINVAL);
   KUNIT_EXPECT_EQ(test, get_proc_info_htable_size(), 1);
   KUNIT_EXPECT_TRUE(test, is_in_proc_info_htable(local_pid));
+
+  process_exit_cleanup(local_pid);
 }

kmod/agnocast_kunit/agnocast_kunit_new_shm.h

@@ -1,11 +1,13 @@
 #pragma once
 #include <kunit/test.h>
 
-#define TEST_CASES_NEW_SHM                                                  \
-  KUNIT_CASE(test_case_new_shm_normal), KUNIT_CASE(test_case_new_shm_many), \
-    KUNIT_CASE(test_case_new_shm_not_aligned), KUNIT_CASE(test_case_new_shm_twice)
+#define TEST_CASES_NEW_SHM                                                             \
+  KUNIT_CASE(test_case_new_shm_normal), KUNIT_CASE(test_case_new_shm_too_big),         \
+    KUNIT_CASE(test_case_new_shm_too_many), KUNIT_CASE(test_case_new_shm_not_aligned), \
+    KUNIT_CASE(test_case_new_shm_twice)
 
 void test_case_new_shm_normal(struct kunit * test);
-void test_case_new_shm_many(struct kunit * test);
+void test_case_new_shm_too_big(struct kunit * test);
+void test_case_new_shm_too_many(struct kunit * test);
 void test_case_new_shm_not_aligned(struct kunit * test);
 void test_case_new_shm_twice(struct kunit * test);

kmod/agnocast_kunit_main.c

@@ -9,6 +9,7 @@
 #include "agnocast_kunit/agnocast_kunit_receive_msg.h"
 #include "agnocast_kunit/agnocast_kunit_subscriber_add.h"
 #include "agnocast_kunit/agnocast_kunit_take_msg.h"
+#include "memory_allocator.h"
 
 #include <kunit/test.h>
 
@@ -55,6 +56,8 @@ static int agnocast_test_suite_init(struct kunit_suite * test_suite)
   ret = agnocast_init_kprobe();
   if (ret < 0) return ret;
 
+  agnocast_init_memory_allocator();
+
   return 0;
 }
 
@@ -64,6 +67,7 @@ static void agnocast_test_suite_exit(struct kunit_suite * test_suite)
   agnocast_exit_kthread();
   agnocast_exit_kprobe();
   agnocast_exit_device();
+  agnocast_exit_memory_allocator();
 }
 
 struct kunit_suite agnocast_test_suite = {

kmod/memory_allocator.c

@@ -0,0 +1,119 @@
+#include "memory_allocator.h"
+
+#include <linux/module.h>
+
+MODULE_LICENSE("Dual BSD/GPL");
+
+#define MEMPOOL_128MB_NUM 1000
+#define MEMPOOL_1GB_NUM 100
+#define MEMPOOL_8GB_NUM 10
+
+static uint64_t mempool_size_128mb = 134217728;  // 128 * 1024 * 1024
+static uint64_t mempool_size_1gb = 1073741824;   // 1 * 1024 * 1024 * 1024
+static uint64_t mempool_size_8gb = 8589934592;   // 8 * 1024 * 1024 * 1024
+
+static pid_t mempool_128mb_pid[MEMPOOL_128MB_NUM];
+static pid_t mempool_1gb_pid[MEMPOOL_1GB_NUM];
+static pid_t mempool_8gb_pid[MEMPOOL_8GB_NUM];
+
+static uint64_t mempool_128mb_addr[MEMPOOL_128MB_NUM];
+static uint64_t mempool_1gb_addr[MEMPOOL_1GB_NUM];
+static uint64_t mempool_8gb_addr[MEMPOOL_8GB_NUM];
+
+void agnocast_init_memory_allocator(void)
+{
+  // TODO(Ryuta Kambe): currently, we assume that the address from 0x40000000000 to 0x8B000000000 is
+  // available
+  uint64_t addr = 0x40000000000;
+
+  for (int i = 0; i < MEMPOOL_128MB_NUM; i++) {
+    mempool_128mb_pid[i] = 0;
+    mempool_128mb_addr[i] = addr;
+    addr += mempool_size_128mb;
+  }
+
+  for (int i = 0; i < MEMPOOL_1GB_NUM; i++) {
+    mempool_1gb_pid[i] = 0;
+    mempool_1gb_addr[i] = addr;
+    addr += mempool_size_1gb;
+  }
+
+  for (int i = 0; i < MEMPOOL_8GB_NUM; i++) {
+    mempool_8gb_pid[i] = 0;
+    mempool_8gb_addr[i] = addr;
+    addr += mempool_size_8gb;
+  }
+}
+
+uint64_t agnocast_assign_memory(const pid_t pid, const uint64_t size)
+{
+  if (size <= mempool_size_128mb) {
+    for (int i = 0; i < MEMPOOL_128MB_NUM; i++) {
+      if (mempool_128mb_pid[i] == 0) {
+        mempool_128mb_pid[i] = pid;
+        return mempool_128mb_addr[i];
+      }
+    }
+  }
+
+  if (size <= mempool_size_1gb) {
+    for (int i = 0; i < MEMPOOL_1GB_NUM; i++) {
+      if (mempool_1gb_pid[i] == 0) {
+        mempool_1gb_pid[i] = pid;
+        return mempool_1gb_addr[i];
+      }
+    }
+  }
+
+  if (size <= mempool_size_8gb) {
+    for (int i = 0; i < MEMPOOL_8GB_NUM; i++) {
+      if (mempool_8gb_pid[i] == 0) {
+        mempool_8gb_pid[i] = pid;
+        return mempool_8gb_addr[i];
+      }
+    }
+  }
+
+  return 0;
+}
+
+int agnocast_free_memory(const pid_t pid)
+{
+  for (int i = 0; i < MEMPOOL_128MB_NUM; i++) {
+    if (mempool_128mb_pid[i] == pid) {
+      mempool_128mb_pid[i] = 0;
+      return 0;
+    }
+  }
+
+  for (int i = 0; i < MEMPOOL_1GB_NUM; i++) {
+    if (mempool_1gb_pid[i] == pid) {
+      mempool_1gb_pid[i] = 0;
+      return 0;
+    }
+  }
+
+  for (int i = 0; i < MEMPOOL_8GB_NUM; i++) {
+    if (mempool_8gb_pid[i] == pid) {
+      mempool_8gb_pid[i] = 0;
+      return 0;
+    }
+  }
+
+  return -1;
+}
+
+void agnocast_exit_memory_allocator(void)
+{
+  for (int i = 0; i < MEMPOOL_128MB_NUM; i++) {
+    mempool_128mb_pid[i] = 0;
+  }
+
+  for (int i = 0; i < MEMPOOL_1GB_NUM; i++) {
+    mempool_1gb_pid[i] = 0;
+  }
+
+  for (int i = 0; i < MEMPOOL_8GB_NUM; i++) {
+    mempool_8gb_pid[i] = 0;
+  }
+}

kmod/memory_allocator.h

@@ -0,0 +1,8 @@
+#pragma once
+
+#include <linux/types.h>
+
+void agnocast_init_memory_allocator(void);
+uint64_t agnocast_assign_memory(const pid_t pid, const uint64_t size);
+int agnocast_free_memory(const pid_t pid);
+void agnocast_exit_memory_allocator(void);