rtmp client: speed up closing

internal/connrtmp/conn.go

@@ -73,11 +73,12 @@ type Conn struct {
 	pathMan             PathMan
 	parent              Parent
 
-	// read mode
+	path       readpublisher.Path
-	ringBuffer *ringbuffer.RingBuffer
+	ringBuffer *ringbuffer.RingBuffer // read
 
 	// in
 	terminate       chan struct{}
+	parentTerminate chan struct{}
 }
 
 // New allocates a Conn.
@@ -106,7 +107,8 @@ func New(
 		conn:                rtmp.NewServerConn(nconn),
 		pathMan:             pathMan,
 		parent:              parent,
-		terminate:           make(chan struct{}),
+		terminate:           make(chan struct{}, 1),
+		parentTerminate:     make(chan struct{}),
 	}
 
 	c.log(logger.Info, "opened")
@@ -117,15 +119,18 @@ func New(
 	return c
 }
 
-// Close closes a Conn.
+// ParentClose closes a Conn.
-func (c *Conn) Close() {
+func (c *Conn) ParentClose() {
 	c.log(logger.Info, "closed")
-	close(c.terminate)
+	close(c.parentTerminate)
 }
 
-// RequestClose closes a Conn.
+// Close closes a Conn.
-func (c *Conn) RequestClose() {
+func (c *Conn) Close() {
-	c.parent.OnConnClose(c)
+	select {
+	case c.terminate <- struct{}{}:
+	default:
+	}
 }
 
 // IsReadPublisher implements readpublisher.ReadPublisher.
@@ -154,34 +159,50 @@ func (c *Conn) run() {
 		defer onConnectCmd.Close()
 	}
 
+	c.ringBuffer = ringbuffer.New(uint64(c.readBufferCount))
+
+	connErr := make(chan error)
+	go func() {
+		connErr <- func() error {
 			c.conn.NetConn().SetReadDeadline(time.Now().Add(c.readTimeout))
 			c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
 			err := c.conn.ServerHandshake()
 			if err != nil {
+				return err
+			}
+
+			if c.conn.IsPublishing() {
+				return c.runPublish()
+			}
+			return c.runRead()
+		}()
+	}()
+
+	select {
+	case err := <-connErr:
+		if err != io.EOF {
 			c.log(logger.Info, "ERR: %s", err)
+		}
+
 		c.conn.NetConn().Close()
 
-		c.parent.OnConnClose(c)
+	case <-c.terminate:
-		<-c.terminate
+		c.ringBuffer.Close()
-		return
+		c.conn.NetConn().Close()
+		<-connErr
 	}
 
-	if c.conn.IsPublishing() {
+	if c.path != nil {
-		c.runPublish()
+		res := make(chan struct{})
-	} else {
+		c.path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-		c.runRead()
+		<-res
 	}
-}
 
-func (c *Conn) runRead() {
+	c.parent.OnConnClose(c)
-	var path readpublisher.Path
+	<-c.parentTerminate
-	var videoTrack *gortsplib.Track
+}
-	var h264Decoder *rtph264.Decoder
-	var audioTrack *gortsplib.Track
-	var audioClockRate int
-	var aacDecoder *rtpaac.Decoder
 
-	err := func() error {
+func (c *Conn) runRead() error {
 	pathName, query := pathNameAndQuery(c.conn.URL())
 
 	sres := make(chan readpublisher.SetupPlayRes)
@@ -198,15 +219,18 @@ func (c *Conn) runRead() {
 	if res.Err != nil {
 		if _, ok := res.Err.(readpublisher.ErrAuthCritical); ok {
 			// wait some seconds to stop brute force attacks
-				select {
+			<-time.After(pauseAfterAuthError)
-				case <-time.After(pauseAfterAuthError):
-				case <-c.terminate:
-				}
 		}
 		return res.Err
 	}
 
-		path = res.Path
+	c.path = res.Path
+
+	var videoTrack *gortsplib.Track
+	var h264Decoder *rtph264.Decoder
+	var audioTrack *gortsplib.Track
+	var audioClockRate int
+	var aacDecoder *rtpaac.Decoder
 
 	for i, t := range res.Tracks {
 		if t.IsH264() {
@@ -235,37 +259,15 @@ func (c *Conn) runRead() {
 	c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
 	c.conn.WriteMetadata(videoTrack, audioTrack)
 
-		return nil
-	}()
-	if err != nil {
-		c.conn.NetConn().Close()
-		c.log(logger.Info, "ERR: %v", err)
-
-		if path != nil {
-			res := make(chan struct{})
-			path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-			<-res
-		}
-
-		c.parent.OnConnClose(c)
-		<-c.terminate
-		return
-	}
-
-	c.ringBuffer = ringbuffer.New(uint64(c.readBufferCount))
-
 	pres := make(chan readpublisher.PlayRes)
-	path.OnReadPublisherPlay(readpublisher.PlayReq{c, pres}) //nolint:govet
+	c.path.OnReadPublisherPlay(readpublisher.PlayReq{c, pres}) //nolint:govet
 	<-pres
 
-	c.log(logger.Info, "is reading from path '%s'", path.Name())
+	c.log(logger.Info, "is reading from path '%s'", c.path.Name())
 
 	// disable read deadline
 	c.conn.NetConn().SetReadDeadline(time.Time{})
 
-	writerDone := make(chan error)
-	go func() {
-		writerDone <- func() error {
 	var videoBuf [][]byte
 	videoDTSEst := h264.NewDTSEstimator()
 
@@ -344,54 +346,19 @@ func (c *Conn) runRead() {
 			}
 		}
 	}
-		}()
-	}()
-
-	select {
-	case err := <-writerDone:
-		c.conn.NetConn().Close()
-
-		if err != io.EOF {
-			c.log(logger.Info, "ERR: %s", err)
-		}
-
-		res := make(chan struct{})
-		path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-		<-res
-
-		c.parent.OnConnClose(c)
-		<-c.terminate
-
-	case <-c.terminate:
-		res := make(chan struct{})
-		path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-		<-res
-
-		c.ringBuffer.Close()
-		c.conn.NetConn().Close()
-		<-writerDone
-	}
 }
 
-func (c *Conn) runPublish() {
+func (c *Conn) runPublish() error {
-	var videoTrack *gortsplib.Track
-	var audioTrack *gortsplib.Track
-	var err error
-	var tracks gortsplib.Tracks
-	var h264Encoder *rtph264.Encoder
-	var aacEncoder *rtpaac.Encoder
-	var path readpublisher.Path
-
-	setupDone := make(chan struct{})
-	go func() {
-		defer close(setupDone)
-		err = func() error {
 	c.conn.NetConn().SetReadDeadline(time.Now().Add(c.readTimeout))
-			videoTrack, audioTrack, err = c.conn.ReadMetadata()
+	videoTrack, audioTrack, err := c.conn.ReadMetadata()
 	if err != nil {
 		return err
 	}
 
+	var tracks gortsplib.Tracks
+	var h264Encoder *rtph264.Encoder
+	var aacEncoder *rtpaac.Encoder
+
 	if videoTrack != nil {
 		h264Encoder = rtph264.NewEncoder(96, nil, nil, nil)
 		tracks = append(tracks, videoTrack)
@@ -425,51 +392,26 @@ func (c *Conn) runPublish() {
 	if res.Err != nil {
 		if _, ok := res.Err.(readpublisher.ErrAuthCritical); ok {
 			// wait some seconds to stop brute force attacks
-					select {
+			<-time.After(pauseAfterAuthError)
-					case <-time.After(pauseAfterAuthError):
-					case <-c.terminate:
-					}
 		}
 		return res.Err
 	}
 
-			path = res.Path
+	c.path = res.Path
-			return nil
-		}()
-	}()
-
-	select {
-	case <-setupDone:
-	case <-c.terminate:
-		c.conn.NetConn().Close()
-		<-setupDone
-	}
-
-	if err != nil {
-		c.conn.NetConn().Close()
-		c.log(logger.Info, "ERR: %s", err)
-
-		c.parent.OnConnClose(c)
-		<-c.terminate
-		return
-	}
 
 	// disable write deadline
 	c.conn.NetConn().SetWriteDeadline(time.Time{})
 
-	readerDone := make(chan error)
+	rresc := make(chan readpublisher.RecordRes)
-	go func() {
+	c.path.OnReadPublisherRecord(readpublisher.RecordReq{Author: c, Res: rresc})
-		readerDone <- func() error {
+	rres := <-rresc
-			resc := make(chan readpublisher.RecordRes)
-			path.OnReadPublisherRecord(readpublisher.RecordReq{Author: c, Res: resc})
-			res := <-resc
 
-			if res.Err != nil {
+	if rres.Err != nil {
-				return res.Err
+		return rres.Err
 	}
 
 	c.log(logger.Info, "is publishing to path '%s', %d %s",
-				path.Name(),
+		c.path.Name(),
 		len(tracks),
 		func() string {
 			if len(tracks) == 1 {
@@ -479,11 +421,11 @@ func (c *Conn) runPublish() {
 		}())
 
 	var onPublishCmd *externalcmd.Cmd
-			if path.Conf().RunOnPublish != "" {
+	if c.path.Conf().RunOnPublish != "" {
 		_, port, _ := net.SplitHostPort(c.rtspAddress)
-				onPublishCmd = externalcmd.New(path.Conf().RunOnPublish,
+		onPublishCmd = externalcmd.New(c.path.Conf().RunOnPublish,
-					path.Conf().RunOnPublishRestart, externalcmd.Environment{
+			c.path.Conf().RunOnPublishRestart, externalcmd.Environment{
-						Path: path.Name(),
+				Path: c.path.Name(),
 				Port: port,
 			})
 	}
@@ -492,14 +434,14 @@ func (c *Conn) runPublish() {
 		if path.Conf().RunOnPublish != "" {
 			onPublishCmd.Close()
 		}
-			}(path)
+	}(c.path)
 
-			rtcpSenders := rtcpsenderset.New(tracks, res.SP.OnFrame)
+	rtcpSenders := rtcpsenderset.New(tracks, rres.SP.OnFrame)
 	defer rtcpSenders.Close()
 
 	onFrame := func(trackID int, payload []byte) {
 		rtcpSenders.OnFrame(trackID, gortsplib.StreamTypeRTP, payload)
-				res.SP.OnFrame(trackID, gortsplib.StreamTypeRTP, payload)
+		rres.SP.OnFrame(trackID, gortsplib.StreamTypeRTP, payload)
 	}
 
 	for {
@@ -564,34 +506,6 @@ func (c *Conn) runPublish() {
 			return fmt.Errorf("ERR: unexpected packet: %v", pkt.Type)
 		}
 	}
-		}()
-	}()
-
-	select {
-	case err := <-readerDone:
-		c.conn.NetConn().Close()
-
-		if err != io.EOF {
-			c.log(logger.Info, "ERR: %s", err)
-		}
-
-		res := make(chan struct{})
-		path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-		<-res
-		path = nil
-
-		c.parent.OnConnClose(c)
-		<-c.terminate
-
-	case <-c.terminate:
-		c.conn.NetConn().Close()
-		<-readerDone
-
-		res := make(chan struct{})
-		path.OnReadPublisherRemove(readpublisher.RemoveReq{c, res}) //nolint:govet
-		<-res
-		path = nil
-	}
 }
 
 func (c *Conn) validateCredentials(

internal/connrtsp/conn.go

@@ -95,8 +95,8 @@ func New(
 	return c
 }
 
-// Close closes a Conn.
+// ParentClose closes a Conn.
-func (c *Conn) Close(err error) {
+func (c *Conn) ParentClose(err error) {
 	if err != io.EOF && !isTeardownErr(err) && !isTerminatedErr(err) {
 		c.log(logger.Info, "ERR: %v", err)
 	}

internal/converterhls/converter.go

@@ -182,14 +182,14 @@ func New(
 	return c
 }
 
-// Close closes a Converter.
+// ParentClose closes a Converter.
-func (c *Converter) Close() {
+func (c *Converter) ParentClose() {
 	c.log(logger.Info, "closed")
 	close(c.terminate)
 }
 
-// RequestClose closes a Converter.
+// Close closes a Converter.
-func (c *Converter) RequestClose() {
+func (c *Converter) Close() {
 	c.parent.OnConverterClose(c)
 }
 

internal/path/path.go

@@ -309,7 +309,7 @@ outer:
 			case readPublisherStateRecord:
 				atomic.AddInt64(pa.stats.CountPublishers, -1)
 			}
-			c.RequestClose()
+			c.Close()
 		}
 	}
 	pa.readPublishersWg.Wait()
@@ -470,7 +470,7 @@ func (pa *Path) removeReadPublisher(c readpublisher.ReadPublisher) {
 		for oc, state := range pa.readPublishers {
 			if state != readPublisherStatePreRemove {
 				pa.removeReadPublisher(oc)
-				oc.RequestClose()
+				oc.Close()
 			}
 		}
 	}
@@ -510,7 +510,7 @@ func (pa *Path) onSourceSetNotReady() {
 	for c, state := range pa.readPublishers {
 		if c != pa.source && state != readPublisherStatePreRemove {
 			pa.removeReadPublisher(c)
-			c.RequestClose()
+			c.Close()
 		}
 	}
 }
@@ -661,10 +661,10 @@ func (pa *Path) onReadPublisherAnnounce(req readpublisher.AnnounceReq) {
 			return
 		}
 
-		pa.Log(logger.Info, "disconnecting existing publisher")
+		pa.Log(logger.Info, "closing existing publisher")
 		curPublisher := pa.source.(readpublisher.ReadPublisher)
 		pa.removeReadPublisher(curPublisher)
-		curPublisher.RequestClose()
+		curPublisher.Close()
 
 		// prevent path closure
 		if pa.closeTimerStarted {

internal/readpublisher/readpublisher.go

@@ -58,7 +58,6 @@ type ReadPublisher interface {
 	IsReadPublisher()
 	IsSource()
 	Close()
-	RequestClose()
 	OnFrame(int, gortsplib.StreamType, []byte)
 }
 

internal/serverhls/server.go

@@ -203,7 +203,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func (s *Server) doConverterClose(c *converterhls.Converter) {
 	delete(s.converters, c.PathName())
-	c.Close()
+	c.ParentClose()
 }
 
 // OnConverterClose is called by converterhls.Converter.

internal/serverrtmp/server.go

@@ -183,6 +183,7 @@ outer:
 
 func (s *Server) doConnClose(c *connrtmp.Conn) {
 	delete(s.conns, c)
+	c.ParentClose()
 	c.Close()
 }
 

internal/serverrtsp/server.go

@@ -221,7 +221,7 @@ func (s *Server) OnConnClose(ctx *gortsplib.ServerHandlerOnConnCloseCtx) {
 	delete(s.conns, ctx.Conn)
 	s.mutex.Unlock()
 
-	c.Close(ctx.Error)
+	c.ParentClose(ctx.Error)
 }
 
 // OnRequest implements gortsplib.ServerHandlerOnRequest.
@@ -270,7 +270,7 @@ func (s *Server) OnSessionClose(ctx *gortsplib.ServerHandlerOnSessionCloseCtx) {
 	delete(s.sessions, ctx.Session)
 	s.mutex.Unlock()
 
-	se.Close()
+	se.ParentClose()
 }
 
 // OnDescribe implements gortsplib.ServerHandlerOnDescribe.

internal/sessionrtsp/session.go

@@ -75,8 +75,8 @@ func New(
 	return s
 }
 
-// Close closes a Session.
+// ParentClose closes a Session.
-func (s *Session) Close() {
+func (s *Session) ParentClose() {
 	switch s.ss.State() {
 	case gortsplib.ServerSessionStatePlay:
 		if s.onReadCmd != nil {
@@ -99,8 +99,8 @@ func (s *Session) Close() {
 	s.log(logger.Info, "closed")
 }
 
-// RequestClose closes a Session.
+// Close closes a Session.
-func (s *Session) RequestClose() {
+func (s *Session) Close() {
 	s.ss.Close()
 }