Use logrus for logging

go.mod

@@ -5,4 +5,5 @@ go 1.13
 require (
 	github.com/grafov/m3u8 v0.11.1
 	github.com/nareix/joy4 v0.0.0-20181022032202-3ddbc8f9d431
+	github.com/sirupsen/logrus v1.4.2
 )

go.sum

@@ -1,4 +1,17 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/grafov/m3u8 v0.11.1 h1:igZ7EBIB2IAsPPazKwRKdbhxcoBKO3lO1UY57PZDeNA=
 github.com/grafov/m3u8 v0.11.1/go.mod h1:nqzOkfBiZJENr52zTVd/Dcl03yzphIMbJqkXGu+u080=
+github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk=
+github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/nareix/joy4 v0.0.0-20181022032202-3ddbc8f9d431 h1:nWhrOsCKdV6bivw03k7MROF2tYzCFGfYBYFrTEHyucs=
 github.com/nareix/joy4 v0.0.0-20181022032202-3ddbc8f9d431/go.mod h1:aFJ1ZwLjvHN4yEzE5Bkz8rD8/d8Vlj3UIuvz2yfET7I=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.4.2 h1:SPIRibHv4MatM3XXNO2BJeFLZwZ2LvZgfQ5+UNI2im4=
+github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
+github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc=
+golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

main.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"io"
-	"log"
 	"os"
 	"path/filepath"
 	"strings"
@@ -13,6 +12,8 @@ import (
 	"github.com/nareix/joy4/av"
 	"github.com/nareix/joy4/format/rtmp"
 	"github.com/nareix/joy4/format/ts"
+
+	log "github.com/sirupsen/logrus"
 )
 
 const addr = ":1935"
@@ -21,7 +22,7 @@ const msPerSegment = 15000
 
 // TODO: replace failln with println / switch to logrus to include stream name in msg
 
-func removeOutdatedSegments(streamName string, playlist *m3u8.MediaPlaylist) error {
+func removeOutdatedSegments(streamLogger *log.Entry, streamName string, playlist *m3u8.MediaPlaylist) error {
 	currentSegments := make(map[string]struct{}, len(playlist.Segments))
 	for _, segment := range playlist.Segments {
 		if segment != nil {
@@ -34,7 +35,11 @@ func removeOutdatedSegments(streamName string, playlist *m3u8.MediaPlaylist) err
 	}
 	for _, segmentFile := range segmentFiles {
 		if _, ok := currentSegments[segmentFile]; !ok {
-			os.Remove(segmentFile)
+			if err := os.Remove(segmentFile); err != nil {
+				streamLogger.Errorln(err)
+			} else {
+				streamLogger.Infof("Removed segment %s\n", segmentFile)
+			}
 		}
 	}
 	return nil
@@ -44,27 +49,32 @@ func main() {
 	server := &rtmp.Server{Addr: addr}
 
 	server.HandlePublish = func(conn *rtmp.Conn) {
-		log.Printf("Handling request %s", conn.URL.RequestURI())
+		log.Infof("Handling request %s\n", conn.URL.RequestURI())
 		if conn.URL.Query().Get("key") != key {
-			log.Println("Key mismatch, aborting request")
+			log.Infoln("Key mismatch, aborting request")
 			return
 		}
 
 		streamName := strings.ReplaceAll(conn.URL.Path, "/", "")
 		if streamName == "" {
-			log.Println("Invalid stream name")
+			log.Errorln("Invalid stream name")
 			return
 		}
 
+		streamLogger := log.WithFields(log.Fields{"stream": streamName})
+
 		// create hls playlist
+		playlistFileName := fmt.Sprintf("%s.m3u8", streamName)
 		playlist, err := m3u8.NewMediaPlaylist(5, 10)
 		if err != nil {
-			log.Fatal(err)
+			streamLogger.Errorln(err)
+			return
 		}
 
 		streams, err := conn.Streams()
 		if err != nil {
-			log.Fatalln(err)
+			streamLogger.Errorln(err)
+			return
 		}
 
 		i := 0
@@ -75,13 +85,15 @@ func main() {
 			segmentName := fmt.Sprintf("%s%04d.ts", streamName, i)
 			outFile, err := os.Create(segmentName)
 			if err != nil {
-				log.Fatalln(err)
+				streamLogger.Errorln(err)
+				return
 			}
 			tsMuxer := ts.NewMuxer(outFile)
 
 			// write header
 			if err := tsMuxer.WriteHeader(streams); err != nil {
-				log.Fatalln(err)
+				streamLogger.Errorln(err)
+				return
 			}
 			// write some data
 			var segmentLength time.Duration = 0
@@ -91,51 +103,62 @@ func main() {
 				var packet av.Packet
 				if packet, err = conn.ReadPacket(); err != nil {
 					if err == io.EOF {
-						log.Println("Client disconnected")
+						streamLogger.Infoln("Client disconnected")
 						clientConnected = false
 						break
 					}
-					log.Fatalln(err)
+					streamLogger.Errorln(err)
+					return
 				}
 				if err = tsMuxer.WritePacket(packet); err != nil {
-					log.Fatalln(err)
+					streamLogger.Errorln(err)
+					return
 				}
 				packetLength = packet.Time - lastPacketTime
 				segmentLength += packetLength
 				lastPacketTime = packet.Time
 			}
 			// write trailer
 			if err := tsMuxer.WriteTrailer(); err != nil {
-				log.Fatalln(err)
+				streamLogger.Errorln(err)
+				return
 			}
-			log.Printf("Successfully wrote segment %s\n", segmentName)
+			streamLogger.Infof("Wrote segment %s\n", segmentName)
 
 			// update playlist
 			playlist.Slide(segmentName, segmentLength.Seconds(), "")
-			playlistFile, err := os.Create(fmt.Sprintf("%s.m3u8", streamName))
+			playlistFile, err := os.Create(playlistFileName)
 			if err != nil {
-				log.Fatalln(err)
+				streamLogger.Errorln(err)
+				return
 			}
 			playlistFile.Write(playlist.Encode().Bytes())
 			playlistFile.Close()
 
 			// cleanup segments
-			if err := removeOutdatedSegments(streamName, playlist); err != nil {
-				log.Println(err)
+			if err := removeOutdatedSegments(streamLogger, streamName, playlist); err != nil {
+				streamLogger.Errorln(err)
+				return
 			}
 
 			// increase counter
 			i++
 		}
 
-		// remove all segments
+		// remove all segments; this is probably not a good idea
 		for _, segment := range playlist.Segments {
 			if segment != nil {
 				if err := os.Remove(segment.URI); err != nil {
-					log.Println(err)
+					streamLogger.Errorln(err)
+					return
 				}
 			}
 		}
+		// remove playlist
+		if err := os.Remove(playlistFileName); err != nil {
+			streamLogger.Error(err)
+			return
+		}
 	}
 
 	log.Printf("Listening on %s", server.Addr)