monibuca/plugin/hls/pkg/record.go

package hls

import (
	"context"
	"fmt"
	"path/filepath"
	"time"

	"m7s.live/v5"
	"m7s.live/v5/pkg/codec"
	"m7s.live/v5/pkg/config"
	"m7s.live/v5/pkg/format"
	mpegts "m7s.live/v5/pkg/format/ts"
	"m7s.live/v5/pkg/storage"
)

func NewRecorder(conf config.Record) m7s.IRecorder {
	return &Recorder{}
}

type Recorder struct {
	m7s.DefaultRecorder
	TsInMemory
	file         storage.File
	segmentCount uint32
	lastTs       time.Duration
	firstSegment bool
}

var CustomFileName = func(job *m7s.RecordJob) string {
	if job.RecConf.Fragment == 0 || job.RecConf.Append {
		return fmt.Sprintf("%s/%s.ts", job.RecConf.FilePath, time.Now().Format("20060102150405"))
	}
	return filepath.Join(job.RecConf.FilePath, time.Now().Format("20060102150405")+".ts")
}

func (r *Recorder) createStream(start time.Time) (err error) {
	r.RecordJob.RecConf.Type = "ts"
	return r.CreateStream(start, CustomFileName)
}

func (r *Recorder) writeTailer(end time.Time) {
	if !r.RecordJob.RecConf.RealTime {
		defer r.TsInMemory.Recycle()
		var err error
		r.file, err = r.RecordJob.GetStorage().CreateFile(context.Background(), r.Event.FilePath)
		if err != nil {
			return
		}
		r.WriteTo(r.file)
	}
	r.file.Close()
	r.WriteTail(end, nil)
}

func (r *Recorder) Dispose() {
	r.writeTailer(time.Now())
}

func (r *Recorder) createNewTs() (err error) {
	if r.RecordJob.RecConf.RealTime {
		r.file, err = r.RecordJob.GetStorage().CreateFile(context.Background(), r.Event.FilePath)
	}
	return
}

func (r *Recorder) writeSegment(ts time.Duration, writeTime time.Time) (err error) {
	if dur := ts - r.lastTs; dur >= r.RecordJob.RecConf.Fragment || r.lastTs == 0 {
		if dur == ts && r.lastTs == 0 { //时间戳不对的情况，首个默认为2s
			dur = time.Duration(2) * time.Second
		}

		// 如果是第一个片段，跳过写入，只记录时间戳
		if r.firstSegment {
			r.lastTs = ts
			r.firstSegment = false
			return nil
		}

		// 结束当前片段的记录
		r.writeTailer(writeTime)

		// 创建新的数据库记录
		err = r.createStream(writeTime)
		if err != nil {
			return
		}

		// 创建新的ts文件
		if err = r.createNewTs(); err != nil {
			return
		}
		if r.RecordJob.RecConf.RealTime {
			r.file.Write(mpegts.DefaultPATPacket)
			r.file.Write(r.PMT)
		}
		r.segmentCount++
		r.lastTs = ts
	}
	return
}

func (r *Recorder) Run() (err error) {
	ctx := &r.RecordJob
	suber := ctx.Subscriber
	startTime := time.Now()
	// 创建第一个片段记录
	if err = r.createStream(startTime); err != nil {
		return
	}

	// 初始化HLS相关结构
	if err = r.createNewTs(); err != nil {
		return
	}
	pesAudio, pesVideo := mpegts.CreatePESWriters()
	r.firstSegment = true

	var audioCodec, videoCodec codec.FourCC
	if suber.Publisher.HasAudioTrack() {
		audioCodec = suber.Publisher.AudioTrack.FourCC()
	}
	if suber.Publisher.HasVideoTrack() {
		videoCodec = suber.Publisher.VideoTrack.FourCC()
	}
	r.WritePMTPacket(audioCodec, videoCodec)
	if ctx.RecConf.RealTime {
		r.file.Write(mpegts.DefaultPATPacket)
		r.file.Write(r.PMT)
	}
	return m7s.PlayBlock(suber, func(audio *format.Mpeg2Audio) (err error) {
		pesAudio.Pts = uint64(suber.AudioReader.AbsTime) * 90
		err = pesAudio.WritePESPacket(audio.Memory, &r.RecyclableMemory)
		if err == nil {
			if ctx.RecConf.RealTime {
				r.RecyclableMemory.WriteTo(r.file)
				r.RecyclableMemory.Recycle()
			}
		}
		return
	}, func(video *mpegts.VideoFrame) (err error) {
		r.Event.Duration = uint32(video.Timestamp.Milliseconds() - r.lastTs.Milliseconds())
		vr := r.RecordJob.Subscriber.VideoReader
		if vr.Value.IDR {
			if err = r.writeSegment(video.Timestamp, vr.Value.WriteTime); err != nil {
				return
			}
		}
		pesVideo.IsKeyFrame = video.IDR
		pesVideo.Pts = uint64(vr.AbsTime+video.GetCTS32()) * 90
		pesVideo.Dts = uint64(vr.AbsTime) * 90
		err = pesVideo.WritePESPacket(video.Memory, &r.RecyclableMemory)
		if err == nil {
			if ctx.RecConf.RealTime {
				r.RecyclableMemory.WriteTo(r.file)
				r.RecyclableMemory.Recycle()
			}
		}
		return
	})
}