Sets the number of CPUs available for a JVM in a container (JDK < 8u191)

proccount.c

/*
 * Copyright 2020 Andrei Pangin
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Compile: gcc -O2 -fPIC -shared -olibproccount.so proccount.c
// Run:     LD_PRELOAD=/path/to/libproccount.so java <args>

#define _GNU_SOURCE

#include <sys/types.h>
#include <sys/stat.h>
#include <dlfcn.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sched.h>


// Read a number from a text file
static long read_long(const char* filename) {
    long result = -1;

    int fd = open(filename, O_RDONLY);
    if (fd >= 0) {
        char buf[64];
        if (read(fd, buf, sizeof(buf)) > 0) {
            result = atol(buf);
        }
        close(fd);
    }

    return result;
}

// Try to discover the number of available CPUs in a Docker container
static int find_container_cpus() {
    long cfs_quota_us = read_long("/sys/fs/cgroup/cpu/cpu.cfs_quota_us");
    long cfs_period_us = read_long("/sys/fs/cgroup/cpu/cpu.cfs_period_us");
    if (cfs_period_us > 0 && cfs_quota_us >= cfs_period_us) {
        return cfs_quota_us / cfs_period_us;
    }

    long cpu_shares = read_long("/sys/fs/cgroup/cpu/cpu.shares");
    if (cpu_shares > 1024) {
        return cpu_shares / 1024;
    }

    return -1;
}


static int container_cpus = -1;

__attribute__((constructor))
static void init_online_cpus() {
    container_cpus = find_container_cpus();
}


// Fake sched_getaffinity() to return the set of [0..container_cpus-1]
int sched_getaffinity(pid_t pid, size_t cpusetsize, cpu_set_t* mask) {
    static void* real_sched_getaffinity = NULL;

    if (container_cpus > 0) {
        CPU_ZERO_S(cpusetsize, mask);
        int i;
        for (i = 0; i < container_cpus; i++) {
            CPU_SET_S(i, cpusetsize, mask);
        }
        return 0;
    }

    if (real_sched_getaffinity == NULL) {
        real_sched_getaffinity = dlsym(RTLD_NEXT, "sched_getaffinity");
    }
    return ((int (*)(pid_t, size_t, cpu_set_t*))real_sched_getaffinity)(pid, cpusetsize, mask);
}

// Fake sysconf(_SC_NPROCESSORS_ONLN) to return container_cpus
long sysconf(int name) {
    static void* real_sysconf = NULL;

    if (name == _SC_NPROCESSORS_ONLN && container_cpus > 0) {
        return container_cpus;
    }

    if (real_sysconf == NULL) {
        real_sysconf = dlsym(RTLD_NEXT, "sysconf");
    }
    return ((long (*)(int))real_sysconf)(name);
}