mirror of
https://github.com/lkmio/lkm.git
synced 2025-10-04 06:46:24 +08:00
封装hls
This commit is contained in:
yangjiechina
293
rtmp/rtmp_stream.go
Normal file
293
rtmp/rtmp_stream.go
Normal file
@@ -0,0 +1,293 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user