mcap/modules/Screen_Record/encoder.jai

// The encoder thread's half: frames of BGRA pixels to H.264 in a fragmented MP4, with FFmpeg
// (modules/ffmpeg). The first encoder in RECORD_ENCODERS that opens at the size wanted is used:
// NVENC, AMD's AMF, Intel's QSV, Windows' Media Foundation, then libx264 on the CPU. NVENC takes the
// BGRA as it is and converts it on the GPU; the others are given NV12 or YUV 4:2:0, converted with
// swscale (BT.709, limited range, as tagged).
//
// With sound, an AAC stream (FFmpeg's own encoder, 160 kb/s, 48 kHz stereo) beside it, written from the
// audio thread: the two streams' packets meet in the file under a lock (mux).
//
// The file is written fragmented, a fragment at each keyframe, so one cut short by a crash still plays
// up to its last fragment. Frames are written only when the screen changes, so a count of frames can
// span any time: a keyframe is made every RECORD_KEYFRAME_SECONDS of time instead, which bounds what a
// crash loses. remux.jai makes the file a regular MP4 once it is done.

RECORD_ENCODERS :: string.["h264_nvenc", "h264_amf", "h264_qsv", "h264_mf", "libx264"];
RECORD_KEYFRAME_SECONDS :: 2;
RECORD_AUDIO_BIT_RATE   :: 160_000;

Encoder :: struct {
  format: *AVFormatContext;
  codec: *AVCodecContext;
  stream: *AVStream;
  frame: *AVFrame;
  packet: *AVPacket;
  sws: *SwsContext;
  name: string;            // the encoder used
  input: AVPixelFormat;    // what it is given: BGR0 as captured, or converted to NV12 or YUV420P
  width, height: s32;
  last_pts: s64 = -1;
  last_key_pts: s64 = -1;  // the last keyframe asked for
  header_written: bool;
  bytes: s64;              // of packets written

  // Sound, when asked for: AAC, written from the audio thread.
  audio: *AVCodecContext;
  audio_stream: *AVStream;
  audio_frame: *AVFrame;
  audio_packet: *AVPacket;
  audio_frame_size: s32;   // samples a frame: AAC's 1024
  mux: Mutex;              // the video and audio threads' packets into the file
  mux_ready: bool;
}

// Opens the file and the first encoder that will take width x height (wanted alone, when given).
// Failures to open an encoder are listed in tried (temp), for the log.
encoder_open :: (e: *Encoder, path: string, width: s32, height: s32, fps: s32, wanted: string, audio := false) -> ok: bool, error: string, tried: string {
  av_log_set_level(AV_LOG_ERROR);
  e.width, e.height = width, height;
  result := avformat_alloc_output_context2(*e.format, null, "mp4", temp_c_string(path));
  if result < 0 || !e.format  return false, av_error("FFmpeg cannot make an MP4", result), "";

  tried: [..] string;
  tried.allocator = temp;
  names: [] string = RECORD_ENCODERS;
  if wanted  names = .[wanted];
  for names {
    ok, why := encoder_try(e, it, width, height, fps);
    if ok  break;
    array_add(*tried, tprint("%: %", it, why));
  }
  tried_text := join(..tried, "; ",, temp);
  if !e.codec {
    encoder_free(e);
    return false, tprint("no H.264 encoder opened (%)", tried_text), tried_text;
  }

  e.stream = avformat_new_stream(e.format, null);
  if !e.stream { encoder_free(e); return false, "FFmpeg cannot add a stream", tried_text; }
  avcodec_parameters_from_context(e.stream.codecpar, e.codec);
  e.stream.time_base = e.codec.time_base;
  if audio {
    audio_ok, audio_error := encoder_open_audio(e);
    if !audio_ok { encoder_free(e); return false, audio_error, tried_text; }
  }
  init(*e.mux, "record mux");
  e.mux_ready = true;

  result = avio_open(*e.format.pb, temp_c_string(path), AVIO_FLAG_WRITE);
  if result < 0 { encoder_free(e); return false, av_error(tprint("FFmpeg cannot write %", path), result), tried_text; }
  options: *AVDictionary;
  av_dict_set(*options, "movflags", "frag_keyframe+empty_moov+default_base_moof", 0);
  result = avformat_write_header(e.format, *options);
  av_dict_free(*options);
  if result < 0 { encoder_free(e); return false, av_error("FFmpeg cannot start the MP4", result), tried_text; }
  e.header_written = true;

  e.packet = av_packet_alloc();
  e.frame = av_frame_alloc();
  if !e.packet || !e.frame { encoder_free(e); return false, "FFmpeg is out of memory", tried_text; }
  if e.input != .BGR0 {
    e.sws = sws_getContext(width, height, .BGR0, width, height, e.input, SWS_POINT, null, null, null);
    if !e.sws { encoder_free(e); return false, "swscale cannot convert to the encoder's format", tried_text; }
    coefficients := cast(*[4] s32) sws_getCoefficients(SWS_CS_ITU709);
    sws_setColorspaceDetails(e.sws, coefficients, 1, coefficients, 0, 0, 1 << 16, 1 << 16);
  }
  return true, "", tried_text;
}

// Encodes one frame; pts in RECORD_TIME_BASE ticks, made later than the last one if it is not.
encoder_write :: (e: *Encoder, pixels: *u8, pts: s64) -> ok: bool, error: string {
  pts_out := max(pts, e.last_pts + 1);
  e.last_pts = pts_out;

  frame := e.frame;
  av_frame_unref(frame);
  frame.format = cast(s32) e.input;
  frame.width  = e.width;
  frame.height = e.height;
  if e.input == .BGR0 {
    // Not reference-counted: avcodec_send_frame copies it, so the buffer goes back to the pool at once.
    frame.data[0] = pixels;
    frame.linesize[0] = e.width * 4;
  } else {
    result := av_frame_get_buffer(frame, 0);
    if result < 0  return false, av_error("FFmpeg cannot make a frame", result);
    source := pixels;
    stride := e.width * 4;
    sws_scale(e.sws, *source, *stride, 0, e.height, frame.data.data, frame.linesize.data);
  }
  frame.pts = pts_out;
  if e.last_key_pts < 0 || pts_out - e.last_key_pts >= RECORD_KEYFRAME_SECONDS * RECORD_TIME_BASE {
    frame.pict_type = .I;   // an IDR, with forced-idr: a new fragment
    e.last_key_pts = pts_out;
  }

  result := avcodec_send_frame(e.codec, frame);
  av_frame_unref(frame);
  if result < 0  return false, av_error(tprint("% cannot encode", e.name), result);
  drained, error := encoder_drain(e);
  return drained, error;
}

// One AAC frame of sound: e.audio_frame_size frames of interleaved stereo float, at pts in samples.
// From the audio thread.
encoder_write_audio :: (e: *Encoder, samples: *float, pts: s64) -> ok: bool, error: string {
  frame := e.audio_frame;
  result := av_frame_make_writable(frame);
  if result < 0  return false, av_error("FFmpeg cannot make a sound frame", result);
  left  := cast(*float) frame.data[0];
  right := cast(*float) frame.data[1];
  for 0..e.audio_frame_size - 1 {
    left[it]  = samples[it * 2];
    right[it] = samples[it * 2 + 1];
  }
  frame.pts = pts;
  result = avcodec_send_frame(e.audio, frame);
  if result < 0  return false, av_error("the AAC encoder", result);
  drained, error := encoder_drain(e, e.audio, e.audio_stream, e.audio_packet);
  return drained, error;
}

// Flushes the encoders, sound first, finishes the file and closes it.
encoder_close :: (e: *Encoder) -> ok: bool, error: string {
  ok, error := true, "";
  if e.audio && e.header_written {
    avcodec_send_frame(e.audio, null);
    ok, error = encoder_drain(e, e.audio, e.audio_stream, e.audio_packet);
  }
  if e.codec && e.header_written {
    avcodec_send_frame(e.codec, null);
    video_ok, video_error := encoder_drain(e);
    if ok  ok, error = video_ok, video_error;
    result := av_write_trailer(e.format);
    if ok && result < 0  ok, error = false, av_error("FFmpeg cannot finish the MP4", result);
  }
  encoder_free(e);
  return ok, error;
}

#scope_file

RECORD_TIME_BASE :: 90000;   // pts ticks a second: MPEG's

encoder_try :: (e: *Encoder, name: string, width: s32, height: s32, fps: s32) -> ok: bool, error: string {
  codec := avcodec_find_encoder_by_name(temp_c_string(name));
  if !codec  return false, "not in this FFmpeg";
  input := encoder_input_format(codec);
  if input == .NONE  return false, "takes no format this can give it";

  c := avcodec_alloc_context3(codec);
  if !c  return false, "out of memory";
  c.width = width;
  c.height = height;
  c.time_base = .{1, RECORD_TIME_BASE};
  c.framerate = .{fps, 1};
  c.pix_fmt = input;
  c.gop_size = fps * 2;
  c.max_b_frames = 0;
  c.sample_aspect_ratio = .{1, 1};
  if input != .BGR0 {   // swscale's conversion; NVENC picks and tags its own for BGRA
    c.colorspace = .BT709;
    c.color_primaries = .BT709;
    c.color_trc = .BT709;
    c.color_range = .MPEG;
  }
  if e.format.oformat.flags & AVFMT_GLOBALHEADER  c.flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

  options: *AVDictionary;
  set :: (key: string, value: string) #expand { av_dict_set(*`options, temp_c_string(key), temp_c_string(value), 0); }
  if name == {
    case "h264_nvenc";  set("preset", "p4"); set("rc", "vbr"); set("cq", "23"); set("forced-idr", "1"); c.bit_rate = 0;
    case "h264_amf";    set("rc", "cqp"); set("qp_i", "22"); set("qp_p", "24"); set("quality", "balanced");
    case "h264_qsv";    set("global_quality", "23"); set("preset", "medium");
    case "h264_mf";     set("rate_control", "quality"); set("quality", "70"); set("scenario", "display_remoting");
    case "libx264";     set("preset", "veryfast"); set("crf", "23"); set("forced-idr", "1");
  }
  result := avcodec_open2(c, codec, *options);
  av_dict_free(*options);
  if result < 0 {
    avcodec_free_context(*c);
    return false, av_error("cannot open", result);
  }
  e.codec = c;
  e.input = input;
  e.name = name;
  return true, "";
}

// BGRA as captured when the encoder takes it, else what swscale converts it to.
encoder_input_format :: (codec: *AVCodec) -> AVPixelFormat {
  if !codec.pix_fmts  return .NONE;
  for wanted: AVPixelFormat.[.BGR0, .BGRA, .NV12, .YUV420P] {
    format := codec.pix_fmts;
    while format.* != .NONE {
      if format.* == wanted  return ifx wanted == .BGRA then .BGR0 else wanted;
      format += 1;
    }
  }
  return .NONE;
}

// The packets an encoder (the video's, unless another is given) has ready, into the file.
encoder_drain :: (e: *Encoder, codec: *AVCodecContext = null, stream: *AVStream = null, packet: *AVPacket = null) -> ok: bool, error: string {
  AVERROR_EAGAIN :: -11;
  if !codec  codec, stream, packet = e.codec, e.stream, e.packet;
  while true {
    result := avcodec_receive_packet(codec, packet);
    if result == AVERROR_EAGAIN || result == cast,no_check(s32) AVERROR_EOF  return true, "";
    if result < 0  return false, av_error(tprint("% failed", ifx codec == e.audio then "the AAC encoder" else e.name), result);
    av_packet_rescale_ts(packet, codec.time_base, stream.time_base);
    packet.stream_index = stream.index;
    lock(*e.mux);
    e.bytes += packet.size;
    result = av_interleaved_write_frame(e.format, packet);   // takes the packet's data
    unlock(*e.mux);
    if result < 0  return false, av_error("FFmpeg cannot write to the MP4", result);
  }
  return true, "";
}

encoder_open_audio :: (e: *Encoder) -> ok: bool, error: string {
  codec := avcodec_find_encoder_by_name("aac");
  if !codec  return false, "this FFmpeg has no AAC encoder";
  c := avcodec_alloc_context3(codec);
  if !c  return false, "out of memory";
  c.sample_fmt = .FLTP;
  c.sample_rate = AUDIO_RATE;
  av_channel_layout_default(*c.ch_layout, AUDIO_CHANNELS);
  c.bit_rate = RECORD_AUDIO_BIT_RATE;
  c.time_base = .{1, AUDIO_RATE};
  if e.format.oformat.flags & AVFMT_GLOBALHEADER  c.flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  result := avcodec_open2(c, codec, null);
  if result < 0 { avcodec_free_context(*c); return false, av_error("the AAC encoder cannot open", result); }
  e.audio = c;
  e.audio_frame_size = ifx c.frame_size > 0 then c.frame_size else 1024;

  e.audio_stream = avformat_new_stream(e.format, null);
  if !e.audio_stream  return false, "FFmpeg cannot add a sound stream";
  avcodec_parameters_from_context(e.audio_stream.codecpar, c);
  e.audio_stream.time_base = c.time_base;

  e.audio_packet = av_packet_alloc();
  e.audio_frame = av_frame_alloc();
  if !e.audio_packet || !e.audio_frame  return false, "FFmpeg is out of memory";
  e.audio_frame.format = cast(s32) AVSampleFormat.FLTP;
  e.audio_frame.nb_samples = e.audio_frame_size;
  e.audio_frame.sample_rate = AUDIO_RATE;
  av_channel_layout_default(*e.audio_frame.ch_layout, AUDIO_CHANNELS);
  result = av_frame_get_buffer(e.audio_frame, 0);
  if result < 0  return false, av_error("FFmpeg cannot make a sound frame", result);
  return true, "";
}

encoder_free :: (e: *Encoder) {
  if e.audio_frame   av_frame_free(*e.audio_frame);
  if e.audio_packet  av_packet_free(*e.audio_packet);
  if e.audio         avcodec_free_context(*e.audio);
  e.audio_stream = null;
  if e.mux_ready { destroy(*e.mux); e.mux_ready = false; }
  if e.sws     sws_freeContext(e.sws);
  if e.frame   av_frame_free(*e.frame);
  if e.packet  av_packet_free(*e.packet);
  if e.codec   avcodec_free_context(*e.codec);
  if e.format {
    if e.format.pb  avio_closep(*e.format.pb);
    avformat_free_context(e.format);
  }
  e.sws, e.format, e.stream = null, null, null;
}

#scope_module

// What FFmpeg says an error code is, after what failed (temp).
av_error :: (what: string, code: s32) -> string {
  buffer: [256] u8;
  av_strerror(code, buffer.data, buffer.count);
  return tprint("% (%)", what, to_string(buffer.data));
}