examples: restart video files when they end (#733)

examples/client-record-format-h264-from-disk/main.go

@@ -2,11 +2,14 @@ package main
 
 import (
 	"crypto/rand"
+	"errors"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"time"
 
+	"github.com/asticode/go-astits"
 	"github.com/bluenviron/gortsplib/v4"
 	"github.com/bluenviron/gortsplib/v4/pkg/description"
 	"github.com/bluenviron/gortsplib/v4/pkg/format"
@@ -38,26 +41,6 @@ func randUint32() (uint32, error) {
 }
 
 func main() {
-	// open a file in MPEG-TS format
-	f, err := os.Open("myvideo.ts")
-	if err != nil {
-		panic(err)
-	}
-	defer f.Close()
-
-	// setup MPEG-TS parser
-	r := &mpegts.Reader{R: f}
-	err = r.Initialize()
-	if err != nil {
-		panic(err)
-	}
-
-	// find the H264 track inside the file
-	track, err := findTrack(r)
-	if err != nil {
-		panic(err)
-	}
-
 	// create a RTSP description that contains a H264 format
 	forma := &format.H264{
 		PayloadTyp:        96,
@@ -72,12 +55,19 @@ func main() {
 
 	// connect to the server, announce the format and start recording
 	c := gortsplib.Client{}
-	err = c.StartRecording("rtsp://myuser:mypass@localhost:8554/mystream", desc)
+	err := c.StartRecording("rtsp://myuser:mypass@localhost:8554/mystream", desc)
 	if err != nil {
 		panic(err)
 	}
 	defer c.Close()
 
+	// open a file in MPEG-TS format
+	f, err := os.Open("myvideo.ts")
+	if err != nil {
+		panic(err)
+	}
+	defer f.Close()
+
 	// setup H264 -> RTP encoder
 	rtpEnc, err := forma.CreateEncoder()
 	if err != nil {
@@ -89,59 +79,92 @@ func main() {
 		panic(err)
 	}
 
-	timeDecoder := mpegts.TimeDecoder{}
-	timeDecoder.Initialize()
-
-	var firstDTS *int64
-	var startTime time.Time
-
-	// setup a callback that is called whenever a H264 access unit is read from the file
-	r.OnDataH264(track, func(pts, dts int64, au [][]byte) error {
-		dts = timeDecoder.Decode(dts)
-		pts = timeDecoder.Decode(pts)
-
-		// sleep between access units
-		if firstDTS != nil {
-			timeDrift := time.Duration(dts-*firstDTS)*time.Second/90000 - time.Since(startTime)
-			if timeDrift > 0 {
-				time.Sleep(timeDrift)
-			}
-		} else {
-			startTime = time.Now()
-			firstDTS = &dts
-		}
-
-		log.Printf("writing access unit with pts=%d dts=%d", pts, dts)
-
-		// wrap the access unit into RTP packets
-		packets, err := rtpEnc.Encode(au)
-		if err != nil {
-			return err
-		}
-
-		// set packet timestamp
-		// we don't have to perform any conversion
-		// since H264 clock rate is the same in both MPEG-TS and RTSP
-		for _, packet := range packets {
-			packet.Timestamp = uint32(int64(randomStart) + pts)
-		}
-
-		// write RTP packets to the server
-		for _, packet := range packets {
-			err := c.WritePacketRTP(desc.Medias[0], packet)
-			if err != nil {
-				return err
-			}
-		}
-
-		return nil
-	})
-
-	// read the MPEG-TS file
 	for {
-		err := r.Read()
+		// setup MPEG-TS parser
+		r := &mpegts.Reader{R: f}
+		err = r.Initialize()
 		if err != nil {
 			panic(err)
 		}
+
+		// find the H264 track inside the file
+		track, err := findTrack(r)
+		if err != nil {
+			panic(err)
+		}
+
+		timeDecoder := mpegts.TimeDecoder{}
+		timeDecoder.Initialize()
+
+		var firstDTS *int64
+		var startTime time.Time
+		var lastRTPTime uint32
+
+		// setup a callback that is called when a H264 access unit is read from the file
+		r.OnDataH264(track, func(pts, dts int64, au [][]byte) error {
+			dts = timeDecoder.Decode(dts)
+			pts = timeDecoder.Decode(pts)
+
+			// sleep between access units
+			if firstDTS != nil {
+				timeDrift := time.Duration(dts-*firstDTS)*time.Second/90000 - time.Since(startTime)
+				if timeDrift > 0 {
+					time.Sleep(timeDrift)
+				}
+			} else {
+				startTime = time.Now()
+				firstDTS = &dts
+			}
+
+			log.Printf("writing access unit with pts=%d dts=%d", pts, dts)
+
+			// wrap the access unit into RTP packets
+			packets, err := rtpEnc.Encode(au)
+			if err != nil {
+				return err
+			}
+
+			// set packet timestamp
+			// we don't have to perform any conversion
+			// since H264 clock rate is the same in both MPEG-TS and RTSP
+			lastRTPTime = uint32(int64(randomStart) + pts)
+			for _, packet := range packets {
+				packet.Timestamp = lastRTPTime
+			}
+
+			// write RTP packets to the server
+			for _, packet := range packets {
+				err := c.WritePacketRTP(desc.Medias[0], packet)
+				if err != nil {
+					return err
+				}
+			}
+
+			return nil
+		})
+
+		// read the file
+		for {
+			err := r.Read()
+			if err != nil {
+				// file has ended
+				if errors.Is(err, astits.ErrNoMorePackets) {
+					log.Printf("file has ended, rewinding")
+
+					// rewind to start position
+					_, err = f.Seek(0, io.SeekStart)
+					if err != nil {
+						panic(err)
+					}
+
+					// keep current timestamp
+					randomStart = lastRTPTime + 1
+
+					break
+				}
+
+				panic(err)
+			}
+		}
 	}
 }