cgroupv2_memlimit.c

#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>

static int write_file(const char *path, const char *s)
{
    int fd = open(path, O_WRONLY);
    if (fd < 0) return -1;
    ssize_t n = write(fd, s, strlen(s));
    close(fd);
    return n == (ssize_t)strlen(s) ? 0 : -1;
}

int main(int argc, char **argv)
{
#if !defined(__linux__)
    fprintf(stderr, "cgroupv2_memlimit is Linux-only");
    return 1;
#else
    if (argc < 3) {
        fprintf(stderr, "usage: %s limit_bytes command [args...]", argv[0]);
        return 1;
    }
    const char *limit = argv[1];
    const char *cg_root = "/sys/fs/cgroup";
    char cg_path[256];
    snprintf(cg_path, sizeof cg_path, "%s/snippet_mem", cg_root);
    mkdir(cg_path, 0755);
    char memmax[512];
    snprintf(memmax, sizeof memmax, "%s/memory.max", cg_path);
    if (write_file(memmax, limit) < 0) {
        perror("memory.max");
        return 1;
    }
    pid_t pid = fork();
    if (pid < 0) {
        perror("fork");
        return 1;
    }
    if (pid == 0) {
        char procs_path[512];
        snprintf(procs_path, sizeof procs_path, "%s/cgroup.procs", cg_path);
        char buf[32];
        snprintf(buf, sizeof buf, "%d", getpid());
        if (write_file(procs_path, buf) < 0) {
            perror("cgroup.procs");
            _exit(1);
        }
        execvp(argv[2], &argv[2]);
        perror("execvp");
        _exit(1);
    } else {
        int status;
        waitpid(pid, &status, 0);
        printf("child exited with status %d", status);
        return 0;
    }
#endif
}