重构媒体轨道等待逻辑

2025-10-06 00:56:58 +08:00 · 2021-06-14 22:15:25 +08:00
parent 96c29bbfa6
commit a0048a9c9d
11 changed files with 764 additions and 494 deletions
--- a/stream.go
+++ b/stream.go
@@ -9,117 +9,122 @@ import (
 	. "github.com/logrusorgru/aurora"
 )

-// Streams 所有的流集合
-var Streams sync.Map
+type StreamCollection struct {
+	sync.RWMutex
+	m map[string]*Stream
+}

-//FindStream 根据流路径查找流
-func FindStream(streamPath string) *Stream {
-	if s, ok := Streams.Load(streamPath); ok {
-		return s.(*Stream)
+func (sc *StreamCollection) GetStream(streamPath string) *Stream {
+	sc.RLock()
+	defer sc.RUnlock()
+	if s, ok := sc.m[streamPath]; ok {
+		return s
 	}
 	return nil
 }
+func (sc *StreamCollection) Delete(streamPath string) {
+	sc.Lock()
+	delete(sc.m, streamPath)
+	sc.Unlock()
+}

-//GetStream 根据流路径获取流，如果不存在则创建一个新的
-func GetStream(streamPath string) (result *Stream) {
-	item, loaded := Streams.LoadOrStore(streamPath, &Stream{
-		StreamPath: streamPath,
-	})
-	result = item.(*Stream)
-	if !loaded {
-		result.Context, result.cancel = context.WithCancel(context.Background())
-		utils.Print(Green("Stream create:"), BrightCyan(streamPath))
+func (sc *StreamCollection) ToList() (r []*Stream) {
+	sc.RLock()
+	defer sc.RUnlock()
+	for _, s := range sc.m {
+		r = append(r, s)
 	}
 	return
 }

-type TrackWaiter struct {
-	Track
-	*sync.Cond
+func init() {
+	Streams.m = make(map[string]*Stream)
 }

-func (tw *TrackWaiter) Ok(t Track) {
-	tw.Track = t
-	tw.Broadcast()
+// Streams 所有的流集合
+var Streams StreamCollection
+
+//FindStream 根据流路径查找流
+func FindStream(streamPath string) *Stream {
+	return Streams.GetStream(streamPath)
 }

 // Stream 流定义
 type Stream struct {
 	context.Context
-	cancel           context.CancelFunc
-	StreamPath       string
-	Type             string    //流类型，来自发布者
-	StartTime        time.Time //流的创建时间
-	*Publisher       `json:"-"`
-	Subscribers      []*Subscriber // 订阅者
-	VideoTracks      sync.Map
-	AudioTracks      sync.Map
-	OriginVideoTrack *VideoTrack //原始视频轨
-	OriginAudioTrack *AudioTrack //原始音频轨
-	subscribeMutex   sync.Mutex
+	StreamPath     string
+	Type           string        //流类型，来自发布者
+	StartTime      time.Time     //流的创建时间
+	Subscribers    []*Subscriber // 订阅者
+	VideoTracks    Tracks
+	AudioTracks    Tracks
+	AutoUnPublish  bool              //	当无人订阅时自动停止发布
+	Transcoding    map[string]string //转码配置，key：目标编码，value：发布者提供的编码
+	subscribeMutex sync.Mutex
+	timeout        *time.Timer //更新时间用来做超时处理
+	Close          func()      `json:"-"`
+	prePayload     uint32      //需要预拼装ByteStream格式的数据的订阅者数量
 }

-func (r *Stream) SetOriginVT(vt *VideoTrack) {
-	r.OriginVideoTrack = vt
-	switch vt.CodecID {
-	case 7:
-		r.AddVideoTrack("h264", vt)
-	case 12:
-		r.AddVideoTrack("h265", vt)
-	}
-}
-func (r *Stream) SetOriginAT(at *AudioTrack) {
-	r.OriginAudioTrack = at
-	switch at.SoundFormat {
-	case 10:
-		r.AddAudioTrack("aac", at)
-	case 7:
-		r.AddAudioTrack("pcma", at)
-	case 8:
-		r.AddAudioTrack("pcmu", at)
-	}
-}
-func (r *Stream) AddVideoTrack(codec string, vt *VideoTrack) *VideoTrack {
-	vt.Stream = r
-	if actual, loaded := r.VideoTracks.LoadOrStore(codec, &TrackWaiter{vt, sync.NewCond(new(sync.Mutex))}); loaded {
-		actual.(*TrackWaiter).Ok(vt)
-	}
-	return vt
+func (r *Stream) Update() {
+	r.timeout.Reset(config.PublishTimeout)
 }

-func (r *Stream) AddAudioTrack(codec string, at *AudioTrack) *AudioTrack {
-	at.Stream = r
-	if actual, loaded := r.AudioTracks.LoadOrStore(codec, &TrackWaiter{at, sync.NewCond(new(sync.Mutex))}); loaded {
-		actual.(*TrackWaiter).Ok(at)
+// Publish 发布者进行发布操作
+func (r *Stream) Publish() bool {
+	Streams.Lock()
+	defer Streams.Unlock()
+	if _, ok := Streams.m[r.StreamPath]; ok {
+		return false
 	}
-	return at
+	r.VideoTracks.Init()
+	r.AudioTracks.Init()
+	var cancel context.CancelFunc
+	customClose := r.Close
+	r.Context, cancel = context.WithCancel(context.Background())
+	var closeOnce sync.Once
+	r.Close = func() {
+		closeOnce.Do(func() {
+			r.timeout.Stop()
+			if customClose != nil {
+				customClose()
+			}
+			cancel()
+			r.VideoTracks.Dispose()
+			r.AudioTracks.Dispose()
+			utils.Print(Yellow("Stream destoryed :"), BrightCyan(r.StreamPath))
+			Streams.Delete(r.StreamPath)
+			TriggerHook(Hook{HOOK_STREAMCLOSE, r})
+		})
+	}
+	r.StartTime = time.Now()
+	Streams.m[r.StreamPath] = r
+	utils.Print(Green("Stream publish:"), BrightCyan(r.StreamPath))
+	r.timeout = time.AfterFunc(config.PublishTimeout, r.Close)
+	//触发钩子
+	TriggerHook(Hook{HOOK_PUBLISH, r})
+	return true
 }

-func (r *Stream) Close() {
-	r.cancel()
-	// if r.OriginVideoTrack != nil {
-	// 	r.OriginVideoTrack.Buffer.Current.Done()
-	// }
-	// if r.OriginAudioTrack != nil {
-	// 	r.OriginAudioTrack.Buffer.Current.Done()
-	// }
-	r.VideoTracks.Range(func(k, v interface{}) bool {
-		v.(*TrackWaiter).Broadcast()
-		if v.(*TrackWaiter).Track != nil {
-			v.(*TrackWaiter).Track.Dispose()
-		}
-		return true
-	})
-	r.AudioTracks.Range(func(k, v interface{}) bool {
-		v.(*TrackWaiter).Broadcast()
-		if v.(*TrackWaiter).Track != nil {
-			v.(*TrackWaiter).Track.Dispose()
-		}
-		return true
-	})
-	utils.Print(Yellow("Stream destoryed :"), BrightCyan(r.StreamPath))
-	Streams.Delete(r.StreamPath)
-	TriggerHook(Hook{HOOK_STREAMCLOSE, r})
+func (r *Stream) WaitVideoTrack(codecs ...string) *VideoTrack {
+	if !config.EnableVideo {
+		return nil
+	}
+	if track := r.VideoTracks.WaitTrack(codecs...); track != nil {
+		return track.(*VideoTrack)
+	}
+	return nil
+}
+
+// TODO: 触发转码逻辑
+func (r *Stream) WaitAudioTrack(codecs ...string) *AudioTrack {
+	if !config.EnableAudio {
+		return nil
+	}
+	if track := r.AudioTracks.WaitTrack(codecs...); track != nil {
+		return track.(*AudioTrack)
+	}
+	return nil
 }

 //Subscribe 订阅流
@@ -129,6 +134,9 @@ func (r *Stream) Subscribe(s *Subscriber) {
 		utils.Print(Sprintf(Yellow("subscribe :%s %s,to Stream %s"), Blue(s.Type), Cyan(s.ID), BrightCyan(r.StreamPath)))
 		s.Context, s.cancel = context.WithCancel(r)
 		r.subscribeMutex.Lock()
+		if s.ByteStreamFormat {
+			r.prePayload++
+		}
 		r.Subscribers = append(r.Subscribers, s)
 		r.subscribeMutex.Unlock()
 		utils.Print(Sprintf(Yellow("%s subscriber %s added remains:%d"), BrightCyan(r.StreamPath), Cyan(s.ID), Blue(len(r.Subscribers))))
@@ -141,12 +149,15 @@ func (r *Stream) UnSubscribe(s *Subscriber) {
 	if r.Err() == nil {
 		var deleted bool
 		r.subscribeMutex.Lock()
+		if s.ByteStreamFormat {
+			r.prePayload--
+		}
 		r.Subscribers, deleted = DeleteSliceItem_Subscriber(r.Subscribers, s)
 		r.subscribeMutex.Unlock()
 		if deleted {
 			utils.Print(Sprintf(Yellow("%s subscriber %s removed remains:%d"), BrightCyan(r.StreamPath), Cyan(s.ID), Blue(len(r.Subscribers))))
 			TriggerHook(Hook{HOOK_UNSUBSCRIBE, s})
-			if len(r.Subscribers) == 0 && (r.Publisher == nil || r.Publisher.AutoUnPublish) {
+			if len(r.Subscribers) == 0 && r.AutoUnPublish {
 				r.Close()
 			}
 		}