fix daemon's concurrency issues

internal/daemon/daemon.go

@@ -24,6 +24,8 @@ type Daemon struct {
 	cancel context.CancelFunc
 
 	pipeline pipeline.Pipeline
+
+	wg sync.WaitGroup
 }
 
 func New(n notify.Notifier) *Daemon {
@@ -41,12 +43,25 @@ func New(n notify.Notifier) *Daemon {
 }
 
 func (d *Daemon) status() pipeline.Status {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
 	if d.pipeline == nil {
 		return pipeline.Idle
 	}
 	return d.pipeline.Status()
 }
 
+func (d *Daemon) stopPipeline() {
+	d.mu.Lock()
+	p := d.pipeline
+	d.pipeline = nil
+	d.mu.Unlock()
+
+	if p != nil {
+		p.Stop()
+	}
+}
+
 func (d *Daemon) Run() error {
 	if err := bus.CheckExistingDaemon(); err != nil {
 		return err
@@ -73,10 +88,11 @@ func (d *Daemon) Run() error {
 		d.cancel()
 	}()
 
-	// Close the listener when context is done
 	go func() {
 		<-d.ctx.Done()
-		ln.Close()
+		if err := ln.Close(); err != nil {
+			log.Printf("Error closing listener: %v", err)
+		}
 	}()
 
 	log.Printf("Daemon started, listening on socket")
@@ -85,18 +101,21 @@ func (d *Daemon) Run() error {
 		c, err := ln.Accept()
 		if err != nil {
 			if d.ctx.Err() != nil {
-				log.Printf("Shutdown requested")
+				log.Printf("Shutdown requested, waiting for connections to finish")
+				d.wg.Wait()
 				return nil
 			}
 			log.Printf("Accept error: %v", err)
 			return fmt.Errorf("accept failed: %w", err)
 		}
+		d.wg.Add(1)
 		go d.handle(c)
 	}
 }
 
 func (d *Daemon) handle(c net.Conn) {
 	defer c.Close()
+	defer d.wg.Done()
 
 	line, err := bufio.NewReader(c).ReadString('\n')
 	if err != nil {
@@ -129,36 +148,34 @@ func (d *Daemon) handle(c net.Conn) {
 }
 
 func (d *Daemon) toggle() {
-	var actionChan chan<- pipeline.Action
-
 	switch d.status() {
 	case pipeline.Idle:
 		p := pipeline.New()
 		p.Run(d.ctx)
+
+		d.mu.Lock()
 		d.pipeline = p
+		d.mu.Unlock()
+
 		go d.notifier.RecordingStarted()
 
 	case pipeline.Recording:
+		d.stopPipeline() // aborted during recording (chunks not sent to transcriber yet)
 		go d.notifier.Aborted()
-		if d.pipeline != nil {
-			d.pipeline.Stop() // aborted during recording (chunks not sent to transcriber yet)
-			d.pipeline = nil
-		}
 
 	case pipeline.Transcribing:
+		d.mu.RLock()
+		if d.pipeline != nil {
+			actionChan := d.pipeline.GetActionCh()
+			d.mu.RUnlock()
+			actionChan <- pipeline.Inject
+		} else {
+			d.mu.RUnlock()
+		}
 		go d.notifier.RecordingEnded()
-		actionChan = d.pipeline.Actions()
-		actionChan <- pipeline.Inject
 
 	case pipeline.Injecting:
+		d.stopPipeline() // aborted during injection
 		go d.notifier.Aborted()
-
-		if d.pipeline != nil {
-			d.pipeline.Stop() // aborted during injection
-			d.pipeline = nil
-		}
-
-	default:
-		d.mu.Unlock()
 	}
 }

internal/notify/notify.go

@@ -56,5 +56,6 @@ type Nop struct{}
 
 func (Nop) RecordingStarted() {}
 func (Nop) RecordingEnded()   {}
+func (Nop) Aborted()          {}
 func (Nop) Transcribing()     {}
 func (Nop) Error(msg string)  {}

internal/pipeline/pipeline.go

@@ -27,14 +27,20 @@ type Pipeline interface {
 	Run(ctx context.Context)
 	Stop()
 	Status() Status
-	Actions() chan<- Action
+	GetActionCh() chan<- Action
 }
 
 type pipeline struct {
 	status   Status
 	actionCh chan Action
-	wg       sync.WaitGroup
-	cancel   context.CancelFunc
+
+	mu sync.RWMutex
+	wg sync.WaitGroup
+
+	cancel context.CancelFunc
+
+	stopOnce sync.Once
+	running  bool
 }
 
 func New() Pipeline {
@@ -44,40 +50,79 @@ func New() Pipeline {
 }
 
 func (p *pipeline) Status() Status {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
 	return p.status
 }
 
-func (p *pipeline) Actions() chan<- Action {
+func (p *pipeline) setStatus(status Status) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.status = status
+}
+
+func (p *pipeline) setCancel(cancel context.CancelFunc) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.cancel = cancel
+}
+
+func (p *pipeline) getCancel() context.CancelFunc {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.cancel
+}
+
+func (p *pipeline) GetActionCh() chan<- Action {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
 	return p.actionCh
 }
 
 func (p *pipeline) Stop() {
-	if p.cancel != nil {
-		p.cancel()
-	}
+	p.stopOnce.Do(func() {
+		cancel := p.getCancel()
+		if cancel != nil {
+			cancel()
+		}
+	})
 	p.wg.Wait()
 }
 
 func (p *pipeline) Run(ctx context.Context) {
+	p.mu.Lock()
+	if p.running {
+		p.mu.Unlock()
+		log.Printf("Pipeline: Already running, ignoring Run() call")
+		return
+	}
+	p.running = true
+	p.mu.Unlock()
+
 	runCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
-	p.cancel = cancel
+	p.setCancel(cancel)
+
 	p.wg.Add(1)
 	go p.run(runCtx)
 }
 
 func (p *pipeline) run(ctx context.Context) {
 	defer func() {
+		log.Printf("Pipeline: defered run")
+		p.mu.Lock()
+		p.running = false
+		p.mu.Unlock()
 		p.wg.Done()
 	}()
 
 	log.Printf("Pipeline: Starting recording")
-	p.status = Recording
+	p.setStatus(Recording)
 
 	recorder := recording.NewDefaultRecorder()
 	frameCh, errCh, err := recorder.Start(ctx)
 	if err != nil {
 		log.Printf("Pipeline: Recording error: %v", err)
-		p.status = Idle
+		p.setStatus(Idle)
 		return
 	}
 	defer recorder.Stop()
@@ -93,12 +138,12 @@ func (p *pipeline) run(ctx context.Context) {
 			log.Printf("Pipeline: Received frame #%d - Size: %d bytes, Timestamp: %v, Total bytes so far: %d",
 				frameCount, len(frame.Data), frame.Timestamp.Format("15:04:05.000"), totalBytes)
 
-			p.status = Transcribing
+			p.setStatus(Transcribing)
 
 		case err := <-errCh:
 			if err != nil {
 				log.Printf("Pipeline: Recording error: %v", err)
-				p.status = Idle
+				p.setStatus(Idle)
 				return
 			}
 
@@ -111,7 +156,13 @@ func (p *pipeline) run(ctx context.Context) {
 				if err := recorder.Stop(); err != nil {
 					log.Printf("Pipeline: Error stopping recorder: %v", err)
 				}
-				p.status = Injecting
+				p.setStatus(Injecting)
+
+				// Simulate injection work then return to idle
+				log.Printf("Pipeline: Simulating injection work")
+				time.Sleep(10 * time.Millisecond)
+				log.Printf("Pipeline: Injection work done, returning to idle")
+				p.setStatus(Idle)
 				return
 			}