修复超时计时器未关闭, 造成Source内存泄漏问题

gb28181/source.go

@@ -83,7 +83,6 @@ func (source *BaseGBSource) Input(data []byte) error {
 	}
 
 	packet := rtp.Packet{}
-	packet.Marshal()
 	_ = packet.Unmarshal(data)
 	return source.deMuxerCtx.Input(packet.Payload)
 }
@@ -249,6 +248,11 @@ func (source *BaseGBSource) Close() {
 	}
 
 	source.PublishSource.Close()
+
+	if source.deMuxerCtx != nil {
+		source.deMuxerCtx.Close()
+		source.deMuxerCtx = nil
+	}
 }
 
 func (source *BaseGBSource) SetConn(conn net.Conn) {

gb28181/source_udp.go

@@ -13,15 +13,6 @@ type UDPSource struct {
 	receiveBuffer *stream.ReceiveBuffer
 }
 
-func NewUDPSource() *UDPSource {
-	u := &UDPSource{
-		receiveBuffer: stream.NewReceiveBuffer(1500, stream.ReceiveBufferUdpBlockCount+50),
-	}
-
-	u.jitterBuffer = stream.NewJitterBuffer(u.OnOrderedRtp)
-	return u
-}
-
 func (u *UDPSource) SetupType() SetupType {
 	return SetupUDP
 }
@@ -43,7 +34,19 @@ func (u *UDPSource) InputRtpPacket(pkt *rtp.Packet) error {
 }
 
 func (u *UDPSource) Close() {
+	// 清空剩余在缓冲区的包
 	u.jitterBuffer.Flush()
-	u.jitterBuffer.Close()
+	u.jitterBuffer.SetHandler(nil)
+
 	u.BaseGBSource.Close()
 }
+
+func NewUDPSource() *UDPSource {
+	u := &UDPSource{
+		receiveBuffer: stream.NewReceiveBuffer(1500, stream.ReceiveBufferUdpBlockCount+50),
+	}
+
+	u.jitterBuffer = stream.NewJitterBuffer()
+	u.jitterBuffer.SetHandler(u.OnOrderedRtp)
+	return u
+}

stream/hook_source.go

@@ -24,14 +24,6 @@ func PreparePublishSource(source Source, hook bool) (*http.Response, utils.HookS
 		return nil, utils.HookStateOccupy
 	}
 
-	if AppConfig.Hooks.IsEnableOnReceiveTimeout() && AppConfig.ReceiveTimeout > 0 {
-		StartReceiveDataTimer(source)
-	}
-
-	if AppConfig.Hooks.IsEnableOnIdleTimeout() && AppConfig.IdleTimeout > 0 {
-		StartIdleTimer(source)
-	}
-
 	source.SetCreateTime(time.Now())
 
 	urls := GetStreamPlayUrls(source.GetID())

stream/source.go

@@ -8,6 +8,7 @@ import (
 	"net"
 	"net/url"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/lkmio/avformat/stream"
@@ -105,10 +106,6 @@ type Source interface {
 	// LastStreamEndTime 返回最近结束拉流时间戳
 	LastStreamEndTime() time.Time
 
-	SetReceiveDataTimer(timer *time.Timer)
-
-	SetIdleTimer(timer *time.Timer)
-
 	IsClosed() bool
 
 	StreamPipe() chan []byte
@@ -139,13 +136,11 @@ type PublishSource struct {
 	pktBuffers       [8]collections.MemoryPool // 推流每路的AVPacket缓存, AVPacket的data从该内存池中分配. 在GOP缓存溢出时,释放池中内存.
 	gopBuffer        GOPBuffer                 // GOP缓存, 音频和视频混合使用, 以视频关键帧为界, 缓存第二个视频关键帧时, 释放前一组gop. 如果不存在视频流, 不缓存音频
 
-	closed     bool // source是否已经关闭
+	closed     atomic.Bool // source是否已经关闭
 	completed  bool        // 所有推流track是否解析完毕, @see writeHeader 函数中赋值为true
 	existVideo bool        // 是否存在视频
 
 	probeTimer *time.Timer // track解析超时计时器, 触发时执行@see writeHeader
-	receiveDataTimer *time.Timer // 收流超时计时器
-	idleTimer        *time.Timer // 拉流空闲计时器
 
 	TransStreams     map[TransStreamID]TransStream     // 所有的输出流, 持有Sink
 	sinks            map[SinkID]Sink                   // 保存所有Sink
@@ -167,7 +162,7 @@ func (s *PublishSource) SetLastPacketTime(time2 time.Time) {
 }
 
 func (s *PublishSource) IsClosed() bool {
-	return s.closed
+	return s.closed.Load()
 }
 
 func (s *PublishSource) StreamPipe() chan []byte {
@@ -178,14 +173,6 @@ func (s *PublishSource) MainContextEvents() chan func() {
 	return s.mainContextEvents
 }
 
-func (s *PublishSource) SetReceiveDataTimer(timer *time.Timer) {
-	s.receiveDataTimer = timer
-}
-
-func (s *PublishSource) SetIdleTimer(timer *time.Timer) {
-	s.idleTimer = timer
-}
-
 func (s *PublishSource) LastStreamEndTime() time.Time {
 	return s.lastStreamEndTime
 }
@@ -448,12 +435,12 @@ func (s *PublishSource) doAddSink(sink Sink) bool {
 		return false
 	}
 
-	// 还没准备好推流, 暂不推流
+	// 还没做好准备(rtsp拉流还在协商sdp中), 暂不推流
 	if !sink.IsReady() {
 		return true
 	}
 
-	// TCP拉流开启异步发包
+	// TCP拉流开启异步发包, 一旦出现网络不好的链路, 其余正常链路不受影响.
 	conn, ok := sink.GetConn().(*transport.Conn)
 	if ok && sink.IsTCPStreaming() && transStream.OutStreamBufferCapacity() > 2 {
 		conn.EnableAsyncWriteMode(transStream.OutStreamBufferCapacity() - 2)
@@ -559,11 +546,11 @@ func (s *PublishSource) SetState(state SessionState) {
 func (s *PublishSource) DoClose() {
 	log.Sugar.Debugf("closing the %s source. id: %s. closed flag: %t", s.Type, s.ID, s.closed)
 
-	if s.closed {
+	if s.closed.Load() {
 		return
 	}
 
-	s.closed = true
+	s.closed.Store(true)
 
 	if s.TransDeMuxer != nil {
 		s.TransDeMuxer.Close()
@@ -584,19 +571,11 @@ func (s *PublishSource) DoClose() {
 		s.gopBuffer = nil
 	}
 
-	// 停止所有计时器
+	// 停止track探测计时器
 	if s.probeTimer != nil {
 		s.probeTimer.Stop()
 	}
 
-	if s.receiveDataTimer != nil {
-		s.receiveDataTimer.Stop()
-	}
-
-	if s.idleTimer != nil {
-		s.idleTimer.Stop()
-	}
-
 	// 关闭录制流
 	if s.recordSink != nil {
 		s.recordSink.Close()
@@ -607,7 +586,9 @@ func (s *PublishSource) DoClose() {
 	// 释放每路转协议流， 将所有sink添加到等待队列
 	_, err := SourceManager.Remove(s.ID)
 	if err != nil {
+		// source不存在, 在创建source时, 未添加到manager中, 目前只有1078流会出现这种情况(tcp连接到端口, 没有推流或推流数据无效, 无法定位到手机号, 以至于无法执行PreparePublishSource函数), 将不再处理后续事情.
 		log.Sugar.Errorf("删除源失败 source:%s err:%s", s.ID, err.Error())
+		return
 	}
 
 	// 关闭所有输出流
@@ -665,9 +646,21 @@ func (s *PublishSource) DoClose() {
 }
 
 func (s *PublishSource) Close() {
+	if s.closed.Load() {
+		return
+	}
+
+	// 同步执行, 确保close后, 主协程已经退出, 不会再处理任何推拉流、查询等任何事情.
+	group := sync.WaitGroup{}
+	group.Add(1)
+
 	s.PostEvent(func() {
 		s.DoClose()
+
+		group.Done()
 	})
+
+	group.Wait()
 }
 
 func (s *PublishSource) OnDiscardPacket(packet utils.AVPacket) {

stream/source_utils.go

@@ -221,7 +221,7 @@ func ExtractAudioPacket(codec utils.AVCodecID, extractStream bool, data []byte,
 }
 
 // StartReceiveDataTimer 启动收流超时计时器
-func StartReceiveDataTimer(source Source) {
+func StartReceiveDataTimer(source Source) *time.Timer {
 	utils.Assert(AppConfig.ReceiveTimeout > 0)
 
 	var receiveDataTimer *time.Timer
@@ -244,10 +244,12 @@ func StartReceiveDataTimer(source Source) {
 		// 对精度没要求
 		receiveDataTimer.Reset(time.Duration(AppConfig.ReceiveTimeout))
 	})
+
+	return receiveDataTimer
 }
 
 // StartIdleTimer 启动拉流空闲计时器
-func StartIdleTimer(source Source) {
+func StartIdleTimer(source Source) *time.Timer {
 	utils.Assert(AppConfig.IdleTimeout > 0)
 
 	var idleTimer *time.Timer
@@ -266,17 +268,42 @@ func StartIdleTimer(source Source) {
 
 		idleTimer.Reset(time.Duration(AppConfig.IdleTimeout))
 	})
+
+	return idleTimer
 }
 
 // LoopEvent 循环读取事件
 func LoopEvent(source Source) {
-	for {
-		select {
-		case data := <-source.StreamPipe():
-			if source.IsClosed() {
-				return
+	// 将超时计时器放在此处开启, 方便在退出的时候关闭
+	var receiveTimer *time.Timer
+	var idleTimer *time.Timer
+
+	defer func() {
+		log.Sugar.Debugf("主协程执行结束 source: %s", source.GetID())
+
+		// 关闭计时器
+		if receiveTimer != nil {
+			receiveTimer.Stop()
+		}
+		if idleTimer != nil {
+			idleTimer.Stop()
+		}
+	}()
+
+	// 开启收流超时计时器
+	if AppConfig.Hooks.IsEnableOnReceiveTimeout() && AppConfig.ReceiveTimeout > 0 {
+		receiveTimer = StartReceiveDataTimer(source)
 	}
 
+	// 开启拉流空闲超时计时器
+	if AppConfig.Hooks.IsEnableOnIdleTimeout() && AppConfig.IdleTimeout > 0 {
+		idleTimer = StartIdleTimer(source)
+	}
+
+	for {
+		select {
+		// 读取推流数据
+		case data := <-source.StreamPipe():
 			if AppConfig.ReceiveTimeout > 0 {
 				source.SetLastPacketTime(time.Now())
 			}
@@ -287,19 +314,13 @@ func LoopEvent(source Source) {
 				return
 			}
 
-			if source.IsClosed() {
-				return
-			}
-
 			break
+		// 切换到主协程,执行该函数. 目的是用于无锁化处理推拉流的连接与断开, 推流源断开, 查询推流源信息等事件. 不要做耗时操作, 否则会影响推拉流.
 		case event := <-source.MainContextEvents():
-			if source.IsClosed() {
-				return
-			}
-
 			event()
 
 			if source.IsClosed() {
+				// 处理推流管道剩余的数据?
 				return
 			}