1 files changed, 286 insertions, 5 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;
 }