From bc8295baeec2069679bbbf481afe2ee9532fb578 Mon Sep 17 00:00:00 2001
From: Alex X <alexey.khit@gmail.com>
Date: Sun, 19 May 2024 11:56:33 +0300
Subject: [PATCH] Improve play audio on RTSP backchannel

---
 internal/ffmpeg/README.md |  7 +++++
 internal/streams/play.go  | 16 +++++++-----
 pkg/rtsp/consumer.go      | 55 +++++++++++++++++++++++----------------
 3 files changed, 49 insertions(+), 29 deletions(-)

diff --git a/internal/ffmpeg/README.md b/internal/ffmpeg/README.md
index 88a91d45..f89996e1 100644
--- a/internal/ffmpeg/README.md
+++ b/internal/ffmpeg/README.md
@@ -45,6 +45,13 @@
 [video4linux2,v4l2 @ 0x7f7de7c58bc0] Compressed:       mjpeg :          Motion-JPEG : 640x480 160x120 176x144 320x176 320x240 352x288 432x240 544x288 640x360 752x416 800x448 800x600 864x480 960x544 960x720 1024x576 1184x656 1280x720 1280x960
 ```
 
+## TTS
+
+```yaml
+streams:
+  tts: ffmpeg:#input=-re -f lavfi -i "flite=text='1 2 3 4 5 6 7 8 9 0'"#audio=pcma
+```
+
 ## Useful links
 
 - https://superuser.com/questions/564402/explanation-of-x264-tune
diff --git a/internal/streams/play.go b/internal/streams/play.go
index 748130f4..7ada66e6 100644
--- a/internal/streams/play.go
+++ b/internal/streams/play.go
@@ -2,6 +2,8 @@ package streams
 
 import (
 	"errors"
+	"time"
+
 	"github.com/AlexxIT/go2rtc/pkg/core"
 )
 
@@ -80,18 +82,20 @@ func (s *Stream) Play(source string) error {
 		s.AddInternalProducer(src)
 		s.AddInternalConsumer(cons)
 
-		go func() {
-			_ = src.Start()
-			_ = dst.Stop()
-			s.RemoveProducer(src)
-		}()
-
 		go func() {
 			_ = dst.Start()
 			_ = src.Stop()
 			s.RemoveInternalConsumer(cons)
 		}()
 
+		go func() {
+			_ = src.Start()
+			// little timeout before stop dst, so the buffer can be transferred
+			time.Sleep(time.Second)
+			_ = dst.Stop()
+			s.RemoveProducer(src)
+		}()
+
 		return nil
 	}
 
diff --git a/pkg/rtsp/consumer.go b/pkg/rtsp/consumer.go
index 4bddd77b..79e2b348 100644
--- a/pkg/rtsp/consumer.go
+++ b/pkg/rtsp/consumer.go
@@ -74,19 +74,38 @@ func (c *Conn) AddTrack(media *core.Media, codec *core.Codec, track *core.Receiv
 	return nil
 }
 
+const (
+	startVideoBuf = 32 * 1024   // 32KB
+	startAudioBuf = 2 * 1024    // 2KB
+	maxBuf        = 1024 * 1024 // 1MB
+	rtpHdr        = 12          // basic RTP header size
+	intHdr        = 4           // interleaved header size
+)
+
 func (c *Conn) packetWriter(codec *core.Codec, channel, payloadType uint8) core.HandlerFunc {
 	var buf []byte
 	var n int
 
 	video := codec.IsVideo()
 	if video {
-		buf = make([]byte, 32*1024) // 32KB
+		buf = make([]byte, startVideoBuf)
 	} else {
-		buf = make([]byte, 2*1024) // 2KB
+		buf = make([]byte, startAudioBuf)
+	}
+
+	flushBuf := func() {
+		if err := c.conn.SetWriteDeadline(time.Now().Add(Timeout)); err != nil {
+			return
+		}
+		//log.Printf("[rtsp] channel:%2d write_size:%6d buffer_size:%6d", channel, n, len(buf))
+		if _, err := c.conn.Write(buf[:n]); err == nil {
+			c.send += n
+		}
+		n = 0
 	}
 
 	handlerFunc := func(packet *rtp.Packet) {
-		if c.state == StateNone || !c.playOK {
+		if c.state == StateNone {
 			return
 		}
 
@@ -106,16 +125,13 @@ func (c *Conn) packetWriter(codec *core.Codec, channel, payloadType uint8) core.
 			packet.Marker = true // better to have marker on all audio packets
 		}
 
-		size := 12 + len(packet.Payload)
+		size := rtpHdr + len(packet.Payload)
 
-		if n+4+size > len(buf) {
-			if len(buf) < 1024*1024 {
-				buf = append(buf, make([]byte, len(buf))...)
+		if l := len(buf); n+intHdr+size > l {
+			if l < maxBuf {
+				buf = append(buf, make([]byte, l)...) // double buffer size
 			} else {
-				if _, err := c.conn.Write(buf[:n]); err == nil {
-					c.send += n
-				}
-				n = 0
+				flushBuf()
 			}
 		}
 
@@ -134,21 +150,14 @@ func (c *Conn) packetWriter(codec *core.Codec, channel, payloadType uint8) core.
 
 		n += 4 + size
 
-		if !packet.Marker {
-			return // collect continious video packets to buffer
-		}
-
-		if err := c.conn.SetWriteDeadline(time.Now().Add(Timeout)); err != nil {
+		if !packet.Marker || !c.playOK {
+			// collect continious video packets to buffer
+			// or wait OK for PLAY command for backchannel
+			//log.Printf("[rtsp] collecting buffer ok=%t", c.playOK)
 			return
 		}
 
-		//log.Printf("[rtsp] channel:%2d write_size:%6d buffer_size:%6d", channel, n, len(buf))
-
-		if _, err := c.conn.Write(buf[:n]); err == nil {
-			c.send += n
-		}
-
-		n = 0
+		flushBuf()
 	}
 
 	if !codec.IsRTP() {