SmartPDCode/pro/mem/mem.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <sys/sysinfo.h>
#include <sys/types.h>
#include <fcntl.h>
#include <errno.h>
#include <signal.h>

#define MAX_LINE 256
#define DEFAULT_INTERVAL 10
#define DEFAULT_DURATION 0
#define FLUSH_INTERVAL 5  // 每5次写入刷新一次文件

volatile sig_atomic_t stop = 0;

void handle_signal(int sig) {
    stop = 1;
}

// 低内存方式获取系统内存信息
int get_memory_info(unsigned long* total, unsigned long* available, 
                   unsigned long* buffers, unsigned long* cached) {
    FILE* fp = fopen("/proc/meminfo", "r");
    if (!fp) return -1;

    char line[MAX_LINE];
    *total = *available = *buffers = *cached = 0;

    while (fgets(line, sizeof(line), fp)) {
        if (strncmp(line, "MemTotal:", 9) == 0) sscanf(line+9, "%lu", total);
        else if (strncmp(line, "MemAvailable:", 13) == 0) sscanf(line+13, "%lu", available);
        else if (strncmp(line, "Buffers:", 8) == 0) sscanf(line+8, "%lu", buffers);
        else if (strncmp(line, "Cached:", 7) == 0) sscanf(line+7, "%lu", cached);
    }

    fclose(fp);
    return 0;
}

// 低内存方式获取进程内存信息
int get_process_mem(pid_t pid, unsigned long* rss, unsigned long* vsz) {
    char path[MAX_LINE];
    snprintf(path, sizeof(path), "/proc/%d/statm", pid);

    FILE* fp = fopen(path, "r");
    if (!fp) return -1;

    if (fscanf(fp, "%lu %lu", vsz, rss) != 2) {
        fclose(fp);
        return -1;
    }

    fclose(fp);
    
    long page_kb = sysconf(_SC_PAGESIZE) / 1024;
    *vsz *= page_kb;
    *rss *= page_kb;
    
    return 0;
}

// 最小化内存占用的CSV写入
void write_csv_entry(FILE* csv, const char* name, pid_t pid, 
                    unsigned long rss, unsigned long vsz,
                    unsigned long total, unsigned long avail,
                    unsigned long buf, unsigned long cache) {
    time_t now = time(NULL);
    struct tm* tm = localtime(&now);
    
    // 直接格式化输出到文件，减少内存缓冲
    fprintf(csv, "%04d-%02d-%02d %02d:%02d:%02d,%s,%d,%lu,%lu,%lu,%lu,%lu,%lu\n",
            tm->tm_year+1900, tm->tm_mon+1, tm->tm_mday,
            tm->tm_hour, tm->tm_min, tm->tm_sec,
            name, pid, rss, vsz, total, avail, buf, cache);
}

int main(int argc, char** argv) {
    if (argc < 4) {
        printf("Usage: %s <name> <pid> <output.csv> [interval=10] [duration=0]\n", argv[0]);
        return 1;
    }

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

    const char* name = argv[1];
    pid_t pid = atoi(argv[2]);
    const char* outfile = argv[3];
    int interval = argc > 4 ? atoi(argv[4]) : DEFAULT_INTERVAL;
    int duration = argc > 5 ? atoi(argv[5]) : DEFAULT_DURATION;

    if (interval <= 0) interval = DEFAULT_INTERVAL;

    // 以最小缓冲模式打开文件
    FILE* csv = fopen(outfile, "a");
    if (!csv) {
        perror("Failed to open output file");
        return 1;
    }

    // 设置文件缓冲模式为行缓冲
    setvbuf(csv, NULL, _IOLBF, 0);

    // 如果是新文件，写入表头
    if (ftell(csv) == 0) {
        fprintf(csv, "Timestamp,Process,PID,RSS(KB),VSZ(KB),Total(KB),Available(KB),Buffers(KB),Cached(KB)\n");
    }

    time_t start = time(NULL);
    int count = 0;

    printf("Monitoring %s (PID: %d), logging to %s\n", name, pid, outfile);
    printf("Interval: %ds, Duration: %ds\n", interval, duration);

    while (!stop) {
        unsigned long rss, vsz, total, avail, buf, cache;
        
        if (get_process_mem(pid, &rss, &vsz)) {
            fprintf(stderr, "Failed to get process memory\n");
        } else if (get_memory_info(&total, &avail, &buf, &cache)) {
            fprintf(stderr, "Failed to get system memory\n");
        } else {
            write_csv_entry(csv, name, pid, rss, vsz, total, avail, buf, cache);
            count++;
            
            // 定期刷新文件缓冲
            if (count % FLUSH_INTERVAL == 0) {
                fflush(csv);
            }
            
            printf(".");
            fflush(stdout);
        }

        // 检查是否达到总时长
        if (duration > 0 && (time(NULL) - start) >= duration) {
            printf("\nReached specified duration\n");
            break;
        }

        // 更精确的睡眠控制
        time_t next = time(NULL) + interval;
        while (time(NULL) < next && !stop) {
            sleep(1);
        }
    }

    fclose(csv);
    printf("\nMonitoring stopped. %d records written.\n", count);
    return 0;
}