tight: Don't block 1 worker while encoding tiles

include/tight.h

@@ -16,6 +16,9 @@
 
 #pragma once
 
+#include "rfb-proto.h"
+#include "vec.h"
+
 #include <unistd.h>
 #include <stdint.h>
 #include <zlib.h>
@@ -25,6 +28,8 @@ struct tight_tile;
 struct pixman_region16;
 struct aml_work;
 
+typedef void (*tight_done_fn)(struct vec* frame, void*);
+
 enum tight_quality {
 	TIGHT_QUALITY_UNSPEC = 0,
 	TIGHT_QUALITY_LOSSLESS,
@@ -44,17 +49,17 @@ struct tight_encoder {
 	z_stream zs[4];
 	struct aml_work* zs_worker[4];
 
-	const struct rfb_pixel_format* dfmt;
+	struct rfb_pixel_format dfmt;
-	const struct rfb_pixel_format* sfmt;
+	struct rfb_pixel_format sfmt;
 	const struct nvnc_fb* fb;
 
 	uint32_t n_rects;
 	uint32_t n_jobs;
 
-	struct vec* dst;
+	struct vec dst;
 
-	pthread_mutex_t wait_mutex;
+	tight_done_fn on_frame_done;
-	pthread_cond_t wait_cond;
+	void* userdata;
 };
 
 int tight_encoder_init(struct tight_encoder* self, uint32_t width,
@@ -64,9 +69,10 @@ void tight_encoder_destroy(struct tight_encoder* self);
 int tight_encoder_resize(struct tight_encoder* self, uint32_t width,
 		uint32_t height);
 
-int tight_encode_frame(struct tight_encoder* self, struct vec* dst,
+int tight_encode_frame(struct tight_encoder* self,
 		const struct rfb_pixel_format* dfmt,
 		const struct nvnc_fb* src,
 		const struct rfb_pixel_format* sfmt,
 		struct pixman_region16* damage,
-		enum tight_quality quality);
+		enum tight_quality quality,
+		tight_done_fn on_done, void* userdata);

src/server.c

@@ -72,6 +72,9 @@ struct fb_update_work {
 int schedule_client_update_fb(struct nvnc_client* client,
 		struct pixman_region16* damage);
 static int send_desktop_resize(struct nvnc_client* client, struct nvnc_fb* fb);
+static enum rfb_encodings choose_frame_encoding(struct nvnc_client* client);
+static enum tight_quality client_get_tight_quality(struct nvnc_client* client);
+static void on_tight_encode_frame_done(struct vec* frame, void* userdata);
 
 #if defined(GIT_VERSION)
 EXPORT const char nvnc_version[] = GIT_VERSION;
@@ -525,7 +528,35 @@ static void process_fb_update_requests(struct nvnc_client* client)
 	pixman_region_init(&client->damage);
 
 	client->is_updating = true;
-	if (schedule_client_update_fb(client, &damage) < 0)
+
+	int rc;
+	enum rfb_encodings encoding = choose_frame_encoding(client);
+
+	// TODO: Check the return value
+	struct rfb_pixel_format server_fmt;
+	rfb_pixfmt_from_fourcc(&server_fmt, fb->fourcc_format);
+
+	switch (encoding) {
+	case RFB_ENCODING_RAW:
+	case RFB_ENCODING_ZRLE:
+		rc = schedule_client_update_fb(client, &damage);
+		break;
+	case RFB_ENCODING_TIGHT:;
+		enum tight_quality quality = client_get_tight_quality(client);
+
+		// TODO: Make sure we don't have use-after-free on client and fb
+		rc = tight_encode_frame(&client->tight_encoder, &client->pixfmt,
+				fb, &server_fmt, &damage, quality,
+				on_tight_encode_frame_done, client);
+
+		pixman_region_fini(&damage);
+		break;
+	default:
+		rc = -1;
+		break;
+	}
+
+	if (rc < 0)
 		client->is_updating = false;
 }
 
@@ -1117,11 +1148,8 @@ static void do_client_update_fb(void* work)
 		raw_encode_frame(&update->frame, &client->pixfmt, fb,
 		                 &update->server_fmt, &update->region);
 		break;
-	case RFB_ENCODING_TIGHT:;
+	case RFB_ENCODING_TIGHT:
-		enum tight_quality quality = client_get_tight_quality(client);
+		abort();
-		tight_encode_frame(&client->tight_encoder, &update->frame,
-				&client->pixfmt, fb, &update->server_fmt,
-				&update->region, quality);
 		break;
 	case RFB_ENCODING_ZRLE:
 		zrle_encode_frame(&client->z_stream, &update->frame,
@@ -1154,6 +1182,12 @@ static void finish_fb_update(struct nvnc_client* client, struct vec* frame)
 	DTRACE_PROBE1(neatvnc, update_fb_done, client);
 }
 
+static void on_tight_encode_frame_done(struct vec* frame, void* userdata)
+{
+	struct nvnc_client* client = userdata;
+	finish_fb_update(client, frame);
+}
+
 static void on_client_update_fb_done(void* work)
 {
 	struct fb_update_work* update = aml_get_userdata(work);

src/tight.c

@@ -73,6 +73,7 @@ struct tight_zs_worker_ctx {
 
 static void do_tight_zs_work(void*);
 static void on_tight_zs_work_done(void*);
+static int schedule_tight_finish(struct tight_encoder* self);
 
 static int tight_encoder_init_stream(z_stream* zs)
 {
@@ -158,21 +159,13 @@ int tight_encoder_init(struct tight_encoder* self, uint32_t width,
 	tight_init_zs_worker(self, 2);
 	tight_init_zs_worker(self, 3);
 
-	pthread_mutex_init(&self->wait_mutex, NULL);
+	aml_require_workers(aml_get_default(), 1);
-	pthread_cond_init(&self->wait_cond, NULL);
-
-	// One worker is blocked while other workers are encoding, so at least
-	// 2 are required.
-	aml_require_workers(aml_get_default(), 2);
 
 	return 0;
 }
 
 void tight_encoder_destroy(struct tight_encoder* self)
 {
-	pthread_cond_destroy(&self->wait_cond);
-	pthread_mutex_destroy(&self->wait_mutex);
-
 	aml_unref(self->zs_worker[3]);
 	aml_unref(self->zs_worker[2]);
 	aml_unref(self->zs_worker[1]);
@@ -256,7 +249,7 @@ static void tight_encode_tile_basic(struct tight_encoder* self,
 	z_stream* zs = &self->zs[zs_index];
 	tile->type = TIGHT_BASIC | TIGHT_STREAM(zs_index);
 
-	int bytes_per_cpixel = calc_bytes_per_cpixel(self->dfmt);
+	int bytes_per_cpixel = calc_bytes_per_cpixel(&self->dfmt);
 	assert(bytes_per_cpixel <= 4);
 	uint8_t row[TSL * 4];
 
@@ -264,7 +257,7 @@ static void tight_encode_tile_basic(struct tight_encoder* self,
 	if (bytes_per_cpixel == 3)
 		rfb_pixfmt_from_fourcc(&cfmt, DRM_FORMAT_XBGR8888);
 	else
-		memcpy(&cfmt, self->dfmt, sizeof(cfmt));
+		memcpy(&cfmt, &self->dfmt, sizeof(cfmt));
 
 	uint32_t* addr = nvnc_fb_get_addr(self->fb);
 	uint32_t stride = nvnc_fb_get_width(self->fb);
@@ -272,8 +265,8 @@ static void tight_encode_tile_basic(struct tight_encoder* self,
 	// TODO: Limit width and hight to the sides
 	for (uint32_t y = y_start; y < y_start + height; ++y) {
 		void* img = addr + x + y * stride;
-		pixel32_to_cpixel(row, &cfmt, img, self->sfmt, bytes_per_cpixel,
+		pixel32_to_cpixel(row, &cfmt, img, &self->sfmt,
-				width);
+				bytes_per_cpixel, width);
 
 		// TODO What to do if the buffer fills up?
 		if (tight_deflate(tile, row, bytes_per_cpixel * width,
@@ -409,10 +402,8 @@ static void on_tight_zs_work_done(void* obj)
 	struct tight_zs_worker_ctx* ctx = aml_get_userdata(obj);
 	struct tight_encoder* self = ctx->encoder;
 
-	pthread_mutex_lock(&self->wait_mutex);
 	if (--self->n_jobs == 0)
-		pthread_cond_signal(&self->wait_cond);
+		schedule_tight_finish(self);
-	pthread_mutex_unlock(&self->wait_mutex);
 }
 
 static int tight_schedule_zs_work(struct tight_encoder* self, int index)
@@ -426,17 +417,11 @@ static int tight_schedule_zs_work(struct tight_encoder* self, int index)
 
 static int tight_schedule_encoding_jobs(struct tight_encoder* self)
 {
-	int rc = -1;
-
-	pthread_mutex_lock(&self->wait_mutex);
 	for (int i = 0; i < 4; ++i)
 		if (tight_schedule_zs_work(self, i) < 0)
-			goto failure;
+			return -1;
 
-	rc = 0;
+	return 0;
-failure:
-	pthread_mutex_unlock(&self->wait_mutex);
-	return rc;
 }
 
 static void tight_finish_tile(struct tight_encoder* self,
@@ -450,11 +435,11 @@ static void tight_finish_tile(struct tight_encoder* self,
 	uint32_t width = tight_tile_width(self, x);
 	uint32_t height = tight_tile_height(self, y);
 
-	encode_rect_head(self->dst, RFB_ENCODING_TIGHT, x, y, width, height);
+	encode_rect_head(&self->dst, RFB_ENCODING_TIGHT, x, y, width, height);
 
-	vec_append(self->dst, &tile->type, sizeof(tile->type));
+	vec_append(&self->dst, &tile->type, sizeof(tile->type));
-	tight_encode_size(self->dst, tile->size);
+	tight_encode_size(&self->dst, tile->size);
-	vec_append(self->dst, tile->buffer, tile->size);
+	vec_append(&self->dst, tile->buffer, tile->size);
 
 	tile->state = TIGHT_TILE_READY;
 }
@@ -467,36 +452,61 @@ static void tight_finish(struct tight_encoder* self)
 				tight_finish_tile(self, x, y);
 }
 
-int tight_encode_frame(struct tight_encoder* self, struct vec* dst,
+static void do_tight_finish(void* obj)
+{
+	// TODO: Make sure there's no use-after-free here
+	struct tight_encoder* self = aml_get_userdata(obj);
+	tight_finish(self);
+}
+
+static void on_tight_finished(void* obj)
+{
+	struct tight_encoder* self = aml_get_userdata(obj);
+	self->on_frame_done(&self->dst, self->userdata);
+}
+
+static int schedule_tight_finish(struct tight_encoder* self)
+{
+	struct aml_work* work = aml_work_new(do_tight_finish, on_tight_finished,
+			self, NULL);
+	if (!work)
+		return -1;
+
+	int rc = aml_start(aml_get_default(), work);
+	aml_unref(work);
+	return rc;
+}
+
+int tight_encode_frame(struct tight_encoder* self,
 		const struct rfb_pixel_format* dfmt,
 		const struct nvnc_fb* src,
 		const struct rfb_pixel_format* sfmt,
 		struct pixman_region16* damage,
-		enum tight_quality quality)
+		enum tight_quality quality,
+		tight_done_fn on_done, void* userdata)
 {
-	self->dfmt = dfmt;
+	memcpy(&self->dfmt, dfmt, sizeof(self->dfmt));
-	self->sfmt = sfmt;
+	memcpy(&self->sfmt, sfmt, sizeof(self->sfmt));
 	self->fb = src;
-	self->dst = dst;
 	self->quality = quality;
+	self->on_frame_done = on_done;
+	self->userdata = userdata;
 
-	vec_clear(dst);
+	uint32_t width = nvnc_fb_get_width(src);
+	uint32_t height = nvnc_fb_get_height(src);
+	int rc = vec_init(&self->dst, width * height * 4);
+	if (rc < 0)
+		return -1;
 
 	self->n_rects = tight_apply_damage(self, damage);
 	assert(self->n_rects > 0);
 
-	encode_rect_count(dst, self->n_rects);
+	encode_rect_count(&self->dst, self->n_rects);
 
-	if (tight_schedule_encoding_jobs(self) < 0)
+	if (tight_schedule_encoding_jobs(self) < 0) {
+		vec_destroy(&self->dst);
 		return -1;
-
+	}
-	// TODO Change architecture so we don't have to wait here
-	pthread_mutex_lock(&self->wait_mutex);
-	while (self->n_jobs != 0)
-		pthread_cond_wait(&self->wait_cond, &self->wait_mutex);
-	pthread_mutex_unlock(&self->wait_mutex);
-
-	tight_finish(self);
 
 	return 0;
 }