fix graceful shutdown, add missing srt init/cleanup calls

Author: iSchluff

Diff for: api/server.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"log"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/voc/srtrelay/config"
@@ -15,6 +16,7 @@ import (
 type Server struct {
 	conf      config.APIConfig
 	srtServer srt.Server
+	done      sync.WaitGroup
 }
 
 func NewServer(conf config.APIConfig, srtServer srt.Server) *Server {
@@ -36,10 +38,17 @@ func (s *Server) Listen(ctx context.Context) error {
 		MaxHeaderBytes: 1 << 14,
 	}
 
+	s.done.Add(1)
 	go func() {
-		log.Fatal(serv.ListenAndServe())
+		defer s.done.Done()
+		err := serv.ListenAndServe()
+		if err != nil {
+			log.Println(err)
+		}
 	}()
+	s.done.Add(1)
 	go func() {
+		defer s.done.Done()
 		<-ctx.Done()
 		ctx2, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 		defer cancel()
@@ -49,6 +58,11 @@ func (s *Server) Listen(ctx context.Context) error {
 	return nil
 }
 
+// Wait blocks until listening sockets have been closed
+func (s *Server) Wait() {
+	s.done.Wait()
+}
+
 func (s *Server) HandleStreams(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 	stats := s.srtServer.GetStatistics()

Diff for: main.go

@@ -8,8 +8,8 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
-	"time"
 
+	"github.com/haivision/srtgo"
 	"github.com/voc/srtrelay/api"
 	"github.com/voc/srtrelay/config"
 	"github.com/voc/srtrelay/relay"
@@ -86,17 +86,18 @@ func main() {
 	// setup graceful shutdown
 	ctx, cancel := context.WithCancel(context.Background())
 	handleSignal(ctx, cancel)
-	defer cancel()
 
 	// create server
+	srtgo.InitSRT()
 	srtServer := srt.NewServer(&serverConfig)
 	err = srtServer.Listen(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	var apiServer *api.Server
 	if conf.API.Enabled {
-		apiServer := api.NewServer(conf.API, srtServer)
+		apiServer = api.NewServer(conf.API, srtServer)
 		err := apiServer.Listen(ctx)
 		if err != nil {
 			log.Fatal(err)
@@ -105,6 +106,9 @@ func main() {
 	}
 
 	// Wait for graceful shutdown
-	<-ctx.Done()
-	time.Sleep(200 * time.Millisecond)
+	srtServer.Wait()
+	if apiServer != nil {
+		apiServer.Wait()
+	}
+	srtgo.CleanupSRT()
 }

Diff for: srt/server.go

@@ -41,6 +41,7 @@ type ServerConfig struct {
 // Server is an interface for a srt relay server
 type Server interface {
 	Listen(context.Context) error
+	Wait()
 	Handle(context.Context, *srtgo.SrtSocket, *net.UDPAddr)
 	GetStatistics() []*relay.StreamStatistics
 	GetSocketStatistics() []*SocketStatistics
@@ -53,6 +54,7 @@ type ServerImpl struct {
 
 	mutex sync.Mutex
 	conns map[*srtConn]bool
+	done  sync.WaitGroup
 }
 
 // NewServer creates a server
@@ -97,6 +99,11 @@ func (s *ServerImpl) Listen(ctx context.Context) error {
 	return nil
 }
 
+// Wait blocks until listening sockets have been closed
+func (s *ServerImpl) Wait() {
+	s.done.Wait()
+}
+
 func (s *ServerImpl) listenCallback(socket *srtgo.SrtSocket, version int, addr *net.UDPAddr, idstring string) bool {
 	var streamid stream.StreamID
 
@@ -130,12 +137,16 @@ func (s *ServerImpl) listenAt(ctx context.Context, host string, port uint16) err
 		return fmt.Errorf("Listen failed: %v", err)
 	}
 
+	s.done.Add(1)
 	go func() {
+		defer s.done.Done()
 		<-ctx.Done()
 		sck.Close()
 	}()
 
+	s.done.Add(1)
 	go func() {
+		defer s.done.Done()
 		for {
 			sock, addr, err := sck.Accept()
 			if err != nil {