From f45f90ed9b190ccb6654c09dfc337969cb64f687 Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Sat, 4 Apr 2020 12:33:18 +0000
Subject: [PATCH] nvnc_fb: Add a mechanism for signaling that a frame is in use

---
 include/fb.h      |  2 ++
 include/neatvnc.h |  3 +++
 src/fb.c          | 14 ++++++++++++++
 src/server.c      | 13 +++++++++----
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/include/fb.h b/include/fb.h
index 260538a..e101236 100644
--- a/include/fb.h
+++ b/include/fb.h
@@ -2,10 +2,12 @@
 
 #include <unistd.h>
 #include <stdint.h>
+#include <stdatomic.h>
 
 struct nvnc_fb {
 	int ref;
 	void* addr;
+	atomic_bool is_locked;
 	size_t size;
 	uint16_t width;
 	uint16_t height;
diff --git a/include/neatvnc.h b/include/neatvnc.h
index 4f700b5..e0d2a30 100644
--- a/include/neatvnc.h
+++ b/include/neatvnc.h
@@ -73,6 +73,9 @@ struct nvnc_fb* nvnc_fb_new(uint16_t width, uint16_t height,
 void nvnc_fb_ref(struct nvnc_fb* fb);
 void nvnc_fb_unref(struct nvnc_fb* fb);
 
+bool nvnc_fb_lock(struct nvnc_fb*);
+void nvnc_fb_unlock(struct nvnc_fb*);
+
 void* nvnc_fb_get_addr(const struct nvnc_fb* fb);
 uint16_t nvnc_fb_get_width(const struct nvnc_fb* fb);
 uint16_t nvnc_fb_get_height(const struct nvnc_fb* fb);
diff --git a/src/fb.c b/src/fb.c
index 189020d..59e44e7 100644
--- a/src/fb.c
+++ b/src/fb.c
@@ -4,6 +4,7 @@
 #include <stdlib.h>
 #include <unistd.h>
 #include <sys/param.h>
+#include <stdatomic.h>
 
 #define UDIV_UP(a, b) (((a) + (b) - 1) / (b))
 #define ALIGN_UP(n, a) (UDIV_UP(n, a) * a)
@@ -39,6 +40,19 @@ struct nvnc_fb* nvnc_fb_new(uint16_t width, uint16_t height,
 	return fb;
 }
 
+EXPORT
+bool nvnc_fb_lock(struct nvnc_fb* fb)
+{
+	bool expected = false;
+	return atomic_compare_exchange_strong(&fb->is_locked, &expected, true);
+}
+
+EXPORT
+void nvnc_fb_unlock(struct nvnc_fb* fb)
+{
+	fb->is_locked = false;
+}
+
 EXPORT
 void* nvnc_fb_get_addr(const struct nvnc_fb* fb)
 {
diff --git a/src/server.c b/src/server.c
index 69393a5..90ed9f4 100644
--- a/src/server.c
+++ b/src/server.c
@@ -474,9 +474,13 @@ static void process_fb_update_requests(struct nvnc_client* client)
 	if (client->is_updating || client->n_pending_requests == 0)
 		return;
 
+	if (!nvnc_fb_lock(client->server->frame))
+		return;
+
 	client->is_updating = true;
 
-	schedule_client_update_fb(client);
+	if (schedule_client_update_fb(client) < 0)
+		nvnc_fb_unlock(client->server->frame);
 }
 
 static int on_client_fb_update_request(struct nvnc_client* client)
@@ -855,7 +859,7 @@ void do_client_update_fb(void* work)
 {
 	struct fb_update_work* update = aml_get_userdata(work);
 	struct nvnc_client* client = update->client;
-	const struct nvnc_fb* fb = update->fb;
+	struct nvnc_fb* fb = update->fb;
 
 	enum rfb_encodings encoding = choose_frame_encoding(client);
 	if (encoding == -1) {
@@ -896,8 +900,10 @@ void on_client_update_fb_done(void* work)
 	struct nvnc_client* client = update->client;
 	struct vec* frame = &update->frame;
 
-	client_ref(client);
+	nvnc_fb_unlock(update->fb);
+	nvnc_fb_unref(update->fb);
 
+	client_ref(client);
 
 	if (client->net_stream->state != STREAM_STATE_CLOSED) {
 		struct rcbuf* payload = rcbuf_new(frame->data, frame->len);
@@ -913,7 +919,6 @@ void on_client_update_fb_done(void* work)
 
 	client->n_pending_requests--;
 	process_fb_update_requests(client);
-	nvnc_fb_unref(update->fb);
 
 	DTRACE_PROBE1(neatvnc, update_fb_done, client);