重构输出流转发, TransStream不再持有Sink

2025-09-27 03:26:01 +08:00 · 2024-10-28 19:15:53 +08:00
parent 9090e28077
commit ec707c8dc1
27 changed files with 894 additions and 747 deletions
--- a/rtc/rtc_sink.go
+++ b/rtc/rtc_sink.go
@@ -1,13 +1,16 @@
 package rtc

 import (
+	"fmt"
+	"github.com/lkmio/avformat/utils"
+	"github.com/lkmio/lkm/log"
 	"github.com/lkmio/lkm/stream"
 	"github.com/pion/webrtc/v3"
 	"github.com/pion/webrtc/v3/pkg/media"
 	"time"
 )

-type sink struct {
+type Sink struct {
 	stream.BaseSink

 	offer  string
@@ -20,31 +23,117 @@ type sink struct {
 	cb func(sdp string)
 }

-func NewSink(id stream.SinkID, sourceId string, offer string, cb func(sdp string)) stream.Sink {
-	return &sink{stream.BaseSink{ID: id, SourceID: sourceId, Protocol: stream.TransStreamRtc}, offer, "", nil, nil, webrtc.ICEConnectionStateNew, cb}
-}
+func (s *Sink) StartStreaming(transStream stream.TransStream) error {
+	// 创建PeerConnection
+	var videoTrack *webrtc.TrackLocalStaticSample
+	s.setTrackCount(transStream.TrackCount())

-func (s *sink) setTrackCount(count int) {
+	connection, err := webrtcApi.NewPeerConnection(webrtc.Configuration{})
+	connection.OnICECandidate(func(candidate *webrtc.ICECandidate) {
+
+	})
+
+	tracks := transStream.GetTracks()
+	for index, track := range tracks {
+		var mimeType string
+		var id string
+		if utils.AVCodecIdH264 == track.CodecId() {
+			mimeType = webrtc.MimeTypeH264
+		} else if utils.AVCodecIdH265 == track.CodecId() {
+			mimeType = webrtc.MimeTypeH265
+		} else if utils.AVCodecIdAV1 == track.CodecId() {
+			mimeType = webrtc.MimeTypeAV1
+		} else if utils.AVCodecIdVP8 == track.CodecId() {
+			mimeType = webrtc.MimeTypeVP8
+		} else if utils.AVCodecIdVP9 == track.CodecId() {
+			mimeType = webrtc.MimeTypeVP9
+		} else if utils.AVCodecIdOPUS == track.CodecId() {
+			mimeType = webrtc.MimeTypeOpus
+		} else if utils.AVCodecIdPCMALAW == track.CodecId() {
+			mimeType = webrtc.MimeTypePCMA
+		} else if utils.AVCodecIdPCMMULAW == track.CodecId() {
+			mimeType = webrtc.MimeTypePCMU
+		} else {
+			log.Sugar.Errorf("codec %s not compatible with webrtc", track.CodecId())
+			continue
+		}
+
+		if utils.AVMediaTypeAudio == track.Type() {
+			id = "audio"
+		} else {
+			id = "video"
+		}
+
+		videoTrack, err = webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: mimeType}, id, "pion")
+		if err != nil {
+			panic(err)
+		} else if _, err := connection.AddTransceiverFromTrack(videoTrack, webrtc.RTPTransceiverInit{Direction: webrtc.RTPTransceiverDirectionSendonly}); err != nil {
+			return err
+		} else if _, err = connection.AddTrack(videoTrack); err != nil {
+			return err
+		}
+
+		s.addTrack(index, videoTrack)
+	}
+
+	if len(connection.GetTransceivers()) == 0 {
+		return fmt.Errorf("no track added")
+	} else if err = connection.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: s.offer}); err != nil {
+		return err
+	}
+
+	complete := webrtc.GatheringCompletePromise(connection)
+	answer, err := connection.CreateAnswer(nil)
+	if err != nil {
+		return err
+	} else if err = connection.SetLocalDescription(answer); err != nil {
+		return err
+	}
+
+	<-complete
+	connection.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
+		s.state = state
+		log.Sugar.Infof("ice state:%v sink:%d source:%s", state.String(), s.GetID(), s.SourceID)
+
+		if state > webrtc.ICEConnectionStateDisconnected {
+			log.Sugar.Errorf("webrtc peer断开连接 sink: %v source :%s", s.GetID(), s.SourceID)
+			s.Close()
+		}
+	})
+
+	s.peer = connection
+	// offer的sdp, 应答给http请求
+	s.cb(connection.LocalDescription().SDP)
+	return nil
+}
+func (s *Sink) setTrackCount(count int) {
 	s.tracks = make([]*webrtc.TrackLocalStaticSample, count)
 }

-func (s *sink) addTrack(index int, track *webrtc.TrackLocalStaticSample) error {
+func (s *Sink) addTrack(index int, track *webrtc.TrackLocalStaticSample) error {
 	s.tracks[index] = track
 	return nil
 }

-func (s *sink) SendHeader(data []byte) error {
-	s.cb(string(data))
-	return nil
-}
-
-func (s *sink) input(index int, data []byte, ts uint32) error {
+func (s *Sink) Write(index int, data [][]byte, ts int64) error {
 	if s.tracks[index] == nil {
 		return nil
 	}

-	return s.tracks[index].WriteSample(media.Sample{
-		Data:     data,
-		Duration: time.Duration(ts) * time.Millisecond,
-	})
+	for _, bytes := range data {
+		err := s.tracks[index].WriteSample(media.Sample{
+			Data:     bytes,
+			Duration: time.Duration(ts) * time.Millisecond,
+		})
+
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func NewSink(id stream.SinkID, sourceId string, offer string, cb func(sdp string)) stream.Sink {
+	return &Sink{stream.BaseSink{ID: id, SourceID: sourceId, Protocol: stream.TransStreamRtc, TCPStreaming: false}, offer, "", nil, nil, webrtc.ICEConnectionStateNew, cb}
 }
--- a/rtc/rtc_stream.go
+++ b/rtc/rtc_stream.go
@@ -1,9 +1,7 @@
 package rtc

 import (
-	"fmt"
 	"github.com/lkmio/avformat/utils"
-	"github.com/lkmio/lkm/log"
 	"github.com/lkmio/lkm/stream"
 	"github.com/pion/interceptor"
 	"github.com/pion/webrtc/v3"
@@ -18,119 +16,27 @@ type transStream struct {
 	stream.BaseTransStream
 }

-func (t *transStream) Input(packet utils.AVPacket) error {
-	for _, iSink := range t.Sinks {
-		sink_ := iSink.(*sink)
-		if sink_.state < webrtc.ICEConnectionStateConnected {
-			continue
+func (t *transStream) Input(packet utils.AVPacket) ([][]byte, int64, bool, error) {
+	t.ClearOutStreamBuffer()
+
+	if utils.AVMediaTypeAudio == packet.MediaType() {
+		t.AppendOutStreamBuffer(packet.Data())
+	} else if utils.AVMediaTypeVideo == packet.MediaType() {
+		if packet.KeyFrame() {
+			extra := t.BaseTransStream.Tracks[packet.Index()].CodecParameters().AnnexBExtraData()
+			t.AppendOutStreamBuffer(extra)
 		}

-		if utils.AVMediaTypeAudio == packet.MediaType() {
-			sink_.input(packet.Index(), packet.Data(), uint32(packet.Duration(1000)))
-		} else if utils.AVMediaTypeVideo == packet.MediaType() {
-			if packet.KeyFrame() {
-				extra := t.BaseTransStream.Tracks[packet.Index()].CodecParameters().AnnexBExtraData()
-				sink_.input(packet.Index(), extra, 0)
-			}
-
-			sink_.input(packet.Index(), packet.AnnexBPacketData(t.BaseTransStream.Tracks[packet.Index()]), uint32(packet.Duration(1000)))
-		}
+		t.AppendOutStreamBuffer(packet.Data())
 	}

-	return nil
-}
-
-func (t *transStream) AddSink(sink_ stream.Sink) error {
-	//创建PeerConnection
-	var videoTrack *webrtc.TrackLocalStaticSample
-	rtcSink := sink_.(*sink)
-	rtcSink.setTrackCount(len(t.Tracks))
-	connection, err := webrtcApi.NewPeerConnection(webrtc.Configuration{})
-	connection.OnICECandidate(func(candidate *webrtc.ICECandidate) {
-
-	})
-
-	for index, track := range t.Tracks {
-		var mimeType string
-		var id string
-		if utils.AVCodecIdH264 == track.CodecId() {
-			mimeType = webrtc.MimeTypeH264
-		} else if utils.AVCodecIdH265 == track.CodecId() {
-			mimeType = webrtc.MimeTypeH265
-		} else if utils.AVCodecIdAV1 == track.CodecId() {
-			mimeType = webrtc.MimeTypeAV1
-		} else if utils.AVCodecIdVP8 == track.CodecId() {
-			mimeType = webrtc.MimeTypeVP8
-		} else if utils.AVCodecIdVP9 == track.CodecId() {
-			mimeType = webrtc.MimeTypeVP9
-		} else if utils.AVCodecIdOPUS == track.CodecId() {
-			mimeType = webrtc.MimeTypeOpus
-		} else if utils.AVCodecIdPCMALAW == track.CodecId() {
-			mimeType = webrtc.MimeTypePCMA
-		} else if utils.AVCodecIdPCMMULAW == track.CodecId() {
-			mimeType = webrtc.MimeTypePCMU
-		} else {
-			log.Sugar.Errorf("codec %d not compatible with webrtc", track.CodecId())
-			continue
-		}
-
-		if utils.AVMediaTypeAudio == track.Type() {
-			id = "audio"
-		} else {
-			id = "video"
-		}
-
-		videoTrack, err = webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: mimeType}, id, "pion")
-		if err != nil {
-			panic(err)
-		} else if _, err := connection.AddTransceiverFromTrack(videoTrack, webrtc.RTPTransceiverInit{Direction: webrtc.RTPTransceiverDirectionSendonly}); err != nil {
-			return err
-		} else if _, err = connection.AddTrack(videoTrack); err != nil {
-			return err
-		}
-
-		rtcSink.addTrack(index, videoTrack)
-	}
-
-	if len(connection.GetTransceivers()) == 0 {
-		return fmt.Errorf("no track added")
-	} else if err = connection.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: rtcSink.offer}); err != nil {
-		return err
-	}
-
-	complete := webrtc.GatheringCompletePromise(connection)
-	answer, err := connection.CreateAnswer(nil)
-	if err != nil {
-		return err
-	} else if err = connection.SetLocalDescription(answer); err != nil {
-		return err
-	}
-
-	<-complete
-	connection.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
-		rtcSink.state = state
-		log.Sugar.Infof("ice state:%v sink:%d source:%s", state.String(), rtcSink.GetID(), rtcSink.SourceID)
-
-		if state > webrtc.ICEConnectionStateDisconnected {
-			log.Sugar.Errorf("webrtc peer断开连接 sink:%v source:%s", rtcSink.GetID(), rtcSink.SourceID)
-			rtcSink.Close()
-		}
-	})
-
-	rtcSink.peer = connection
-	rtcSink.SendHeader([]byte(connection.LocalDescription().SDP))
-	return t.BaseTransStream.AddSink(sink_)
+	return t.OutBuffer[:t.OutBufferSize], int64(uint32(packet.Duration(1000))), utils.AVMediaTypeVideo == packet.MediaType() && packet.KeyFrame(), nil
 }

 func (t *transStream) WriteHeader() error {
 	return nil
 }

-func NewTransStream() stream.TransStream {
-	t := &transStream{}
-	return t
-}
-
 func InitConfig() {
 	setting := webrtc.SettingEngine{}
 	var ips []string
@@ -145,11 +51,11 @@ func InitConfig() {
 		panic(err)
 	}

-	//设置公网ip和监听端口
+	// 设置公网ip和监听端口
 	setting.SetICEUDPMux(webrtc.NewICEUDPMux(nil, udpListener))
 	setting.SetNAT1To1IPs(ips, webrtc.ICECandidateTypeHost)

-	//注册音视频编码器
+	// 注册音视频编码器
 	m := &webrtc.MediaEngine{}
 	if err := m.RegisterDefaultCodecs(); err != nil {
 		panic(err)
@@ -163,6 +69,11 @@ func InitConfig() {
 	webrtcApi = webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(i), webrtc.WithSettingEngine(setting))
 }

+func NewTransStream() stream.TransStream {
+	t := &transStream{}
+	return t
+}
+
 func TransStreamFactory(source stream.Source, protocol stream.TransStreamProtocol, streams []utils.AVStream) (stream.TransStream, error) {
 	return NewTransStream(), nil
 }