cuframes-composer/examples/simple_record.c

/* simple_record — Phase 1 smoke test.
 *
 * Подписывается на один cuframes-источник, кодирует каждый кадр через
 * NVENC и пишет H.264 Annex-B byte stream в файл. Завершение по SIGINT.
 *
 * Цель — доказать end-to-end zero-copy: NV12 frame из VMM-памяти publisher'а
 * напрямую попадает в NVENC без промежуточных копий через CPU.
 *
 * Использование:
 *   simple_record --key cam-parking --out test.h264 [--fps 25] [--bitrate 5000]
 *
 * Полученный файл проверяется через ffmpeg:
 *   ffprobe test.h264          → должно показать h264 stream
 *   ffmpeg -i test.h264 ...    → должен декодироваться
 *
 * Лицензия: LGPL-2.1+
 */

#include "../include/cuframes_composer/nvenc.h"
#include "../include/cuframes_composer/source.h"

#include <cuda.h>

#include <errno.h>
#include <getopt.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

static volatile sig_atomic_t g_stop = 0;

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

/* user-data, передаваемая encoder callback'у */
typedef struct write_ctx {
    FILE *fp;
    uint64_t bytes_written;
    uint64_t frames_encoded;
    uint64_t idr_count;
} write_ctx_t;

/* Encoder callback — пишем H.264 Annex-B bytes как есть. */
static void on_bitstream(const uint8_t *bs, size_t size, int64_t pts_ns,
                          int is_idr, void *user)
{
    (void)pts_ns;
    write_ctx_t *ctx = (write_ctx_t *)user;
    if (fwrite(bs, 1, size, ctx->fp) != size) {
        fprintf(stderr, "[simple_record] fwrite failed: %s\n", strerror(errno));
    } else {
        ctx->bytes_written += size;
        ctx->frames_encoded++;
        if (is_idr) ctx->idr_count++;
    }
}

static const char *cu_err(CUresult r)
{
    const char *s = NULL;
    cuGetErrorString(r, &s);
    return s ? s : "unknown";
}

#define CUCHECK(expr) do { \
    CUresult _r = (expr); \
    if (_r != CUDA_SUCCESS) { \
        fprintf(stderr, "[simple_record] %s failed: %s\n", #expr, cu_err(_r)); \
        return 1; \
    } \
} while (0)

int main(int argc, char **argv)
{
    const char *key = NULL;
    const char *out_path = NULL;
    int fps = 25;
    int bitrate_kbps = 5000;
    int max_seconds = 0;  /* 0 = до SIGINT */

    static struct option opts[] = {
        {"key",     required_argument, 0, 'k'},
        {"out",     required_argument, 0, 'o'},
        {"fps",     required_argument, 0, 'f'},
        {"bitrate", required_argument, 0, 'b'},
        {"seconds", required_argument, 0, 's'},
        {"help",    no_argument,       0, 'h'},
        {0, 0, 0, 0},
    };
    int c;
    while ((c = getopt_long(argc, argv, "k:o:f:b:s:h", opts, NULL)) != -1) {
        switch (c) {
        case 'k': key = optarg; break;
        case 'o': out_path = optarg; break;
        case 'f': fps = atoi(optarg); break;
        case 'b': bitrate_kbps = atoi(optarg); break;
        case 's': max_seconds = atoi(optarg); break;
        case 'h':
        default:
            fprintf(stderr,
                "Использование: %s --key <cuframes-key> --out <file.h264>\n"
                "                [--fps 25] [--bitrate 5000] [--seconds N]\n",
                argv[0]);
            return c == 'h' ? 0 : 1;
        }
    }
    if (!key || !out_path) {
        fprintf(stderr, "[simple_record] требуются --key и --out\n");
        return 1;
    }

    signal(SIGINT,  on_sigint);
    signal(SIGTERM, on_sigint);

    /* 1) CUDA primary context на устройстве 0. cuframes-subscriber и NVENC
     *    оба должны работать в одном контексте. */
    CUCHECK(cuInit(0));
    CUdevice dev;
    CUCHECK(cuDeviceGet(&dev, 0));
    CUcontext ctx;
    CUCHECK(cuDevicePrimaryCtxRetain(&ctx, dev));
    CUCHECK(cuCtxPushCurrent(ctx));

    /* 2) Открыть source. */
    cfc_source_config_t scfg = {
        .key = key,
        .consumer_name = "simple_record",
        .cuda_device = 0,
    };
    cfc_source_t *src = NULL;
    if (cfc_source_create(&scfg, &src) != 0) {
        fprintf(stderr, "[simple_record] cfc_source_create failed\n");
        return 1;
    }

    /* 3) Ждать первый кадр чтобы узнать размер. До 30 секунд. */
    cfc_source_snapshot_t snap = { 0 };
    int waited_ms = 0;
    while (!g_stop && waited_ms < 30000) {
        cfc_source_get_latest(src, &snap);
        if (snap.state == CFC_SOURCE_ACTIVE && snap.width > 0) break;
        struct timespec ts = {.tv_sec = 0, .tv_nsec = 50 * 1000 * 1000};
        nanosleep(&ts, NULL);
        waited_ms += 50;
    }
    if (g_stop) goto cleanup_src;
    if (snap.width <= 0 || snap.height <= 0) {
        fprintf(stderr,
                "[simple_record] не дождался первого кадра за 30с (state=%d)\n",
                snap.state);
        goto cleanup_src;
    }
    fprintf(stderr,
            "[simple_record] первый кадр: %dx%d pitch=%d → создаю encoder\n",
            snap.width, snap.height, snap.pitch_y);

    /* 4) Создать encoder под полученный размер. */
    cfc_encoder_config_t ecfg = {
        .cuda_ctx = ctx,
        .width = snap.width,
        .height = snap.height,
        .fps_num = fps,
        .fps_den = 1,
        .bitrate_kbps = bitrate_kbps,
        .gop_size = fps,           /* IDR раз в секунду — стандарт для RTSP */
        .num_b_frames = 0,         /* low-latency: B-кадры мешают */
        .preset = "ll",
    };
    cfc_encoder_t *enc = NULL;
    if (cfc_encoder_create(&ecfg, &enc) != 0) {
        fprintf(stderr, "[simple_record] cfc_encoder_create failed\n");
        goto cleanup_src;
    }

    /* 5) Открыть выходной файл. */
    write_ctx_t wctx = { 0 };
    wctx.fp = fopen(out_path, "wb");
    if (!wctx.fp) {
        fprintf(stderr, "[simple_record] fopen(%s) failed: %s\n",
                out_path, strerror(errno));
        goto cleanup_enc;
    }

    /* 6) Главный цикл — забираем кадры по seq, кодируем. */
    uint64_t last_seq = 0;
    int64_t start_us;
    struct timespec ts_start;
    clock_gettime(CLOCK_MONOTONIC, &ts_start);
    start_us = (int64_t)ts_start.tv_sec * 1000000 + ts_start.tv_nsec / 1000;

    fprintf(stderr, "[simple_record] начало записи в %s (Ctrl+C для остановки)\n",
            out_path);

    while (!g_stop) {
        cfc_source_get_latest(src, &snap);
        if (snap.state != CFC_SOURCE_ACTIVE) {
            /* Источник не active — короткий sleep и снова. */
            struct timespec ts = {.tv_sec = 0, .tv_nsec = 10 * 1000 * 1000};
            nanosleep(&ts, NULL);
            continue;
        }
        if (snap.seq == last_seq) {
            /* Тот же кадр что и был — ждём новый. Спим четверть кадрового
             * интервала чтобы не крутить CPU впустую. */
            int sleep_ns = 1000000000 / fps / 4;
            struct timespec ts = {.tv_sec = 0, .tv_nsec = sleep_ns};
            nanosleep(&ts, NULL);
            continue;
        }
        last_seq = snap.seq;

        if (cfc_encoder_encode_frame(enc, snap.ptr, snap.pitch_y,
                                      snap.pts_ns, on_bitstream, &wctx) != 0) {
            fprintf(stderr, "[simple_record] encode_frame failed\n");
            break;
        }

        /* Прогресс каждые 100 кадров. */
        if (wctx.frames_encoded > 0 && wctx.frames_encoded % 100 == 0) {
            struct timespec now;
            clock_gettime(CLOCK_MONOTONIC, &now);
            int64_t now_us = (int64_t)now.tv_sec * 1000000 + now.tv_nsec / 1000;
            double elapsed_s = (now_us - start_us) / 1e6;
            fprintf(stderr,
                    "[simple_record] %llu кадров, %llu IDR, %.1f МБ за %.1fс (%.1f fps)\n",
                    (unsigned long long)wctx.frames_encoded,
                    (unsigned long long)wctx.idr_count,
                    wctx.bytes_written / 1048576.0,
                    elapsed_s,
                    wctx.frames_encoded / elapsed_s);
        }

        if (max_seconds > 0) {
            struct timespec now;
            clock_gettime(CLOCK_MONOTONIC, &now);
            int64_t elapsed = ((int64_t)now.tv_sec * 1000000 + now.tv_nsec / 1000) - start_us;
            if (elapsed / 1000000 >= max_seconds) {
                fprintf(stderr, "[simple_record] достигнут лимит %dс\n", max_seconds);
                break;
            }
        }
    }

    fprintf(stderr, "[simple_record] flush encoder\n");
    cfc_encoder_flush(enc, on_bitstream, &wctx);

    fprintf(stderr,
            "[simple_record] итого: %llu кадров, %llu IDR, %.2f МБ\n",
            (unsigned long long)wctx.frames_encoded,
            (unsigned long long)wctx.idr_count,
            wctx.bytes_written / 1048576.0);

    fclose(wctx.fp);

cleanup_enc:
    cfc_encoder_destroy(enc);
cleanup_src:
    cfc_source_destroy(src);
    cuCtxPopCurrent(NULL);
    cuDevicePrimaryCtxRelease(dev);
    return 0;
}