mjpeg_encoder: configure mjpeg quality and frame rate according to a given bit rate

Previously, the mjpeg quality was always 70. The frame rate was
tuned according to the frames' congestion in the pipe.
This patch sets the quality and frame rate according to
a given bit rate and the size of the first encoded frames.

The following patches will introduce an adaptive video streaming, in which
the bit rate, the quality, and the frame rate, change in response to
different parameters.

Patches that make red_worker adopt this feature will also follow.

3 files changed, 312 insertions, 7 deletions

diff --git a/server/mjpeg_encoder.c b/server/mjpeg_encoder.c
index b812ba05..00b721db 100644
--- a/server/mjpeg_encoder.c
+++ b/server/mjpeg_encoder.c
@@ -24,27 +24,92 @@
 #include <jerror.h>
 #include <jpeglib.h>
 
+#define MJPEG_MAX_FPS 25
+#define MJPEG_MIN_FPS 1
+
+#define MJPEG_QUALITY_SAMPLE_NUM 7
+static const int mjpeg_quality_samples[MJPEG_QUALITY_SAMPLE_NUM] = {20, 30, 40, 50, 60, 70, 80};
+
+#define MJPEG_LEGACY_STATIC_QUALITY_ID 5 // jpeg quality 70
+
+#define MJPEG_IMPROVE_QUALITY_FPS_STRICT_TH 10
+#define MJPEG_IMPROVE_QUALITY_FPS_PERMISSIVE_TH 5
+
+typedef struct MJpegEncoderQualityEval {
+    uint64_t encoded_size_by_quality[MJPEG_QUALITY_SAMPLE_NUM];
+    /* lower limit for the current evaluation round */
+    int min_quality_id;
+    int min_quality_fps; // min fps for the given quality
+    /* upper limit for the current evaluation round */
+    int max_quality_id;
+    int max_quality_fps; // max fps for the given quality
+    /* tracking the best sampled fps so far */
+    int max_sampled_fps;
+    int max_sampled_fps_quality_id;
+} MJpegEncoderQualityEval;
+
+/*
+ * Adjusting the stream jpeg quality and frame rate (fps):
+ * When during_quality_eval=TRUE, we compress different frames with different
+ * jpeg quality. By considering (1) the resulting compression ratio, and (2) the available
+ * bit rate, we evaulate the max frame frequency for the stream with the given quality,
+ * and we choose the highest quality that will allow a reasonable frame rate.
+ * during_quality_eval is set for new streams and can also be set any time we want
+ * to re-evaluate the stream parameters (e.g., when the bit rate and/or
+ * compressed frame size significantly change).
+ */
+typedef struct MJpegEncoderRateControl {
+    int during_quality_eval;
+    MJpegEncoderQualityEval quality_eval_data;
+
+    uint64_t byte_rate;
+    int quality_id;
+    uint32_t fps;
+
+    uint64_t last_enc_size;
+} MJpegEncoderRateControl;
+
 struct MJpegEncoder {
     uint8_t *row;
     uint32_t row_size;
     int first_frame;
-    int quality;
 
     struct jpeg_compress_struct cinfo;
     struct jpeg_error_mgr jerr;
 
     unsigned int bytes_per_pixel; /* bytes per pixel of the input buffer */
     void (*pixel_converter)(uint8_t *src, uint8_t *dest);
+
+    int rate_control_is_active;
+    MJpegEncoderRateControl rate_control;
+    MJpegEncoderRateControlCbs cbs;
+    void *cbs_opaque;
 };
 
-MJpegEncoder *mjpeg_encoder_new(void)
+static inline void mjpeg_encoder_reset_quality(MJpegEncoder *encoder, int quality_id, uint32_t fps);
+static uint32_t get_max_fps(uint64_t frame_size, uint64_t bytes_per_sec);
+
+MJpegEncoder *mjpeg_encoder_new(int bit_rate_control, uint64_t starting_bit_rate,
+                                MJpegEncoderRateControlCbs *cbs, void *opaque)
 {
     MJpegEncoder *enc;
 
+    spice_assert(!bit_rate_control || (cbs && cbs->get_roundtrip_ms && cbs->get_source_fps));
+
     enc = spice_new0(MJpegEncoder, 1);
 
     enc->first_frame = TRUE;
-    enc->quality = 70;
+    enc->rate_control_is_active = bit_rate_control;
+    enc->rate_control.byte_rate = starting_bit_rate / 8;
+    if (bit_rate_control) {
+        enc->cbs = *cbs;
+        enc->cbs_opaque = opaque;
+        mjpeg_encoder_reset_quality(enc, MJPEG_QUALITY_SAMPLE_NUM / 2, 5);
+        enc->rate_control.during_quality_eval = TRUE;
+    } else {
+        mjpeg_encoder_reset_quality(enc, MJPEG_LEGACY_STATIC_QUALITY_ID, MJPEG_MAX_FPS);
+    }
+
     enc->cinfo.err = jpeg_std_error(&enc->jerr);
     jpeg_create_compress(&enc->cinfo);
 
@@ -191,10 +256,214 @@ spice_jpeg_mem_dest(j_compress_ptr cinfo,
 }
 /* end of code from libjpeg */
 
+static inline uint32_t mjpeg_encoder_get_latency(MJpegEncoder *encoder)
+{
+    return encoder->cbs.get_roundtrip_ms ?
+        encoder->cbs.get_roundtrip_ms(encoder->cbs_opaque) / 2 : 0;
+}
+
+static uint32_t get_max_fps(uint64_t frame_size, uint64_t bytes_per_sec)
+{
+    double fps;
+    double send_time_ms;
+
+    if (!bytes_per_sec) {
+        return 0;
+    }
+    send_time_ms = frame_size * 1000.0 / bytes_per_sec;
+    fps = send_time_ms ? 1000 / send_time_ms : MJPEG_MAX_FPS;
+    return fps;
+}
+
+static inline void mjpeg_encoder_reset_quality(MJpegEncoder *encoder, int quality_id, uint32_t fps)
+{
+    MJpegEncoderRateControl *rate_control = &encoder->rate_control;
+
+    rate_control->during_quality_eval = FALSE;
+
+    if (rate_control->quality_id != quality_id) {
+        rate_control->last_enc_size = 0;
+    }
+    rate_control->quality_id = quality_id;
+    memset(&rate_control->quality_eval_data, 0, sizeof(MJpegEncoderQualityEval));
+    rate_control->quality_eval_data.max_quality_id = MJPEG_QUALITY_SAMPLE_NUM - 1;
+    rate_control->quality_eval_data.max_quality_fps = MJPEG_MAX_FPS;
+    rate_control->fps = MAX(MJPEG_MIN_FPS, fps);
+    rate_control->fps = MIN(MJPEG_MAX_FPS, rate_control->fps);
+}
+
+#define QUALITY_WAS_EVALUATED(encoder, quality) \
+    ((encoder)->rate_control.quality_eval_data.encoded_size_by_quality[(quality)] != 0)
+
+/*
+ * Adjust the stream's jpeg quality and frame rate.
+ * We evaluate the compression ratio of different jpeg qualities;
+ * We compress successive frames with different qualities,
+ * and then we estimate the stream frame rate according to the currently
+ * evaluated jpeg quality and available bit rate.
+ *
+ * During quality evaluation, mjpeg_encoder_eval_quality is called before a new
+ * frame is encoded. mjpeg_encoder_eval_quality examines the encoding size of
+ * the previously encoded frame, and determines whether to continue evaluation
+ * (and chnages the quality for the frame that is going to be encoded),
+ * or stop evaluation (and sets the quality and frame rate for the stream).
+ * When qualities are scanned, we assume monotonicity of compression ratio
+ * as a function of jpeg quality. When we reach a quality with too small, or
+ * big enough compression ratio, we stop the evaluation and set the stream parameters.
+*/
+static inline void mjpeg_encoder_eval_quality(MJpegEncoder *encoder)
+{
+    MJpegEncoderRateControl *rate_control;
+    MJpegEncoderQualityEval *quality_eval;
+    uint32_t fps, src_fps;
+    uint64_t enc_size;
+    uint32_t final_quality_id;
+    uint32_t final_fps;
+    uint64_t final_quality_enc_size;
+
+    rate_control = &encoder->rate_control;
+    quality_eval = &rate_control->quality_eval_data;
+
+    spice_assert(rate_control->during_quality_eval);
+
+    /* retrieving the encoded size of the last encoded frame */
+    enc_size = quality_eval->encoded_size_by_quality[rate_control->quality_id];
+    if (enc_size == 0) {
+        spice_debug("size info missing");
+        return;
+    }
+
+    src_fps = encoder->cbs.get_source_fps(encoder->cbs_opaque);
+
+    fps = get_max_fps(enc_size, rate_control->byte_rate);
+    spice_debug("mjpeg %p: jpeg %d: %.2f (KB) fps %d src-fps %u",
+                encoder,
+                mjpeg_quality_samples[rate_control->quality_id],
+                enc_size / 1024.0,
+                fps,
+                src_fps);
+
+    if (fps > quality_eval->max_sampled_fps ||
+        ((fps == quality_eval->max_sampled_fps || fps >= src_fps) &&
+         rate_control->quality_id > quality_eval->max_sampled_fps_quality_id)) {
+        quality_eval->max_sampled_fps = fps;
+        quality_eval->max_sampled_fps_quality_id = rate_control->quality_id;
+    }
+
+    /*
+     * Choosing whether to evaluate another quality, or to complete evaluation
+     * and set the stream parameters according to one of the qualities that
+     * were already sampled.
+     */
+
+    if (rate_control->quality_id > MJPEG_QUALITY_SAMPLE_NUM / 2 &&
+        fps < MJPEG_IMPROVE_QUALITY_FPS_STRICT_TH &&
+        fps < src_fps) {
+        /*
+         * When the jpeg quality is bigger than the median quality, prefer a reasonable
+         * frame rate over improving the quality
+         */
+        spice_debug("fps < %d && (fps < src_fps), quality %d",
+                MJPEG_IMPROVE_QUALITY_FPS_STRICT_TH,
+                mjpeg_quality_samples[rate_control->quality_id]);
+        if (QUALITY_WAS_EVALUATED(encoder, rate_control->quality_id - 1)) {
+            /* the next worse quality was already evaluated and it passed the frame
+             * rate thresholds (we know that, because we continued evaluating a better
+             * quality) */
+            rate_control->quality_id--;
+            goto complete_sample;
+        } else {
+            /* evaluate the next worse quality */
+            rate_control->quality_id--;
+        }
+    } else if ((fps > MJPEG_IMPROVE_QUALITY_FPS_PERMISSIVE_TH &&
+                fps >= 0.66 * quality_eval->min_quality_fps) || fps >= src_fps) {
+        /* When the jpeg quality is worse than the median one (see first condition), we allow a less
+           strict threshold for fps, in order to improve the jpeg quality */
+        if (rate_control->quality_id + 1 == MJPEG_QUALITY_SAMPLE_NUM ||
+            rate_control->quality_id >= quality_eval->max_quality_id ||
+            QUALITY_WAS_EVALUATED(encoder, rate_control->quality_id + 1)) {
+            /* best quality has been reached, or the next (better) quality was
+             * already evaluated and didn't pass the fps thresholds */
+            goto complete_sample;
+        } else {
+            if (rate_control->quality_id == MJPEG_QUALITY_SAMPLE_NUM / 2 &&
+                fps < MJPEG_IMPROVE_QUALITY_FPS_STRICT_TH &&
+                fps < src_fps) {
+                goto complete_sample;
+            }
+            /* evaluate the next quality as well*/
+            rate_control->quality_id++;
+        }
+    } else { // very small frame rate, try to improve by downgrading the quality
+        if (rate_control->quality_id == 0 ||
+            rate_control->quality_id <= quality_eval->min_quality_id) {
+            goto complete_sample;
+        } else if (QUALITY_WAS_EVALUATED(encoder, rate_control->quality_id - 1)) {
+            rate_control->quality_id--;
+            goto complete_sample;
+        } else {
+            /* evaluate the next worse quality */
+            rate_control->quality_id--;
+        }
+    }
+    return;
+
+complete_sample:
+    if (quality_eval->max_sampled_fps != 0) {
+        /* covering a case were monotonicity was violated and we sampled
+           a better jepg quality, with better frame rate. */
+        final_quality_id = MAX(rate_control->quality_id,
+                               quality_eval->max_sampled_fps_quality_id);
+    } else {
+        final_quality_id = rate_control->quality_id;
+    }
+    final_quality_enc_size = quality_eval->encoded_size_by_quality[final_quality_id];
+    final_fps = get_max_fps(final_quality_enc_size,
+                            rate_control->byte_rate);
+
+    if (final_quality_id == quality_eval->min_quality_id) {
+        final_fps = MAX(final_fps, quality_eval->min_quality_fps);
+    }
+    if (final_quality_id == quality_eval->max_quality_id) {
+        final_fps = MIN(final_fps, quality_eval->max_quality_fps);
+    }
+    mjpeg_encoder_reset_quality(encoder, final_quality_id, final_fps);
+
+    spice_debug("MJpeg quality sample end %p: quality %d fps %d",
+                encoder, mjpeg_quality_samples[rate_control->quality_id], rate_control->fps);
+}
+
+static void mjpeg_encoder_adjust_params_to_bit_rate(MJpegEncoder *encoder)
+{
+    MJpegEncoderRateControl *rate_control;
+
+    if (!encoder->rate_control_is_active) {
+        return;
+    }
+
+    rate_control = &encoder->rate_control;
+
+    if (!rate_control->last_enc_size) {
+        spice_debug("missing sample size");
+        return;
+    }
+
+    if (rate_control->during_quality_eval) {
+        MJpegEncoderQualityEval *quality_eval = &rate_control->quality_eval_data;
+        quality_eval->encoded_size_by_quality[rate_control->quality_id] = rate_control->last_enc_size;
+        mjpeg_encoder_eval_quality(encoder);
+    }
+}
+
 int mjpeg_encoder_start_frame(MJpegEncoder *encoder, SpiceBitmapFmt format,
                               int width, int height,
                               uint8_t **dest, size_t *dest_len)
 {
+    uint32_t quality;
+
+    mjpeg_encoder_adjust_params_to_bit_rate(encoder);
+
     encoder->cinfo.in_color_space   = JCS_RGB;
     encoder->cinfo.input_components = 3;
     encoder->pixel_converter = NULL;
@@ -245,7 +514,8 @@ int mjpeg_encoder_start_frame(MJpegEncoder *encoder, SpiceBitmapFmt format,
     encoder->cinfo.image_height     = height;
     jpeg_set_defaults(&encoder->cinfo);
     encoder->cinfo.dct_method       = JDCT_IFAST;
-    jpeg_set_quality(&encoder->cinfo, encoder->quality, TRUE);
+    quality = mjpeg_quality_samples[encoder->rate_control.quality_id];
+    jpeg_set_quality(&encoder->cinfo, quality, TRUE);
     jpeg_start_compress(&encoder->cinfo, encoder->first_frame);
 
     return TRUE;
@@ -271,6 +541,7 @@ int mjpeg_encoder_encode_scanline(MJpegEncoder *encoder, uint8_t *src_pixels,
     }
     if (scanlines_written == 0) { /* Not enough space */
         jpeg_abort_compress(&encoder->cinfo);
+        encoder->rate_control.last_enc_size = 0;
         return 0;
     }
 
@@ -284,5 +555,15 @@ size_t mjpeg_encoder_end_frame(MJpegEncoder *encoder)
     jpeg_finish_compress(&encoder->cinfo);
 
     encoder->first_frame = FALSE;
-    return dest->pub.next_output_byte - dest->buffer;
+    encoder->rate_control.last_enc_size = dest->pub.next_output_byte - dest->buffer;
+
+    return encoder->rate_control.last_enc_size;
+}
+
+uint32_t mjpeg_encoder_get_fps(MJpegEncoder *encoder)
+{
+    if (!encoder->rate_control_is_active) {
+        spice_warning("bit rate control is not active");
+    }
+    return encoder->rate_control.fps;
 }
diff --git a/server/mjpeg_encoder.h b/server/mjpeg_encoder.h
index b9a2ed7a..902dcbe7 100644
--- a/server/mjpeg_encoder.h
+++ b/server/mjpeg_encoder.h
@@ -23,7 +23,21 @@
 
 typedef struct MJpegEncoder MJpegEncoder;
 
-MJpegEncoder *mjpeg_encoder_new(void);
+/*
+ * Callbacks required for controling and adjusting
+ * the stream bit rate:
+ * get_roundtrip_ms: roundtrip time in milliseconds
+ * get_source_fps: the input frame rate (#frames per second), i.e.,
+ * the rate of frames arriving from the guest to spice-server,
+ * before any drops.
+ */
+typedef struct MJpegEncoderRateControlCbs {
+    uint32_t (*get_roundtrip_ms)(void *opaque);
+    uint32_t (*get_source_fps)(void *opaque);
+} MJpegEncoderRateControlCbs;
+
+MJpegEncoder *mjpeg_encoder_new(int bit_rate_control, uint64_t starting_bit_rate,
+                                MJpegEncoderRateControlCbs *cbs, void *opaque);
 void mjpeg_encoder_destroy(MJpegEncoder *encoder);
 
 uint8_t mjpeg_encoder_get_bytes_per_pixel(MJpegEncoder *encoder);
@@ -39,5 +53,15 @@ int mjpeg_encoder_encode_scanline(MJpegEncoder *encoder, uint8_t *src_pixels,
                                   size_t image_width);
 size_t mjpeg_encoder_end_frame(MJpegEncoder *encoder);
 
+/*
+ * bit rate control
+ */
+
+/*
+ * The recommended output frame rate (per second) for the
+ * current available bit rate.
+ */
+uint32_t mjpeg_encoder_get_fps(MJpegEncoder *encoder);
+
 
 #endif
diff --git a/server/red_worker.c b/server/red_worker.c
index 4c7cca71..23d08a82 100644
--- a/server/red_worker.c
+++ b/server/red_worker.c
@@ -2887,7 +2887,7 @@ static void red_display_create_stream(DisplayChannelClient *dcc, Stream *stream)
     agent->drops = 0;
     agent->fps = MAX_FPS;
     reset_rate(dcc, agent);
-    agent->mjpeg_encoder = mjpeg_encoder_new();
+    agent->mjpeg_encoder = mjpeg_encoder_new(FALSE, 0, NULL, NULL);
     red_channel_client_pipe_add(&dcc->common.base, &agent->create_item);
 }