From f18ae2b972c80d7c5d273191582ea629851d4c10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Thu, 2 Sep 2021 00:00:56 +0200
Subject: [PATCH] add waitgroups.

---
 internal/capture/screencast.go |  8 +++++++-
 internal/capture/stream.go     |  8 +++++++-
 internal/desktop/manager.go    | 18 +++++++++++-------
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/internal/capture/screencast.go b/internal/capture/screencast.go
index bec88894..179e8b32 100644
--- a/internal/capture/screencast.go
+++ b/internal/capture/screencast.go
@@ -16,6 +16,7 @@ import (
 type ScreencastManagerCtx struct {
 	logger      zerolog.Logger
 	mu          sync.Mutex
+	wg          sync.WaitGroup
 	pipelineStr string
 	pipeline    *gst.Pipeline
 	enabled     bool
@@ -39,8 +40,11 @@ func screencastNew(enabled bool, pipelineStr string) *ScreencastManagerCtx {
 		emitUpdate:  make(chan bool),
 	}
 
+	manager.wg.Add(1)
+
 	go func() {
 		manager.logger.Debug().Msg("started emitting samples")
+		defer manager.wg.Done()
 
 		ticker := time.NewTicker(screencastTimeout)
 		defer ticker.Stop()
@@ -66,10 +70,12 @@ func screencastNew(enabled bool, pipelineStr string) *ScreencastManagerCtx {
 }
 
 func (manager *ScreencastManagerCtx) shutdown() {
-	manager.logger.Info().Msgf("shutting down")
+	manager.logger.Info().Msgf("shutdown")
 
 	manager.destroyPipeline()
+
 	manager.emitStop <- true
+	manager.wg.Wait()
 }
 
 func (manager *ScreencastManagerCtx) Enabled() bool {
diff --git a/internal/capture/stream.go b/internal/capture/stream.go
index 34f6a91f..f8af877c 100644
--- a/internal/capture/stream.go
+++ b/internal/capture/stream.go
@@ -15,6 +15,7 @@ import (
 type StreamManagerCtx struct {
 	logger      zerolog.Logger
 	mu          sync.Mutex
+	wg          sync.WaitGroup
 	codec       codec.RTPCodec
 	pipelineStr func() string
 	pipeline    *gst.Pipeline
@@ -42,8 +43,11 @@ func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string)
 		started:     false,
 	}
 
+	manager.wg.Add(1)
+
 	go func() {
 		manager.logger.Debug().Msg("started emitting samples")
+		defer manager.wg.Done()
 
 		for {
 			select {
@@ -66,7 +70,7 @@ func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string)
 }
 
 func (manager *StreamManagerCtx) shutdown() {
-	manager.logger.Info().Msgf("shutting down")
+	manager.logger.Info().Msgf("shutdown")
 
 	manager.emitMu.Lock()
 	for key := range manager.listeners {
@@ -75,7 +79,9 @@ func (manager *StreamManagerCtx) shutdown() {
 	manager.emitMu.Unlock()
 
 	manager.destroyPipeline()
+
 	manager.emitStop <- true
+	manager.wg.Wait()
 }
 
 func (manager *StreamManagerCtx) Codec() codec.RTPCodec {
diff --git a/internal/desktop/manager.go b/internal/desktop/manager.go
index 5909a9f5..7a1fab7b 100644
--- a/internal/desktop/manager.go
+++ b/internal/desktop/manager.go
@@ -18,6 +18,7 @@ var mu = sync.Mutex{}
 
 type DesktopManagerCtx struct {
 	logger   zerolog.Logger
+	wg       sync.WaitGroup
 	shutdown chan bool
 	emmiter  events.EventEmmiter
 	config   *config.Desktop
@@ -39,14 +40,11 @@ func (manager *DesktopManagerCtx) Start() {
 
 	xorg.GetScreenConfigurations()
 
-	manager.logger.Info().
+	err := xorg.ChangeScreenSize(manager.config.ScreenWidth, manager.config.ScreenHeight, manager.config.ScreenRate)
+	manager.logger.Err(err).
 		Str("screen_size", fmt.Sprintf("%dx%d@%d", manager.config.ScreenWidth, manager.config.ScreenHeight, manager.config.ScreenRate)).
 		Msgf("setting initial screen size")
 
-	if err := xorg.ChangeScreenSize(manager.config.ScreenWidth, manager.config.ScreenHeight, manager.config.ScreenRate); err != nil {
-		manager.logger.Err(err).Msg("unable to set initial screen size")
-	}
-
 	go xevent.EventLoop(manager.config.Display)
 
 	// In case it was opened
@@ -61,14 +59,17 @@ func (manager *DesktopManagerCtx) Start() {
 			Msg("X event error occured")
 	})
 
+	manager.wg.Add(1)
+
 	go func() {
+		defer manager.wg.Done()
+
 		ticker := time.NewTicker(1 * time.Second)
 		defer ticker.Stop()
 
 		for {
 			select {
 			case <-manager.shutdown:
-				xorg.DisplayClose()
 				return
 			case <-ticker.C:
 				xorg.CheckKeys(time.Second * 10)
@@ -90,8 +91,11 @@ func (manager *DesktopManagerCtx) OnAfterScreenSizeChange(listener func()) {
 }
 
 func (manager *DesktopManagerCtx) Shutdown() error {
-	manager.logger.Info().Msgf("desktop shutting down")
+	manager.logger.Info().Msgf("shutdown")
 
 	manager.shutdown <- true
+	manager.wg.Wait()
+
+	xorg.DisplayClose()
 	return nil
 }