src/internal/screencapturer/wholescreencapturer.go

@@ -57,10 +57,13 @@ type wholeScreenCapturer struct {
 	latest string
 	done   chan struct{}
 	sync.WaitGroup
+	mu sync.RWMutex
 }
 
 // Get implements ScreenCapturer.
 func (w *wholeScreenCapturer) Get() (filepath string, err error) {
+	w.mu.RLock()
+	defer w.mu.Unlock()
 	if !w.initialized {
 		return "", ErrNotInitialized
 	}
@@ -185,6 +188,8 @@ func (w *wholeScreenCapturer) captureAndSave() (string, error) {
 }
 
 func (w *wholeScreenCapturer) addToBuffer(fp string) {
+	w.mu.Lock()
+    defer w.mu.Unlock()
 	w.files = append(w.files, fp)
 	w.latest = fp
 
@@ -210,8 +215,25 @@ func NewWholeScreenCapturer(lc fx.Lifecycle, config *config.Config, log *logger.
 		done:   make(chan struct{}),
 	}
 	lc.Append(fx.StopHook(func(ctx context.Context) error {
-		log.Debug("Stopping capturer (fx hook)")
-		return capturer.Stop()
+		if (!capturer.initialized) {
+			log.Debug("wholescreencapturer not initialized, nothing to do")
+			return nil
+		}
+
+		log.Debug("stopping wholescreencapturer")
+		
+		err := capturer.Stop()
+		if err != nil {
+			log.Error("failed to stop wholescreencapturer gracefully")
+			return err
+		}
+
+		err = os.RemoveAll(capturer.tempDirectory)
+		if err != nil {
+			log.Error("failed to remove temp directory")
+			return err
+		}
+		return nil
 	}))
 	return capturer
 }