mjpeg_encoder: rework output buffer allocation

When encoding a frame, red_worker passes an allocated buffer to
libjpeg where it should encode the frame. When it fails, a new
bigger buffer is allocated and the encoding is restarted from
scratch. However, it's possible to use libjpeg to realloc this
buffer if it gets too small during the encoding process. Make use
of this feature, especially since it will make it easier to encore
one line at a time instead of a full frame in subsequent commits.

3 files changed, 112 insertions, 26 deletions

diff --git a/server/mjpeg_encoder.c b/server/mjpeg_encoder.c
index 30f15efb..ae48da53 100644
--- a/server/mjpeg_encoder.c
+++ b/server/mjpeg_encoder.c
@@ -21,6 +21,7 @@
 
 #include "red_common.h"
 #include "mjpeg_encoder.h"
+#include <jerror.h>
 #include <jpeglib.h>
 
 struct MJpegEncoder {
@@ -74,33 +75,120 @@ size_t mjpeg_encoder_get_frame_stride(MJpegEncoder *encoder)
     return encoder->stride;
 }
 
-static void init_destination(j_compress_ptr cinfo)
+
+/* code from libjpeg 8 to handle compression to a memory buffer
+ *
+ * Copyright (C) 1994-1996, Thomas G. Lane.
+ * Modified 2009 by Guido Vollbeding.
+ * This file is part of the Independent JPEG Group's software.
+ */
+typedef struct {
+  struct jpeg_destination_mgr pub; /* public fields */
+
+  unsigned char ** outbuffer;	/* target buffer */
+  unsigned long * outsize;
+  unsigned char * newbuffer;	/* newly allocated buffer */
+  uint8_t * buffer;		/* start of buffer */
+  size_t bufsize;
+} mem_destination_mgr;
+
+static void init_mem_destination(j_compress_ptr cinfo)
 {
 }
 
-static boolean empty_output_buffer(j_compress_ptr cinfo)
+static boolean empty_mem_output_buffer(j_compress_ptr cinfo)
 {
-    return FALSE;
+  size_t nextsize;
+  uint8_t * nextbuffer;
+  mem_destination_mgr *dest = (mem_destination_mgr *) cinfo->dest;
+
+  /* Try to allocate new buffer with double size */
+  nextsize = dest->bufsize * 2;
+  nextbuffer = malloc(nextsize);
+
+  if (nextbuffer == NULL)
+    ERREXIT1(cinfo, JERR_OUT_OF_MEMORY, 10);
+
+  memcpy(nextbuffer, dest->buffer, dest->bufsize);
+
+  if (dest->newbuffer != NULL)
+    free(dest->newbuffer);
+
+  dest->newbuffer = nextbuffer;
+
+  dest->pub.next_output_byte = nextbuffer + dest->bufsize;
+  dest->pub.free_in_buffer = dest->bufsize;
+
+  dest->buffer = nextbuffer;
+  dest->bufsize = nextsize;
+
+  return TRUE;
+}
+
+static void term_mem_destination(j_compress_ptr cinfo)
+{
+  mem_destination_mgr *dest = (mem_destination_mgr *) cinfo->dest;
+
+  *dest->outbuffer = dest->buffer;
+  *dest->outsize = dest->bufsize - dest->pub.free_in_buffer;
 }
 
-static void term_destination(j_compress_ptr cinfo)
+/*
+ * Prepare for output to a memory buffer.
+ * The caller may supply an own initial buffer with appropriate size.
+ * Otherwise, or when the actual data output exceeds the given size,
+ * the library adapts the buffer size as necessary.
+ * The standard library functions malloc/free are used for allocating
+ * larger memory, so the buffer is available to the application after
+ * finishing compression, and then the application is responsible for
+ * freeing the requested memory.
+ */
+
+static void
+jpeg_mem_dest (j_compress_ptr cinfo,
+	       unsigned char ** outbuffer, unsigned long * outsize)
 {
+  mem_destination_mgr *dest;
+#define OUTPUT_BUF_SIZE  4096	/* choose an efficiently fwrite'able size */
+
+  if (outbuffer == NULL || outsize == NULL)	/* sanity check */
+    ERREXIT(cinfo, JERR_BUFFER_SIZE);
+
+  /* The destination object is made permanent so that multiple JPEG images
+   * can be written to the same buffer without re-executing jpeg_mem_dest.
+   */
+  if (cinfo->dest == NULL) {	/* first time for this JPEG object? */
+    cinfo->dest = spice_malloc(sizeof(mem_destination_mgr));
+  }
+
+  dest = (mem_destination_mgr *) cinfo->dest;
+  dest->pub.init_destination = init_mem_destination;
+  dest->pub.empty_output_buffer = empty_mem_output_buffer;
+  dest->pub.term_destination = term_mem_destination;
+  dest->outbuffer = outbuffer;
+  dest->outsize = outsize;
+  dest->newbuffer = NULL;
+
+  if (*outbuffer == NULL || *outsize == 0) {
+    /* Allocate initial buffer */
+    dest->newbuffer = *outbuffer = malloc(OUTPUT_BUF_SIZE);
+    if (dest->newbuffer == NULL)
+      ERREXIT1(cinfo, JERR_OUT_OF_MEMORY, 10);
+    *outsize = OUTPUT_BUF_SIZE;
+  }
+
+  dest->pub.next_output_byte = dest->buffer = *outbuffer;
+  dest->pub.free_in_buffer = dest->bufsize = *outsize;
 }
+/* end of code from libjpeg */
 
 int mjpeg_encoder_encode_frame(MJpegEncoder *encoder,
-                               uint8_t *buffer, size_t buffer_len)
+                               uint8_t **buffer, size_t *buffer_len)
 {
-    struct jpeg_destination_mgr destmgr;
     uint8_t *frame;
     int n;
 
-    destmgr.next_output_byte = buffer;
-    destmgr.free_in_buffer = buffer_len;
-    destmgr.init_destination = init_destination;
-    destmgr.empty_output_buffer = empty_output_buffer;
-    destmgr.term_destination = term_destination;
-
-    encoder->cinfo.dest = &destmgr;
+    jpeg_mem_dest(&encoder->cinfo, buffer, buffer_len);
 
     encoder->cinfo.image_width      = encoder->width;
     encoder->cinfo.image_height     = encoder->height;
@@ -125,5 +213,5 @@ int mjpeg_encoder_encode_frame(MJpegEncoder *encoder,
     jpeg_finish_compress(&encoder->cinfo);
 
     encoder->first_frame = FALSE;
-    return destmgr.next_output_byte - buffer;
+    return encoder->cinfo.dest->next_output_byte - *buffer;
 }
diff --git a/server/mjpeg_encoder.h b/server/mjpeg_encoder.h
index da1966d7..cd8f6af4 100644
--- a/server/mjpeg_encoder.h
+++ b/server/mjpeg_encoder.h
@@ -29,7 +29,7 @@ void mjpeg_encoder_destroy(MJpegEncoder *encoder);
 uint8_t *mjpeg_encoder_get_frame(MJpegEncoder *encoder);
 size_t mjpeg_encoder_get_frame_stride(MJpegEncoder *encoder);
 int mjpeg_encoder_encode_frame(MJpegEncoder *encoder,
-                               uint8_t *buffer, size_t buffer_len);
+                               uint8_t **buffer, size_t *buffer_len);
 
 
 #endif
diff --git a/server/red_worker.c b/server/red_worker.c
index 07d1d9ff..817fc411 100644
--- a/server/red_worker.c
+++ b/server/red_worker.c
@@ -7426,6 +7426,7 @@ static inline int red_send_stream_data(DisplayChannel *display_channel,
 
     StreamAgent *agent = &display_channel->stream_agents[stream - worker->streams_buf];
     uint64_t time_now = red_now();
+    size_t outbuf_size;
     if (time_now - agent->last_send_time < (1000 * 1000 * 1000) / agent->fps) {
         agent->frames--;
         return TRUE;
@@ -7440,19 +7441,16 @@ static inline int red_send_stream_data(DisplayChannel *display_channel,
         return FALSE;
     }
 
-    while ((n = mjpeg_encoder_encode_frame(stream->mjpeg_encoder,
-                                           display_channel->send_data.stream_outbuf,
-                                           display_channel->send_data.stream_outbuf_size)) == 0) {
-        uint8_t *new_buf;
-        size_t new_size;
 
-        new_size = display_channel->send_data.stream_outbuf_size * 2;
-        new_buf = spice_malloc(new_size);
-
-        free(display_channel->send_data.stream_outbuf);
-        display_channel->send_data.stream_outbuf = new_buf;
-        display_channel->send_data.stream_outbuf_size = new_size;
+    outbuf_size = display_channel->send_data.stream_outbuf_size;
+    n = mjpeg_encoder_encode_frame(stream->mjpeg_encoder,
+                                   &display_channel->send_data.stream_outbuf,
+                                   &outbuf_size);
+    if (n == 0) {
+        red_printf("failed to encode frame, out of memory?");
+        return FALSE;
     }
+    display_channel->send_data.stream_outbuf_size = outbuf_size;
 
     red_channel_init_send_data(channel, SPICE_MSG_DISPLAY_STREAM_DATA, NULL);