/*
 * mem_stress.c — 匿名页内存加压工具 (iOS / HarmonyOS / macOS / Linux 通用)
 *
 * 功能:
 *   按入参大小 mmap 一块匿名私有内存 (MAP_PRIVATE|MAP_ANONYMOUS), 并逐页写入
 *   触页 (touch) 提交物理页, 形成真实内存压力; 之后按参数保持一段时间
 *   (或直到收到 SIGINT/SIGTERM) 再释放。
 *
 *   可选 -k 周期性重触: 每个周期向所有页写入新值, 迫使被 XNU compressor /
 *   Linux zram 压缩回收的页重新解压驻留, 并把干净页重新弄脏, 维持压力。
 *
 * 用法:
 *   mem_stress [-f] [-k sec] [-t sec] [-qv] <size>[K|M|G|T]
 *
 * 构建:
 *   host      : cc -O2 -o mem_stress mem_stress.c
 *   iOS       : xcrun -sdk iphoneos clang -arch arm64 -O2 \
 *               -miphoneos-version-min=12.0 -o mem_stress mem_stress.c
 *   HarmonyOS : <OHOS_NDK>/llvm/bin/clang --target=aarch64-linux-ohos \
 *               --sysroot=<OHOS_NDK>/sysroot -O2 -o mem_stress mem_stress.c
 *
 * 说明:
 *   - 纯 POSIX C, 无第三方依赖; 仅 RSS/footprint 读取按平台分支。
 *   - 单位为二进制: K=1024, M=1024^2, G=1024^3, T=1024^4。
 *   - iOS 上单进程匿名保留有 ~4GiB 上限 (14.8 实测), 超限 mmap 直接失败;
 *     更常见的是压力过大先被 jetsam 以 SIGKILL 终止 —— 属预期行为。
 */

#if !defined(__APPLE__)
#define _POSIX_C_SOURCE 200809L
#endif

#include <errno.h>
#include <inttypes.h>
#include <limits.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <time.h>
#include <unistd.h>

#if defined(__APPLE__)
#include <mach/mach.h>
#endif

#ifndef MAP_ANONYMOUS
#define MAP_ANONYMOUS MAP_ANON
#endif

static volatile sig_atomic_t g_stop = 0;
static size_t g_page = 4096;

/* ---------------- 公共小工具 ---------------- */

static void on_signal(int sig)
{
    (void)sig;
    g_stop = 1;
}

static double now_sec(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
}

static double to_mib(unsigned long long bytes)
{
    return (double)bytes / 1048576.0;
}

/*
 * 解析大小参数: "1024" "100K" "512M" "2G" "1T", 兼容 B/iB 后缀与大小写。
 * 二进制单位 (K=1024)。成功返回 0, 失败返回 -1。
 */
static int parse_size(const char *s, unsigned long long *out)
{
    if (s == NULL || *s == '\0' || *s == '-' || *s == '+')
        return -1;

    errno = 0;
    char *end = NULL;
    unsigned long long v = strtoull(s, &end, 10);
    if (errno != 0 || end == s)
        return -1;

    const char *u = end;
    unsigned long long mult = 1ULL;

    if (*u == 'K' || *u == 'k' || *u == 'M' || *u == 'm' ||
        *u == 'G' || *u == 'g' || *u == 'T' || *u == 't') {
        switch (*u | 0x20) {
        case 'k': mult = 1ULL << 10; break;
        case 'm': mult = 1ULL << 20; break;
        case 'g': mult = 1ULL << 30; break;
        case 't': mult = 1ULL << 40; break;
        }
        u++;
        if (*u == 'i')                 /* KiB / MiB / GiB ... */
            u++;
        if (*u == 'B' || *u == 'b')
            u++;
    } else if (*u == 'B' || *u == 'b') {
        u++;                           /* "100B" = 100 字节 */
    }
    if (*u != '\0')
        return -1;
    if (v == 0)
        return -1;
    if (v > ULLONG_MAX / mult)
        return -1;

    *out = v * mult;
    return 0;
}

/* 当前驻留集 (字节); 取不到返回 0 (调用方按 0 跳过显示) */
static unsigned long long get_rss_bytes(void)
{
#if defined(__APPLE__)
    struct task_vm_info vm;
    mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
    if (task_info(mach_task_self(), TASK_VM_INFO,
                  (task_info_t)&vm, &count) != KERN_SUCCESS)
        return 0;
    return (unsigned long long)vm.resident_size;
#else
    /* Linux / OpenHarmony(musl): /proc/self/statm 第二列 = 驻留页数 */
    FILE *f = fopen("/proc/self/statm", "r");
    if (f == NULL)
        return 0;
    unsigned long long total = 0, resident = 0;
    int n = fscanf(f, "%llu %llu", &total, &resident);
    fclose(f);
    if (n != 2)
        return 0;
    return resident * (unsigned long long)g_page;
#endif
}

#if defined(__APPLE__)
/* iOS jetsam 的判定口径: phys_footprint = 驻留脏页 + 已压缩脏页 + IOKit 映射 */
static unsigned long long get_footprint_bytes(void)
{
    struct task_vm_info vm;
    mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
    if (task_info(mach_task_self(), TASK_VM_INFO,
                  (task_info_t)&vm, &count) != KERN_SUCCESS)
        return 0;
    return (unsigned long long)vm.phys_footprint;
}
#endif

/* 峰值驻留集 (字节); Darwin 的 ru_maxrss 单位是字节, Linux 系是 KiB */
static unsigned long long get_peak_rss_bytes(void)
{
    struct rusage ru;
    if (getrusage(RUSAGE_SELF, &ru) != 0)
        return 0;
#if defined(__APPLE__)
    return (unsigned long long)ru.ru_maxrss;
#else
    return (unsigned long long)ru.ru_maxrss * 1024ULL;
#endif
}

/*
 * 解析秒数参数 (strtod 严格校验, 拒绝垃圾输入)。成功返回 0。
 * 注: 不用 getopt —— BSD/musl 的 getopt 不重排 argv, "size 在前、选项在后"
 * 的写法会解析失败; 手写解析保证各平台行为一致。
 */
static int parse_seconds(const char *s, double *out)
{
    if (s == NULL || *s == '\0')
        return -1;
    errno = 0;
    char *end = NULL;
    double v = strtod(s, &end);
    if (errno != 0 || end == s || *end != '\0' || v < 0.0)
        return -1;
    *out = v;
    return 0;
}

static void usage(FILE *out, const char *prog)
{
    fprintf(out,
        "mem_stress - anonymous memory pressure tool (iOS/HarmonyOS/macOS/Linux)\n"
        "\n"
        "usage: %s [-f] [-k sec] [-t sec] [-qv] <size>\n"
        "\n"
        "  <size>    bytes of anonymous memory to commit\n"
        "            e.g. 268435456 | 100K | 512M | 1500M | 2G | 1T\n"
        "            (binary units: K=1024, M=1024^2, G=1024^3; B/iB suffix ok)\n"
        "  -f        full fill: memset the whole region (default: 1 byte per page)\n"
        "  -k sec    re-touch all pages every <sec> while holding, to defeat\n"
        "            compressor/zram reclaim (0 = off, default off)\n"
        "  -t sec    hold duration after touch, then release\n"
        "            (default: hold until SIGINT/SIGTERM)\n"
        "  -q        quiet (suppress progress lines)\n"
        "  -v        verbose progress\n"
        "\n"
        "examples:\n"
        "  %s 512M -t 30          commit 512M, hold 30s, release\n"
        "  %s 1500M -t 60 -k 5    hold 60s, re-touch every 5s against compression\n",
        prog, prog, prog);
}

int main(int argc, char **argv)
{
    int full_fill = 0, verbose = 0, quiet = 0;
    double keep_interval = 0.0;
    double hold_secs = -1.0;           /* -1 = 一直持有直到信号 */
    const char *size_arg = NULL;

    for (int i = 1; i < argc; i++) {
        const char *a = argv[i];
        if (strcmp(a, "-f") == 0) {
            full_fill = 1;
        } else if (strcmp(a, "-q") == 0) {
            quiet = 1;
        } else if (strcmp(a, "-v") == 0) {
            verbose = 1;
        } else if (strcmp(a, "-h") == 0 || strcmp(a, "--help") == 0) {
            usage(stdout, argv[0]);
            return 0;
        } else if (strcmp(a, "-k") == 0 || strcmp(a, "-t") == 0) {
            double v = 0.0;
            if (i + 1 >= argc || parse_seconds(argv[i + 1], &v) != 0) {
                fprintf(stderr, "error: %s expects a non-negative number of seconds\n", a);
                return 2;
            }
            i++;
            if (a[1] == 'k')
                keep_interval = v;
            else
                hold_secs = v;
        } else if (a[0] == '-' && a[1] != '\0') {
            fprintf(stderr, "error: unknown option '%s'\n", a);
            usage(stderr, argv[0]);
            return 2;
        } else {
            if (size_arg != NULL) {
                fprintf(stderr, "error: multiple sizes given ('%s' and '%s')\n",
                        size_arg, a);
                return 2;
            }
            size_arg = a;
        }
    }
    if (size_arg == NULL) {
        usage(stderr, argv[0]);
        return 2;
    }

    unsigned long long req = 0;
    if (parse_size(size_arg, &req) != 0) {
        fprintf(stderr, "error: bad size '%s' (try 512M / 2G / 65536K)\n",
                size_arg);
        return 2;
    }

    long ps = sysconf(_SC_PAGESIZE);
    if (ps > 0)
        g_page = (size_t)ps;

    if (req > (unsigned long long)(SIZE_MAX - (g_page - 1))) {
        fprintf(stderr, "error: size %llu bytes exceeds this platform's address space\n",
                req);
        return 2;
    }
    size_t size = ((size_t)req + g_page - 1) & ~((size_t)g_page - 1);
    unsigned long long pages = (unsigned long long)(size / g_page);

    setvbuf(stdout, NULL, _IOLBF, 0);  /* 管道/ssh 下进度实时可见 */

    struct sigaction sa;
    memset(&sa, 0, sizeof sa);
    sa.sa_handler = on_signal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;                   /* 不重启: nanosleep 被信号打断立刻返回 */
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);

    double t_start = now_sec();

    /* ---- 申请匿名区 ---- */
    double t0 = now_sec();
    char *mem = mmap(NULL, size, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    double map_s = now_sec() - t0;
    if (mem == MAP_FAILED) {
        fprintf(stderr, "error: mmap %llu bytes failed: %s\n",
                (unsigned long long)size, strerror(errno));
#if defined(__APPLE__)
        fprintf(stderr, "hint: iOS caps a single process's anonymous reservation "
                        "(~4GiB observed on 14.8); try a smaller size.\n");
#endif
        return 1;
    }

    if (!quiet)
        printf("alloc  : %llu bytes (%.2f MiB), %llu page(s) x %zu B, "
               "mmap %.3f ms\n",
               (unsigned long long)size, to_mib(size), pages, g_page,
               map_s * 1e3);

    /* ---- 触页提交物理内存 ---- */
    unsigned long long done = 0;       /* 已触页数 (page-touch 模式) */
    t0 = now_sec();
    if (full_fill) {
        memset(mem, 0xA5, size);       /* 整块填充 (不可中途打断) */
    } else {
        /* 每页首字节写一次即触发缺页提交; 默认每 10% 汇报, -v 每 1% */
        unsigned long long report_every = pages / (verbose ? 100 : 10);
        if (report_every == 0)
            report_every = 1;
        for (size_t off = 0; off < size; off += g_page) {
            if (g_stop)
                break;
            mem[off] = (char)0xA5;
            done++;
            if (!quiet && (done % report_every) == 0)
                printf("touch  : %3llu%% (%8.2f MiB) elapsed %.2f s\n",
                       done * 100ULL / pages,
                       to_mib(done * (unsigned long long)g_page),
                       now_sec() - t0);
        }
    }
    double touch_s = now_sec() - t0;

    if (g_stop && !full_fill) {
        printf("touch  : interrupted at %.2f MiB (%llu/%llu pages)\n",
               to_mib(done * (unsigned long long)g_page), done, pages);
    } else {
        printf("touch  : %s %llu pages (%.2f MiB) in %.3f s",
               full_fill ? "filled" : "touched", pages, to_mib(size), touch_s);
        if (touch_s > 0.0)
            printf(" [%.1f MiB/s, %.1fK pages/s]",
                   to_mib(size) / touch_s,
                   (double)pages / touch_s / 1000.0);
        printf("\n");
    }

    {
        unsigned long long rss = get_rss_bytes();
#if defined(__APPLE__)
        unsigned long long fp = get_footprint_bytes();
        if (rss || fp)
            printf("rss    : resident %.2f MiB, phys_footprint %.2f MiB\n",
                   to_mib(rss), to_mib(fp));
#else
        if (rss)
            printf("rss    : resident %.2f MiB\n", to_mib(rss));
#endif
    }

    /* ---- 保持 (可选周期重触) ---- */
    if (!g_stop) {
        if (hold_secs >= 0.0) {
            if (!quiet)
                printf("hold   : %.3f s%s\n", hold_secs,
                       keep_interval > 0.0 ? ", re-touch on" : "");
        } else if (!quiet) {
            printf("hold   : until SIGINT/SIGTERM (press Ctrl-C to release)\n");
        }

        double hold0 = now_sec();
        double deadline = (hold_secs >= 0.0) ? hold0 + hold_secs : -1.0;
        double next_cycle = (keep_interval > 0.0) ? hold0 + keep_interval : -1.0;
        unsigned cycle = 0;

        while (!g_stop) {
            double now = now_sec();
            if (deadline > 0.0 && now >= deadline)
                break;
            if (next_cycle > 0.0 && now >= next_cycle) {
                cycle++;
                /* 写入与上一周期不同的值: 强制压缩页解压回 RAM、干净页重新变脏 */
                unsigned char pat = (unsigned char)(0xA0 | (cycle & 0x0F));
                if (full_fill) {
                    memset(mem, pat, size);
                } else {
                    for (size_t off = 0; off < size && !g_stop; off += g_page)
                        mem[off] = (char)pat;
                }
                if (!quiet) {
                    printf("cycle  : %u re-touched %.2f MiB", cycle, to_mib(size));
                    unsigned long long r2 = get_rss_bytes();
                    if (r2)
                        printf(", rss %.2f MiB", to_mib(r2));
                    printf("\n");
                }
                next_cycle += keep_interval;
                continue;
            }
            struct timespec ts = { 0, 20 * 1000 * 1000 };  /* 20ms 轮询 */
            nanosleep(&ts, NULL);
        }
    }

    /* ---- 释放 ---- */
    munmap(mem, size);
    printf("release: munmap ok, total %.3f s, peak rss %.2f MiB\n",
           now_sec() - t_start, to_mib(get_peak_rss_bytes()));
    return 0;
}