mirror of
https://github.com/lkmio/lkm.git
synced 2025-09-27 11:32:26 +08:00
294 lines
7.0 KiB
Go
294 lines
7.0 KiB
Go
package rtmp
|
||
|
||
import (
|
||
"github.com/yangjiechina/avformat/libflv"
|
||
"github.com/yangjiechina/avformat/librtmp"
|
||
"github.com/yangjiechina/avformat/utils"
|
||
"github.com/yangjiechina/live-server/stream"
|
||
)
|
||
|
||
type TransStream struct {
|
||
stream.TransStreamImpl
|
||
chunkSize int
|
||
//sequence header
|
||
header []byte
|
||
headerSize int
|
||
muxer *libflv.Muxer
|
||
|
||
//只存在音频流
|
||
onlyAudio bool
|
||
audioChunk librtmp.Chunk
|
||
videoChunk librtmp.Chunk
|
||
|
||
//只需要缓存一组GOP+第2组GOP的第一个合并写切片
|
||
//当缓存到第2组GOP的第二个合并写切片时，将上一个GOP缓存释放掉
|
||
//使用2块内存池，分别缓存2个GOP，保证内存连续，一次发送
|
||
//不开启GOP缓存和只有音频包的情况下，创建使用一个MemoryPool
|
||
memoryPool [2]stream.MemoryPool
|
||
|
||
//当前合并写切片的缓存时长
|
||
segmentDuration int
|
||
//当前合并写切片位于memoryPool的开始偏移量
|
||
segmentOffset int
|
||
//前一个包的时间戳
|
||
prePacketTS int64
|
||
|
||
firstVideoPacket bool
|
||
|
||
//发送未完整切片的Sinks
|
||
//当AddSink时，还未缓存到一组切片，有多少先发多少. 后续切片未满之前的生成的rtmp包都将直接发送给sink.
|
||
//只要满了一组切片后，这些sink都不单独发包, 统一发送切片.
|
||
incompleteSinks []stream.ISink
|
||
}
|
||
|
||
func NewTransStream(chunkSize int) stream.ITransStream {
|
||
transStream := &TransStream{chunkSize: chunkSize, TransStreamImpl: stream.TransStreamImpl{Sinks: make(map[stream.SinkId]stream.ISink, 64)}}
|
||
return transStream
|
||
}
|
||
|
||
func (t *TransStream) Input(packet utils.AVPacket) {
|
||
utils.Assert(t.TransStreamImpl.Completed)
|
||
|
||
var data []byte
|
||
var chunk *librtmp.Chunk
|
||
var videoPkt bool
|
||
var videoKey bool
|
||
var length int
|
||
//rtmp chunk消息体的数据大小
|
||
var payloadSize int
|
||
|
||
if utils.AVMediaTypeAudio == packet.MediaType() {
|
||
data = packet.Data()
|
||
length = len(data)
|
||
chunk = &t.audioChunk
|
||
payloadSize += 2 + length
|
||
} else if utils.AVMediaTypeVideo == packet.MediaType() {
|
||
//首帧必须要视频关键帧
|
||
if !t.firstVideoPacket {
|
||
if !packet.KeyFrame() {
|
||
return
|
||
}
|
||
|
||
t.firstVideoPacket = true
|
||
}
|
||
|
||
videoPkt = true
|
||
videoKey = packet.KeyFrame()
|
||
data = packet.AVCCPacketData()
|
||
length = len(data)
|
||
chunk = &t.videoChunk
|
||
payloadSize += 5 + length
|
||
}
|
||
|
||
//即不开启GOP缓存又不合并发送. 直接使用AVPacket的预留头封装发送
|
||
if !stream.AppConfig.GOPCache || t.onlyAudio {
|
||
//首帧视频帧必须要关键帧
|
||
return
|
||
}
|
||
|
||
if videoKey {
|
||
tmp := t.memoryPool[0]
|
||
head, _ := tmp.Data()
|
||
if len(head) > t.segmentOffset {
|
||
for _, sink := range t.Sinks {
|
||
sink.Input(head[t.segmentOffset:])
|
||
}
|
||
}
|
||
|
||
t.memoryPool[0].Clear()
|
||
//交替使用缓存
|
||
t.memoryPool[0] = t.memoryPool[1]
|
||
t.memoryPool[1] = tmp
|
||
|
||
t.segmentDuration = 0
|
||
t.segmentOffset = 0
|
||
}
|
||
|
||
//分配内存
|
||
t.memoryPool[0].Mark()
|
||
allocate := t.memoryPool[0].Allocate(12 + payloadSize + (payloadSize / t.chunkSize))
|
||
|
||
//写chunk头
|
||
chunk.Length = payloadSize
|
||
chunk.Timestamp = uint32(packet.Pts())
|
||
n := chunk.ToBytes(allocate)
|
||
utils.Assert(n == 12)
|
||
|
||
//写flv
|
||
ct := packet.Pts() - packet.Dts()
|
||
if videoPkt {
|
||
n += t.muxer.WriteVideoData(allocate[12:], uint32(ct), packet.KeyFrame(), false)
|
||
} else {
|
||
n += t.muxer.WriteAudioData(allocate[12:], false)
|
||
}
|
||
|
||
first := true
|
||
for length > 0 {
|
||
var min int
|
||
if first {
|
||
min = utils.MinInt(length, t.chunkSize-5)
|
||
first = false
|
||
} else {
|
||
min = utils.MinInt(length, t.chunkSize)
|
||
}
|
||
|
||
copy(allocate[n:], data[:min])
|
||
n += min
|
||
|
||
length -= min
|
||
data = data[min:]
|
||
|
||
//写一个ChunkType3用作分割
|
||
if length > 0 {
|
||
if videoPkt {
|
||
allocate[n] = (0x3 << 6) | byte(librtmp.ChunkStreamIdVideo)
|
||
} else {
|
||
allocate[n] = (0x3 << 6) | byte(librtmp.ChunkStreamIdAudio)
|
||
}
|
||
n++
|
||
}
|
||
}
|
||
|
||
rtmpData := t.memoryPool[0].Fetch()[:n]
|
||
t.segmentDuration += int(packet.Pts() - t.prePacketTS)
|
||
t.prePacketTS = packet.Pts()
|
||
|
||
//给不完整切片的Sink补齐包
|
||
if len(t.incompleteSinks) > 0 {
|
||
for _, sink := range t.incompleteSinks {
|
||
sink.Input(rtmpData)
|
||
}
|
||
|
||
if t.segmentDuration >= stream.AppConfig.MergeWriteLatency {
|
||
head, tail := t.memoryPool[0].Data()
|
||
utils.Assert(len(tail) == 0)
|
||
|
||
t.segmentOffset = len(head)
|
||
t.segmentDuration = 0
|
||
t.incompleteSinks = nil
|
||
}
|
||
|
||
return
|
||
}
|
||
|
||
if t.segmentDuration < stream.AppConfig.MergeWriteLatency {
|
||
return
|
||
}
|
||
|
||
head, tail := t.memoryPool[0].Data()
|
||
utils.Assert(len(tail) == 0)
|
||
for _, sink := range t.Sinks {
|
||
sink.Input(head[t.segmentOffset:])
|
||
}
|
||
|
||
t.segmentOffset = len(head)
|
||
t.segmentDuration = 0
|
||
|
||
//当缓存到第2组GOP的第二个合并写切片时，将上一个GOP缓存释放掉
|
||
if t.segmentOffset > len(head) && t.memoryPool[1] != nil && !t.memoryPool[1].Empty() {
|
||
t.memoryPool[1].Clear()
|
||
}
|
||
}
|
||
|
||
func (t *TransStream) AddSink(sink stream.ISink) {
|
||
utils.Assert(t.headerSize > 0)
|
||
|
||
t.TransStreamImpl.AddSink(sink)
|
||
//发送sequence header
|
||
sink.Input(t.header[:t.headerSize])
|
||
|
||
if !stream.AppConfig.GOPCache || t.onlyAudio {
|
||
return
|
||
}
|
||
|
||
//发送当前内存池已有的合并写切片
|
||
if t.segmentOffset > 0 {
|
||
data, tail := t.memoryPool[0].Data()
|
||
utils.Assert(len(data) > 0)
|
||
utils.Assert(len(tail) == 0)
|
||
sink.Input(data[:t.segmentOffset])
|
||
return
|
||
}
|
||
|
||
//发送上一组GOP
|
||
if t.memoryPool[1] != nil && !t.memoryPool[1].Empty() {
|
||
data, tail := t.memoryPool[0].Data()
|
||
utils.Assert(len(data) > 0)
|
||
utils.Assert(len(tail) == 0)
|
||
sink.Input(data)
|
||
return
|
||
}
|
||
|
||
//不足一个合并写切片, 有多少发多少
|
||
data, tail := t.memoryPool[0].Data()
|
||
utils.Assert(len(tail) == 0)
|
||
if len(data) > 0 {
|
||
sink.Input(data)
|
||
t.incompleteSinks = append(t.incompleteSinks, sink)
|
||
}
|
||
}
|
||
|
||
func (t *TransStream) WriteHeader() error {
|
||
utils.Assert(t.Tracks != nil)
|
||
utils.Assert(!t.TransStreamImpl.Completed)
|
||
|
||
var audioStream utils.AVStream
|
||
var videoStream utils.AVStream
|
||
var audioCodecId utils.AVCodecID
|
||
var videoCodecId utils.AVCodecID
|
||
|
||
for _, track := range t.Tracks {
|
||
if utils.AVMediaTypeAudio == track.Type() {
|
||
audioStream = track
|
||
audioCodecId = audioStream.CodecId()
|
||
t.audioChunk = librtmp.NewAudioChunk()
|
||
} else if utils.AVMediaTypeVideo == track.Type() {
|
||
videoStream = track
|
||
videoCodecId = videoStream.CodecId()
|
||
t.videoChunk = librtmp.NewVideoChunk()
|
||
}
|
||
}
|
||
|
||
utils.Assert(audioStream != nil || videoStream != nil)
|
||
|
||
//初始化
|
||
t.TransStreamImpl.Completed = true
|
||
t.header = make([]byte, 1024)
|
||
t.muxer = libflv.NewMuxer(audioCodecId, videoCodecId, 0, 0, 0)
|
||
|
||
if stream.AppConfig.GOPCache {
|
||
//创建2块内存
|
||
t.memoryPool[0] = stream.NewMemoryPoolWithDirect(1024*4000, true)
|
||
t.memoryPool[1] = stream.NewMemoryPoolWithDirect(1024*4000, true)
|
||
} else {
|
||
|
||
}
|
||
|
||
var n int
|
||
if audioStream != nil {
|
||
n += t.muxer.WriteAudioData(t.header[12:], true)
|
||
extra := audioStream.Extra()
|
||
copy(t.header[n+12:], extra)
|
||
n += len(extra)
|
||
|
||
t.audioChunk.Length = n
|
||
t.audioChunk.ToBytes(t.header)
|
||
n += 12
|
||
}
|
||
|
||
if videoStream != nil {
|
||
tmp := n
|
||
n += t.muxer.WriteVideoData(t.header[n+12:], 0, false, true)
|
||
extra := videoStream.Extra()
|
||
copy(t.header[n+12:], extra)
|
||
n += len(extra)
|
||
|
||
t.videoChunk.Length = 5 + len(extra)
|
||
t.videoChunk.ToBytes(t.header[tmp:])
|
||
n += 12
|
||
}
|
||
|
||
t.headerSize = n
|
||
return nil
|
||
}
|