mirror of
https://github.com/lkmio/lkm.git
synced 2025-09-27 03:26:01 +08:00
396 lines
8.9 KiB
Go
396 lines
8.9 KiB
Go
package stream
|
||
|
||
import (
|
||
"fmt"
|
||
"github.com/lkmio/avformat"
|
||
"github.com/lkmio/avformat/collections"
|
||
"github.com/lkmio/lkm/log"
|
||
"net"
|
||
"net/url"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/lkmio/avformat/utils"
|
||
)
|
||
|
||
var (
|
||
StreamEndInfoBride func(source string, tracks []*Track, streams map[TransStreamID]TransStream) *StreamEndInfo
|
||
)
|
||
|
||
// Source 对推流源的封装
|
||
type Source interface {
|
||
// GetID 返回SourceID
|
||
GetID() string
|
||
|
||
SetID(id string)
|
||
|
||
// Input 输入推流数据
|
||
Input(data []byte) error
|
||
|
||
// GetType 返回推流类型
|
||
GetType() SourceType
|
||
|
||
SetType(sourceType SourceType)
|
||
|
||
// OriginTracks 返回所有的推流track
|
||
OriginTracks() []*Track
|
||
|
||
SetState(state SessionState)
|
||
|
||
// Close 关闭Source
|
||
// 关闭推流网络链路, 停止一切封装和转发流以及转码工作
|
||
// 将Sink添加到等待队列
|
||
Close()
|
||
|
||
// IsCompleted 所有推流track是否解析完毕
|
||
IsCompleted() bool
|
||
|
||
Init(receiveQueueSize int)
|
||
|
||
RemoteAddr() string
|
||
|
||
String() string
|
||
|
||
State() SessionState
|
||
|
||
// UrlValues 返回推流url参数
|
||
UrlValues() url.Values
|
||
|
||
// SetUrlValues 设置推流url参数
|
||
SetUrlValues(values url.Values)
|
||
|
||
// PostEvent 切换到主协程执行当前函数
|
||
postEvent(cb func())
|
||
|
||
executeSyncEvent(cb func())
|
||
|
||
// LastPacketTime 返回最近收流时间戳
|
||
LastPacketTime() time.Time
|
||
|
||
SetLastPacketTime(time2 time.Time)
|
||
|
||
IsClosed() bool
|
||
|
||
StreamPipe() chan []byte
|
||
|
||
MainContextEvents() chan func()
|
||
|
||
CreateTime() time.Time
|
||
|
||
SetCreateTime(time time.Time)
|
||
|
||
GetBitrateStatistics() *BitrateStatistics
|
||
|
||
ProbeTimeout()
|
||
|
||
GetTransStreamPublisher() TransStreamPublisher
|
||
}
|
||
|
||
type PublishSource struct {
|
||
ID string
|
||
Type SourceType
|
||
state SessionState
|
||
Conn net.Conn
|
||
|
||
streamPipe *NonBlockingChannel[[]byte] // 推流数据管道
|
||
mainContextEvents chan func() // 切换到主协程执行函数的事件管道
|
||
streamPublisher TransStreamPublisher // 解析出来的AVStream和AVPacket, 交由streamPublisher处理
|
||
|
||
TransDemuxer avformat.Demuxer // 负责从推流协议中解析出AVStream和AVPacket
|
||
originTracks TrackManager // 推流的音视频Streams
|
||
|
||
closed atomic.Bool // 是否已经关闭
|
||
completed atomic.Bool // 推流track是否解析完毕, @see writeHeader 函数中赋值为true
|
||
existVideo bool // 是否存在视频
|
||
|
||
lastPacketTime time.Time // 最近收到推流包的时间
|
||
urlValues url.Values // 推流url携带的参数
|
||
createTime time.Time // source创建时间
|
||
statistics *BitrateStatistics // 码流统计
|
||
streamLogger avformat.OnUnpackStream2FileHandler
|
||
}
|
||
|
||
func (s *PublishSource) SetLastPacketTime(time2 time.Time) {
|
||
s.lastPacketTime = time2
|
||
}
|
||
|
||
func (s *PublishSource) IsClosed() bool {
|
||
return s.closed.Load()
|
||
}
|
||
|
||
func (s *PublishSource) StreamPipe() chan []byte {
|
||
return s.streamPipe.Channel
|
||
}
|
||
|
||
func (s *PublishSource) MainContextEvents() chan func() {
|
||
return s.mainContextEvents
|
||
}
|
||
|
||
func (s *PublishSource) LastPacketTime() time.Time {
|
||
return s.lastPacketTime
|
||
}
|
||
|
||
func (s *PublishSource) GetID() string {
|
||
return s.ID
|
||
}
|
||
|
||
func (s *PublishSource) SetID(id string) {
|
||
s.ID = id
|
||
if s.streamPublisher != nil {
|
||
s.streamPublisher.SetSourceID(id)
|
||
}
|
||
}
|
||
|
||
func (s *PublishSource) Init(receiveQueueSize int) {
|
||
s.SetState(SessionStateHandshakeSuccess)
|
||
|
||
// 初始化事件接收管道
|
||
// -2是为了保证从管道取到流, 到处理完流整个过程安全的, 不会被覆盖
|
||
s.streamPipe = NewNonBlockingChannel[[]byte](receiveQueueSize - 1)
|
||
s.mainContextEvents = make(chan func(), 128)
|
||
s.statistics = NewBitrateStatistics()
|
||
s.streamPublisher = NewTransStreamPublisher(s.ID)
|
||
// 设置探测时长
|
||
s.TransDemuxer.SetProbeDuration(AppConfig.ProbeTimeout)
|
||
}
|
||
|
||
func (s *PublishSource) Input(data []byte) error {
|
||
s.streamPipe.Post(data)
|
||
s.statistics.Input(len(data))
|
||
return nil
|
||
}
|
||
|
||
func (s *PublishSource) OriginTracks() []*Track {
|
||
return s.originTracks.All()
|
||
}
|
||
|
||
func (s *PublishSource) SetState(state SessionState) {
|
||
s.state = state
|
||
}
|
||
|
||
func (s *PublishSource) DoClose() {
|
||
log.Sugar.Debugf("closing the %s source. id: %s. closed flag: %t", s.Type, s.ID, s.closed.Load())
|
||
|
||
if s.closed.Load() {
|
||
return
|
||
}
|
||
|
||
s.closed.Store(true)
|
||
|
||
// 关闭推流源的解复用器, 不再接收数据
|
||
if s.TransDemuxer != nil {
|
||
s.TransDemuxer.Close()
|
||
s.TransDemuxer = nil
|
||
}
|
||
|
||
// 等传输流发布器关闭结束
|
||
s.streamPublisher.close()
|
||
|
||
// 释放解复用器
|
||
// 释放转码器
|
||
// 释放每路转协议流， 将所有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
|
||
}
|
||
|
||
// 异步hook
|
||
go func() {
|
||
if s.Conn != nil {
|
||
_ = s.Conn.Close()
|
||
s.Conn = nil
|
||
}
|
||
|
||
HookPublishDoneEvent(s)
|
||
}()
|
||
}
|
||
|
||
func (s *PublishSource) Close() {
|
||
if s.closed.Load() {
|
||
return
|
||
}
|
||
|
||
// 同步执行, 确保close后, 主协程已经退出, 不会再处理任何推拉流、查询等任何事情.
|
||
s.executeSyncEvent(func() {
|
||
s.DoClose()
|
||
})
|
||
}
|
||
|
||
// 解析完所有track后, 创建各种输出流
|
||
func (s *PublishSource) writeHeader() {
|
||
if s.completed.Load() {
|
||
fmt.Printf("添加Stream失败 Source: %s已经WriteHeader", s.ID)
|
||
return
|
||
}
|
||
|
||
s.completed.Store(true)
|
||
|
||
s.streamPublisher.Post(&StreamEvent{StreamEventTypeTrackCompleted, nil})
|
||
|
||
if len(s.originTracks.All()) == 0 {
|
||
log.Sugar.Errorf("没有一路track, 删除source: %s", s.ID)
|
||
s.DoClose()
|
||
return
|
||
}
|
||
}
|
||
|
||
func (s *PublishSource) IsCompleted() bool {
|
||
return s.completed.Load()
|
||
}
|
||
|
||
// NotTrackAdded 返回该index对应的track是否没有添加
|
||
func (s *PublishSource) NotTrackAdded(index int) bool {
|
||
for _, track := range s.originTracks.All() {
|
||
if track.Stream.Index == index {
|
||
return false
|
||
}
|
||
}
|
||
|
||
return true
|
||
}
|
||
|
||
func (s *PublishSource) OnNewTrack(track avformat.Track) {
|
||
if AppConfig.Debug {
|
||
s.streamLogger.Path = "dump/" + strings.ReplaceAll(s.ID, "/", "_")
|
||
s.streamLogger.OnNewTrack(track)
|
||
}
|
||
|
||
stream := track.GetStream()
|
||
|
||
if s.completed.Load() {
|
||
log.Sugar.Warnf("添加%s track失败,已经WriteHeader. source: %s", stream.MediaType, s.ID)
|
||
return
|
||
} else if !s.NotTrackAdded(stream.Index) {
|
||
log.Sugar.Warnf("添加%s track失败,已经添加索引为%d的track. source: %s", stream.MediaType, stream.Index, s.ID)
|
||
return
|
||
}
|
||
|
||
newTrack := NewTrack(stream, 0, 0)
|
||
s.originTracks.Add(newTrack)
|
||
|
||
if utils.AVMediaTypeVideo == stream.MediaType {
|
||
s.existVideo = true
|
||
}
|
||
|
||
s.streamPublisher.Post(&StreamEvent{StreamEventTypeTrack, newTrack})
|
||
}
|
||
|
||
func (s *PublishSource) OnTrackComplete() {
|
||
if AppConfig.Debug {
|
||
s.streamLogger.OnTrackComplete()
|
||
}
|
||
|
||
s.writeHeader()
|
||
}
|
||
|
||
func (s *PublishSource) OnTrackNotFind() {
|
||
if AppConfig.Debug {
|
||
s.streamLogger.OnTrackNotFind()
|
||
}
|
||
|
||
log.Sugar.Errorf("no tracks found. source id: %s", s.ID)
|
||
}
|
||
|
||
func (s *PublishSource) OnPacket(packet *avformat.AVPacket) {
|
||
if AppConfig.Debug {
|
||
s.streamLogger.OnPacket(packet)
|
||
}
|
||
|
||
// track超时，忽略推流数据
|
||
if s.NotTrackAdded(packet.Index) {
|
||
s.TransDemuxer.DiscardHeadPacket(packet.BufferIndex)
|
||
return
|
||
}
|
||
|
||
packetPtr := collections.NewReferenceCounter(packet)
|
||
packetPtr.Refer() // 引用计数加1
|
||
|
||
packets := s.originTracks.FindWithType(packet.MediaType).Packets
|
||
packets.Add(packetPtr)
|
||
s.streamPublisher.Post(&StreamEvent{StreamEventTypePacket, packetPtr})
|
||
|
||
// 释放未引用的AVPacket
|
||
for packets.Size() > 0 {
|
||
if packets.Get(0).UseCount() > 1 {
|
||
break
|
||
}
|
||
|
||
packets.Remove(0).Release()
|
||
s.TransDemuxer.DiscardHeadPacket(packet.BufferIndex)
|
||
}
|
||
}
|
||
|
||
func (s *PublishSource) GetType() SourceType {
|
||
return s.Type
|
||
}
|
||
|
||
func (s *PublishSource) SetType(sourceType SourceType) {
|
||
s.Type = sourceType
|
||
}
|
||
|
||
func (s *PublishSource) RemoteAddr() string {
|
||
if s.Conn == nil {
|
||
return ""
|
||
}
|
||
|
||
return s.Conn.RemoteAddr().String()
|
||
}
|
||
|
||
func (s *PublishSource) String() string {
|
||
return fmt.Sprintf("source: %s type: %s conn: %s ", s.ID, s.Type.String(), s.RemoteAddr())
|
||
}
|
||
|
||
func (s *PublishSource) State() SessionState {
|
||
return s.state
|
||
}
|
||
|
||
func (s *PublishSource) UrlValues() url.Values {
|
||
return s.urlValues
|
||
}
|
||
|
||
func (s *PublishSource) SetUrlValues(values url.Values) {
|
||
s.urlValues = values
|
||
}
|
||
|
||
func (s *PublishSource) postEvent(cb func()) {
|
||
s.mainContextEvents <- cb
|
||
}
|
||
|
||
func (s *PublishSource) executeSyncEvent(cb func()) {
|
||
group := sync.WaitGroup{}
|
||
group.Add(1)
|
||
|
||
s.postEvent(func() {
|
||
cb()
|
||
group.Done()
|
||
})
|
||
|
||
group.Wait()
|
||
}
|
||
|
||
func (s *PublishSource) CreateTime() time.Time {
|
||
return s.createTime
|
||
}
|
||
|
||
func (s *PublishSource) SetCreateTime(time time.Time) {
|
||
s.createTime = time
|
||
}
|
||
|
||
func (s *PublishSource) GetBitrateStatistics() *BitrateStatistics {
|
||
return s.statistics
|
||
}
|
||
|
||
func (s *PublishSource) ProbeTimeout() {
|
||
if s.TransDemuxer != nil {
|
||
s.TransDemuxer.ProbeComplete()
|
||
}
|
||
}
|
||
|
||
func (s *PublishSource) GetTransStreamPublisher() TransStreamPublisher {
|
||
return s.streamPublisher
|
||
}
|