修复rtmp推流结束后不释放内存问题

rtmp/rtmp_session.go

@@ -61,7 +61,7 @@ func (s *Session) OnPlay(app, stream_ string, response chan utils.HookState) {
 	streamName, values := stream.ParseUrl(stream_)
 
 	sourceId := s.generateSourceId(app, streamName)
-	sink := NewSink(stream.GenerateSinkId(s.conn.RemoteAddr()), sourceId, s.conn)
+	sink := NewSink(stream.GenerateSinkId(s.conn.RemoteAddr()), sourceId, s.conn, s.stack)
 	sink.SetUrlValues(values)
 
 	log.Sugar.Infof("rtmp onplay app:%s stream:%s sink:%v conn:%s", app, stream_, sink.Id(), s.conn.RemoteAddr().String())
@@ -109,7 +109,7 @@ func (s *Session) Close() {
 			s.receiveBuffer = nil
 		}
 	} else {
-		sink := s.handle.(stream.Sink)
+		sink := s.handle.(*Sink)
 		log.Sugar.Infof("rtmp拉流结束 %s", sink.PrintInfo())
 		sink.Close()
 	}

rtmp/rtmp_sink.go

@@ -1,11 +1,33 @@
 package rtmp
 
 import (
+	"github.com/yangjiechina/avformat/librtmp"
 	"github.com/yangjiechina/avformat/utils"
 	"github.com/yangjiechina/lkm/stream"
 	"net"
 )
 
-func NewSink(id stream.SinkId, sourceId string, conn net.Conn) stream.Sink {
-	return &stream.BaseSink{Id_: id, SourceId_: sourceId, State_: stream.SessionStateCreate, Protocol_: stream.ProtocolRtmp, Conn: conn, DesiredAudioCodecId_: utils.AVCodecIdNONE, DesiredVideoCodecId_: utils.AVCodecIdNONE}
+type Sink struct {
+	stream.BaseSink
+	stack *librtmp.Stack
+}
+
+func (s *Sink) Start() {
+	_ = s.stack.SendStreamBeginChunk(s.Conn)
+}
+
+func (s *Sink) Flush() {
+	_ = s.stack.SendStreamEOFChunk(s.Conn)
+}
+
+func (s *Sink) Close() {
+	s.stack = nil
+	s.BaseSink.Close()
+}
+
+func NewSink(id stream.SinkId, sourceId string, conn net.Conn, stack *librtmp.Stack) stream.Sink {
+	return &Sink{
+		BaseSink: stream.BaseSink{Id_: id, SourceId_: sourceId, State_: stream.SessionStateCreate, Protocol_: stream.ProtocolRtmp, Conn: conn, DesiredAudioCodecId_: utils.AVCodecIdNONE, DesiredVideoCodecId_: utils.AVCodecIdNONE},
+		stack:    stack,
+	}
 }

rtmp/rtmp_stream.go

@@ -18,7 +18,11 @@ type transStream struct {
 	muxer      libflv.Muxer
 	audioChunk librtmp.Chunk
 	videoChunk librtmp.Chunk
-	mwBuffer   stream.MergeWritingBuffer
+
+	//合并写内存泄露问题: 推流结束后, mwBuffer的data一直释放不掉, 只有拉流全部断开之后, 才会释放该内存.
+	//起初怀疑是代码层哪儿有问题, 但是测试发现如果将合并写切片再拷贝一次发送 给sink, 推流结束后，mwBuffer的data内存块释放没问题, 只有拷贝的内存块未释放. 所以排除了代码层造成内存泄露的可能性.
+	//看来是conn在write后还会持有data. 查阅代码发现, 的确如此. 向fd发送数据前buffer会引用data, 但是后续没有赋值为nil, 取消引用. https://github.com/golang/go/blob/d38f1d13fa413436d38d86fe86d6a146be44bb84/src/internal/poll/fd_windows.go#L694
+	mwBuffer stream.MergeWritingBuffer
 }
 
 func (t *transStream) Input(packet utils.AVPacket) error {
@@ -179,7 +183,6 @@ func (t *transStream) Close() error {
 	if len(segment) > 0 {
 		t.SendPacket(segment)
 	}
-
 	return nil
 }
 

stream/sink.go

@@ -60,6 +60,10 @@ type Sink interface {
 	UrlValues() url.Values
 
 	SetUrlValues(values url.Values)
+
+	Start()
+
+	Flush()
 }
 
 // GenerateSinkId 根据网络地址生成SinkId IPV4使用一个uint64, IPV6使用String
@@ -233,3 +237,10 @@ func (s *BaseSink) UrlValues() url.Values {
 func (s *BaseSink) SetUrlValues(values url.Values) {
 	s.urlValues = values
 }
+
+func (s *BaseSink) Start() {
+
+}
+
+func (s *BaseSink) Flush() {
+}

stream/source.go

@@ -499,6 +499,10 @@ func (s *PublishSource) doClose() {
 					AddSinkToWaitingQueue(s.Id_, sink)
 				}
 			}
+
+			if SessionStateClose != sink.State() {
+				sink.Flush()
+			}
 		})
 	}
 

stream/trans_stream.go

@@ -56,6 +56,7 @@ func (t *BaseTransStream) AddTrack(stream utils.AVStream) error {
 
 func (t *BaseTransStream) AddSink(sink Sink) error {
 	t.Sinks[sink.Id()] = sink
+	sink.Start()
 	return nil
 }