git.sesse.net Git - nageru/blob - x264_encoder.h

   1 // A wrapper around x264, to encode video in higher quality than Quick Sync
   2 // can give us. We maintain a queue of uncompressed Y'CbCr frames (of 50 frames,
   3 // so a little under 100 MB at 720p), then have a separate thread pull out
   4 // those threads as fast as we can to give it to x264 for encoding.
   5 //
   6 // The encoding threads are niced down because mixing is more important than
   7 // encoding; if we lose frames in mixing, we'll lose frames to disk _and_
   8 // to the stream, as where if we lose frames in encoding, we'll lose frames
   9 // to the stream only, so the latter is strictly better. More importantly,
  10 // this allows speedcontrol to do its thing without disturbing the mixer.
  11
  12 #ifndef _X264ENCODE_H
  13 #define _X264ENCODE_H 1
  14
  15 #include <sched.h>
  16 #include <stdint.h>
  17 #include <x264.h>
  18 #include <atomic>
  19 #include <chrono>
  20 #include <condition_variable>
  21 #include <memory>
  22 #include <mutex>
  23 #include <queue>
  24 #include <string>
  25 #include <thread>
  26 #include <unordered_map>
  27 #include <vector>
  28
  29 extern "C" {
  30 #include <libavformat/avformat.h>
  31 }
  32
  33 #include <movit/image_format.h>
  34
  35 #include "defs.h"
  36 #include "print_latency.h"
  37 #include "x264_dynamic.h"
  38
  39 class Mux;
  40 class X264SpeedControl;
  41
  42 class X264Encoder {
  43 public:
  44         X264Encoder(AVOutputFormat *oformat);  // Does not take ownership.
  45
  46         // Called after the last frame. Will block; once this returns,
  47         // the last data is flushed.
  48         ~X264Encoder();
  49
  50         // Must be called before first frame. Does not take ownership.
  51         void add_mux(Mux *mux) { muxes.push_back(mux); }
  52
  53         // <data> is taken to be raw NV12 data of WIDTHxHEIGHT resolution.
  54         // Does not block.
  55         void add_frame(int64_t pts, int64_t duration, movit::YCbCrLumaCoefficients ycbcr_coefficients, const uint8_t *data, const ReceivedTimestamps &received_ts);
  56
  57         std::string get_global_headers() const {
  58                 while (!x264_init_done) {
  59                         sched_yield();
  60                 }
  61                 return global_headers;
  62         }
  63
  64         void change_bitrate(unsigned rate_kbit) {
  65                 new_bitrate_kbit = rate_kbit;
  66         }
  67
  68 private:
  69         struct QueuedFrame {
  70                 int64_t pts, duration;
  71                 movit::YCbCrLumaCoefficients ycbcr_coefficients;
  72                 uint8_t *data;
  73                 ReceivedTimestamps received_ts;
  74         };
  75         void encoder_thread_func();
  76         void init_x264();
  77         void encode_frame(QueuedFrame qf);
  78
  79         // One big memory chunk of all 50 (or whatever) frames, allocated in
  80         // the constructor. All data functions just use pointers into this
  81         // pool.
  82         std::unique_ptr<uint8_t[]> frame_pool;
  83
  84         std::vector<Mux *> muxes;
  85         bool wants_global_headers;
  86
  87         std::string global_headers;
  88         std::string buffered_sei;  // Will be output before first frame, if any.
  89
  90         std::thread encoder_thread;
  91         std::atomic<bool> x264_init_done{false};
  92         std::atomic<bool> should_quit{false};
  93         X264Dynamic dyn;
  94         x264_t *x264;
  95         std::unique_ptr<X264SpeedControl> speed_control;
  96
  97         std::function<void(x264_param_t *)> bitrate_override_func;
  98
  99         std::atomic<unsigned> new_bitrate_kbit{0};  // 0 for no change.
 100
 101         // Protects everything below it.
 102         std::mutex mu;
 103
 104         // Frames that are not being encoded or waiting to be encoded,
 105         // so that add_frame() can use new ones.
 106         std::queue<uint8_t *> free_frames;
 107
 108         // Frames that are waiting to be encoded (ie., add_frame() has been
 109         // called, but they are not picked up for encoding yet).
 110         std::queue<QueuedFrame> queued_frames;
 111
 112         // Whenever the state of <queued_frames> changes.
 113         std::condition_variable queued_frames_nonempty;
 114
 115         // Key is the pts of the frame.
 116         std::unordered_map<int64_t, ReceivedTimestamps> frames_being_encoded;
 117
 118         // Metrics.
 119         std::atomic<int64_t> metric_x264_queued_frames{0};
 120         std::atomic<int64_t> metric_x264_max_queued_frames{X264_QUEUE_LENGTH};
 121         std::atomic<int64_t> metric_x264_dropped_frames{0};
 122         std::atomic<int64_t> metric_x264_output_frames_i{0};
 123         std::atomic<int64_t> metric_x264_output_frames_p{0};
 124         std::atomic<int64_t> metric_x264_output_frames_b{0};
 125
 126         static constexpr size_t crf_buckets = 50;
 127         std::atomic<int64_t> metric_x264_crf[crf_buckets]{{0}};
 128         std::atomic<double> metric_x264_crf_sum{0.0};
 129 };
 130
 131 #endif  // !defined(_X264ENCODE_H)