Add display destroy listener in screenshooter, destroy display in rootston

3 files changed, 49 insertions, 10 deletions

diff --git a/include/wlr/types/wlr_screenshooter.h b/include/wlr/types/wlr_screenshooter.h
index 4bda3d3c..e8b59669 100644
--- a/include/wlr/types/wlr_screenshooter.h
+++ b/include/wlr/types/wlr_screenshooter.h
@@ -5,6 +5,13 @@
 struct wlr_screenshooter {
 	struct wl_global *wl_global;
 	struct wlr_renderer *renderer;
+	struct wl_list screenshots; // wlr_screenshot::link
+
+	struct wl_listener display_destroy;
+
+	struct {
+		struct wl_signal destroy;
+	} events;
 
 	void *data;
 };
@@ -12,6 +19,7 @@ struct wlr_screenshooter {
 struct wlr_screenshot {
 	struct wl_resource *resource;
 	struct wl_resource *output_resource;
+	struct wl_list link;
 
 	struct wlr_output *output;
 	struct wlr_screenshooter *screenshooter;
diff --git a/rootston/main.c b/rootston/main.c
index 46548094..3b65a067 100644
--- a/rootston/main.c
+++ b/rootston/main.c
@@ -51,6 +51,7 @@ int main(int argc, char **argv) {
 	if (!wlr_backend_start(server.backend)) {
 		wlr_log(L_ERROR, "Failed to start backend");
 		wlr_backend_destroy(server.backend);
+		wl_display_destroy(server.wl_display);
 		return 1;
 	}
 
@@ -69,5 +70,6 @@ int main(int argc, char **argv) {
 
 	wl_display_run(server.wl_display);
 	wlr_backend_destroy(server.backend);
+	wl_display_destroy(server.wl_display);
 	return 0;
 }
diff --git a/types/wlr_screenshooter.c b/types/wlr_screenshooter.c
index 94b45384..d62820e6 100644
--- a/types/wlr_screenshooter.c
+++ b/types/wlr_screenshooter.c
@@ -26,6 +26,18 @@ struct screenshot_state {
 	struct wl_listener frame_listener;
 };
 
+static void screenshot_destroy(struct wlr_screenshot *screenshot) {
+	wl_list_remove(&screenshot->link);
+	free(screenshot);
+}
+
+static void handle_screenshot_resource_destroy(
+		struct wl_resource *screenshot_resource) {
+	struct wlr_screenshot *screenshot =
+		wl_resource_get_user_data(screenshot_resource);
+	screenshot_destroy(screenshot);
+}
+
 static void output_frame_notify(struct wl_listener *listener, void *_data) {
 	struct screenshot_state *state = wl_container_of(listener, state,
 		frame_listener);
@@ -102,7 +114,8 @@ static void screenshooter_shoot(struct wl_client *client,
 		return;
 	}
 	wl_resource_set_implementation(screenshot->resource, NULL, screenshot,
-		NULL);
+		handle_screenshot_resource_destroy);
+	wl_list_insert(&screenshooter->screenshots, &screenshot->link);
 
 	wlr_log(L_DEBUG, "new screenshot %p (res %p)", screenshot,
 		screenshot->resource);
@@ -144,6 +157,25 @@ static void screenshooter_bind(struct wl_client *wl_client, void *data,
 		screenshooter, NULL);
 }
 
+void wlr_screenshooter_destroy(struct wlr_screenshooter *screenshooter) {
+	if (!screenshooter) {
+		return;
+	}
+	wl_signal_emit(&screenshooter->events.destroy, screenshooter);
+	struct wlr_screenshot *screenshot, *tmp;
+	wl_list_for_each_safe(screenshot, tmp, &screenshooter->screenshots, link) {
+		screenshot_destroy(screenshot);
+	}
+	wl_global_destroy(screenshooter->wl_global);
+	free(screenshooter);
+}
+
+static void handle_display_destroy(struct wl_listener *listener, void *data) {
+	struct wlr_screenshooter *screenshooter =
+		wl_container_of(listener, screenshooter, display_destroy);
+	wlr_screenshooter_destroy(screenshooter);
+}
+
 struct wlr_screenshooter *wlr_screenshooter_create(struct wl_display *display,
 		struct wlr_renderer *renderer) {
 	struct wlr_screenshooter *screenshooter =
@@ -153,6 +185,12 @@ struct wlr_screenshooter *wlr_screenshooter_create(struct wl_display *display,
 	}
 	screenshooter->renderer = renderer;
 
+	wl_list_init(&screenshooter->screenshots);
+	wl_signal_init(&screenshooter->events.destroy);
+
+	screenshooter->display_destroy.notify = handle_display_destroy;
+	wl_display_add_destroy_listener(display, &screenshooter->display_destroy);
+
 	screenshooter->wl_global = wl_global_create(display,
 		&orbital_screenshooter_interface, 1, screenshooter, screenshooter_bind);
 	if (screenshooter->wl_global == NULL) {
@@ -162,12 +200,3 @@ struct wlr_screenshooter *wlr_screenshooter_create(struct wl_display *display,
 
 	return screenshooter;
 }
-
-void wlr_screenshooter_destroy(struct wlr_screenshooter *screenshooter) {
-	if (!screenshooter) {
-		return;
-	}
-	// TODO: this segfault (wl_display->registry_resource_list is not init)
-	// wl_global_destroy(screenshooter->wl_global);
-	free(screenshooter);
-}