GPUSuitabilityEst/code/kernel_module/main.c

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <time.h>
#include <sys/resource.h>
#include <pthread.h>
#include <bpf/libbpf.h>
#include "ebpf.skel.h"
#include "event.h"
#include "data_structure/list_entry.h"
#include "container_worker.h"

#define MAX_CMD_LEN 512

LIST_ENTRY g_ContainerList;

PCONTAINER_WORKER
FindContainer(
    const char *ContainerId)
{
    if (IsListEmpty(&g_ContainerList))
        return NULL;

    PLIST_ENTRY Entry = g_ContainerList.flink;

    while (Entry != &g_ContainerList)
    {

        PCONTAINER_WORKER Worker =
            CONTAINING_RECORD(Entry, CONTAINER_WORKER, ListEntry);

        if (strcmp(Worker->ContainerId, ContainerId) == 0)
            return Worker;

        Entry = Entry->flink;
    }

    return NULL;
}

void InsertContainer(
    const char *ContainerId,
    pid_t ContainerPid)
{
    PCONTAINER_WORKER Worker =
        (PCONTAINER_WORKER)malloc(sizeof(CONTAINER_WORKER));

    if (!Worker)
        return;

    memset(Worker, 0, sizeof(CONTAINER_WORKER));

    strncpy(Worker->ContainerId,
            ContainerId,
            MAX_CONTAINER_ID - 1);

    Worker->ContainerPid = ContainerPid;
    Worker->IsRunning = 1;
    pthread_create(&Worker->MonitorThread, NULL,
                   ContainerWorkerEntry, Worker);

    InsertTailList(&g_ContainerList, &Worker->ListEntry);
}

void RemoveContainer(
    PCONTAINER_WORKER Worker)
{

    Worker->IsRunning = 0;
    pthread_join(Worker->MonitorThread, NULL);
    RemoveEntryList(&Worker->ListEntry);

    free(Worker);
}

void *delayed_free(void *arg)
{
    PCONTAINER_WORKER Worker = (PCONTAINER_WORKER)arg;

    usleep(EXIT_DELAY_MICROSECONDS);

    RemoveContainer(Worker);
    return NULL;
}

void RefreshContainerList(void)
{

    FILE *Pipe = popen(
        "docker inspect -f '{{.Name}} {{.State.Pid}}' $(docker ps -q) 2>/dev/null",
        "r");

    if (!Pipe)
    {
        perror("docker inspect failed");
        return;
    }

    char Buffer[65536] = {0};
    size_t Offset = 0;
    while (fgets(Buffer + Offset, sizeof(Buffer) - Offset, Pipe))
    {
        Offset = strlen(Buffer); // 更新 Offset
        if (Offset >= sizeof(Buffer) - 1)
            break;
    }
    pclose(Pipe);

    if (!IsListEmpty(&g_ContainerList))
    {
        PLIST_ENTRY Entry = g_ContainerList.flink;
        while (Entry != &g_ContainerList)
        {
            PCONTAINER_WORKER Worker =
                CONTAINING_RECORD(Entry, CONTAINER_WORKER, ListEntry);
            Worker->ContainerPid = -1;
            Entry = Entry->flink;
        }
    }

    char *Line = strtok(Buffer, "\n");
    while (Line)
    {
        char ContainerId[MAX_CONTAINER_ID];
        int ContainerPid = -1;

        if (sscanf(Line, "%64s %d", ContainerId, &ContainerPid) != 2)
        {
            Line = strtok(NULL, "\n");
            continue;
        }

        if (ContainerId[0] == '/')
            memmove(ContainerId, ContainerId + 1, strlen(ContainerId));

        if (ContainerPid <= 0 || strncmp(ContainerId, "k8s", 3) == 0)
        {
            Line = strtok(NULL, "\n");
            continue;
        }

        PCONTAINER_WORKER Worker = FindContainer(ContainerId);
        if (Worker)
        {
            Worker->ContainerPid = ContainerPid;
        }
        else
        {
            InsertContainer(ContainerId, ContainerPid);
        }

        Line = strtok(NULL, "\n");
    }

    PLIST_ENTRY Entry = g_ContainerList.flink;
    while (Entry != &g_ContainerList)
    {
        PCONTAINER_WORKER Worker =
            CONTAINING_RECORD(Entry, CONTAINER_WORKER, ListEntry);
        PLIST_ENTRY Next = Entry->flink;

        if (Worker->ContainerPid == -1)
        {
            if(!Worker->Exited){

            pthread_t tid;
            pthread_create(&tid, NULL, delayed_free, Worker);
            pthread_detach(tid);
            }
            //

            Worker->Exited = 1;
        }

        Entry = Next;
    }
}

static volatile sig_atomic_t exiting = 0;

static void handle_signal(int sig)
{
    exiting = 1;
}

static int handle_event(void *ctx, void *data, size_t data_sz)
{
    struct event *e = data;
    pid_t pid = e->pid;

    EbpfOpCode opcode = e->op;

    if (IsListEmpty(&g_ContainerList))
    {
         return 0;
    }

    int print = 0;

    PLIST_ENTRY Entry = g_ContainerList.flink;
    while (Entry != &g_ContainerList)
    {

        PCONTAINER_WORKER Worker =
            CONTAINING_RECORD(Entry, CONTAINER_WORKER, ListEntry);

        do
        {
            if (pid == Worker->ContainerPid)
            {
                print = 1;
                break;
            }

            if (Worker->threads == NULL)
                break;

            PLIST_ENTRY Threads = Worker->threads->flink;

            while (Threads != Worker->threads)
            {
                thread_perf_list *Thread =
                    CONTAINING_RECORD(Threads, thread_perf_list, entry);
                if (pid == Thread->tid)
                {
                    print = 1;
                    break;
                }
                Threads = Threads->flink;
            }
        } while (0);

        if (print)
        {
            if (opcode == OPEN)
                Worker->LibGPUCount++;
            if (opcode == MMAP && strcmp(e->comm,"runc:[2:INIT]"))
            {
                Worker->MMAPSize += (unsigned long long)e->size;
            }
            if (opcode == MUNMAP && strcmp(e->comm,"runc:[2:INIT]"))
            {
                Worker->MMAPSize -= (unsigned long long)e->size;
            }
            break;
        }
        Entry = Entry->flink;
    }
    
    // if (print && e->op == MUNMAP  && strcmp(e->comm,"python"))
    // {

    //     printf(
    //         "pid=%u comm=%s size=%llu path=%s\n",
    //         pid,
    //         e->comm,
    //         (unsigned long long)e->size,
    //         e->path[0] ? e->path : "-");
    // }

    return 0;
}

int main()
{
    InitializeListHead(&g_ContainerList);

    struct ebpf_bpf *skel;
    struct ring_buffer *rb = NULL;
    int err;

    struct rlimit rlim = {
        .rlim_cur = RLIM_INFINITY,
        .rlim_max = RLIM_INFINITY,
    };
    setrlimit(RLIMIT_MEMLOCK, &rlim);

    signal(SIGINT, handle_signal);
    signal(SIGTERM, handle_signal);

    skel = ebpf_bpf__open();
    if (!skel)
    {
        fprintf(stderr, "Failed to open BPF skeleton\n");
        return 1;
    }

    do
    {
        err = ebpf_bpf__load(skel);
        if (err)
        {
            fprintf(stderr, "Failed to load BPF skeleton: %d\n", err);
            break;
        }

        err = ebpf_bpf__attach(skel);
        if (err)
        {
            fprintf(stderr, "Failed to attach BPF programs: %d\n", err);
            break;
        }

        rb = ring_buffer__new(
            bpf_map__fd(skel->maps.events),
            handle_event,
            NULL,
            NULL);
        if (!rb)
        {
            fprintf(stderr, "Failed to create ring buffer\n");
            break;
        }
        setvbuf(stdout, NULL, _IONBF, 0);
        printf("time,cid,tid,total,scalar,scalar_mac,pack_128,pack_256,pack_512,vector_mac,int_all,int_128,int_256,ex_ret_brn,ex_ret_brn_misp,libgpuopen,mmapsize\n");
        while (!exiting)
        {
            RefreshContainerList();
            err = ring_buffer__poll(rb, 10);
            if (err == -EINTR)
                break;
            if (err < 0)
            {
                fprintf(stderr, "ring_buffer__poll error: %d\n", err);
                break;
            }
            usleep(1000000);
        }
    } while (0);

    ring_buffer__free(rb);
    ebpf_bpf__destroy(skel);
    return err < 0 ? -err : 0;

    return 0;
}