From 1359426f1b4194e3ccf37b612cec5964abcbb7b4 Mon Sep 17 00:00:00 2001 From: Han Gyoung-Su Date: Sun, 22 Jun 2025 18:38:46 +0900 Subject: [PATCH] Feature/drain (#18) * feat: mp4 split issue * feat: drain process for HLS, webm, mp4, WHEP --- go.mod | 4 +- media/streamer/egress/hls/handler.go | 46 +++++++++++---- media/streamer/egress/record/mp4/handler.go | 28 +++++++-- media/streamer/egress/record/webm/handler.go | 18 ++++++ media/streamer/egress/whep/whep.go | 17 ++++++ media/streamer/ingress/rtmp/server.go | 4 ++ media/streamer/processes/transcoder.go | 60 ++++++++++++++++++++ 7 files changed, 159 insertions(+), 18 deletions(-) diff --git a/go.mod b/go.mod index 38a5234..39768e9 100644 --- a/go.mod +++ b/go.mod @@ -1,10 +1,9 @@ module liveflow -go 1.21 +go 1.23 require ( github.com/asticode/go-astiav v0.19.0 - github.com/asticode/go-astits v1.13.0 github.com/at-wat/ebml-go v0.17.1 github.com/bluenviron/gohlslib v1.4.0 github.com/deepch/vdk v0.0.27 @@ -26,6 +25,7 @@ require ( require ( github.com/abema/go-mp4 v1.2.0 // indirect github.com/asticode/go-astikit v0.43.0 // indirect + github.com/asticode/go-astits v1.13.0 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/bluenviron/mediacommon v1.11.1-0.20240525122142-20163863aa75 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect diff --git a/media/streamer/egress/hls/handler.go b/media/streamer/egress/hls/handler.go index 9bd4cc8..38b4063 100644 --- a/media/streamer/egress/hls/handler.go +++ b/media/streamer/egress/hls/handler.go @@ -4,9 +4,10 @@ import ( "context" "errors" "fmt" - "liveflow/media/streamer/processes" "time" + "liveflow/media/streamer/processes" + "github.com/asticode/go-astiav" "github.com/bluenviron/gohlslib" "github.com/bluenviron/gohlslib/pkg/codecs" @@ -99,11 +100,43 @@ func (h *HLS) Start(ctx context.Context, source hub.Source) error { h.onVideo(ctx, data.H264Video) } } + if audioTranscodingProcess != nil { + log.Info(ctx, "draining audio transcoding process for HLS") + packets, err := audioTranscodingProcess.Drain() + if err != nil { + log.Error(ctx, err, "failed to drain audio transcoder for HLS") + } + for _, packet := range packets { + h.onAudio(ctx, source, &hub.AACAudio{ + Data: packet.Data, + SequenceHeader: false, + MPEG4AudioConfigBytes: h.mpeg4AudioConfigBytes, + MPEG4AudioConfig: h.mpeg4AudioConfig, + PTS: packet.PTS, + DTS: packet.DTS, + AudioClockRate: uint32(packet.SampleRate), + }) + } + log.Info(ctx, "audio transcoding process for HLS drained") + } log.Info(ctx, "[HLS] end of streamID: ", source.StreamID()) + if h.muxer != nil { + h.muxer.Close() + } }() return nil } +func (h *HLS) onVideo(ctx context.Context, h264Video *hub.H264Video) { + if h.muxer != nil { + au, _ := h264parser.SplitNALUs(h264Video.Data) + err := h.muxer.WriteH264(time.Now(), time.Duration(h264Video.RawDTS())*time.Millisecond, au) + if err != nil { + log.Errorf(ctx, "failed to write h264: %v", err) + } + } +} + func (h *HLS) onAudio(ctx context.Context, source hub.Source, aacAudio *hub.AACAudio) { if len(aacAudio.MPEG4AudioConfigBytes) > 0 { if h.muxer == nil { @@ -125,17 +158,6 @@ func (h *HLS) onAudio(ctx context.Context, source hub.Source, aacAudio *hub.AACA h.muxer.WriteMPEG4Audio(time.Now(), time.Duration(aacAudio.RawDTS())*time.Millisecond, [][]byte{audioData}) } } - -func (h *HLS) onVideo(ctx context.Context, h264Video *hub.H264Video) { - if h.muxer != nil { - au, _ := h264parser.SplitNALUs(h264Video.Data) - err := h.muxer.WriteH264(time.Now(), time.Duration(h264Video.RawDTS())*time.Millisecond, au) - if err != nil { - log.Errorf(ctx, "failed to write h264: %v", err) - } - } -} - func (h *HLS) onOPUSAudio(ctx context.Context, source hub.Source, audioTranscodingProcess *processes.AudioTranscodingProcess, opusAudio *hub.OPUSAudio) { packets, err := audioTranscodingProcess.Process(&processes.MediaPacket{ Data: opusAudio.Data, diff --git a/media/streamer/egress/record/mp4/handler.go b/media/streamer/egress/record/mp4/handler.go index f7e3d6c..c6712e8 100644 --- a/media/streamer/egress/record/mp4/handler.go +++ b/media/streamer/egress/record/mp4/handler.go @@ -118,11 +118,29 @@ func (m *MP4) Start(ctx context.Context, source hub.Source) error { } } } - err = m.muxer.WriteTrailer() - if err != nil { - log.Error(ctx, err, "failed to write trailer") + + if audioTranscodingProcess != nil { + log.Info(ctx, "draining audio transcoding process") + // Drain the remaining frames from the transcoder + packets, err := audioTranscodingProcess.Drain() + if err != nil { + log.Error(ctx, err, "failed to drain audio transcoder") + } + for _, packet := range packets { + m.onAudio(ctx, &hub.AACAudio{ + Data: packet.Data, + SequenceHeader: false, + MPEG4AudioConfigBytes: m.mpeg4AudioConfigBytes, + MPEG4AudioConfig: m.mpeg4AudioConfig, + PTS: packet.PTS, + DTS: packet.DTS, + AudioClockRate: uint32(packet.SampleRate), + }) + } + log.Info(ctx, "audio transcoding process drained") + } else { + log.Info(ctx, "no audio transcoding process to drain") } - log.Info(ctx, "mp4 file closed") }() return nil } @@ -153,6 +171,7 @@ func (m *MP4) createNewFile(ctx context.Context) error { // closeFile closes the current MP4 file and muxer func (m *MP4) closeFile(ctx context.Context) { if m.muxer != nil { + log.Info(ctx, "writing mp4 trailer") err := m.muxer.WriteTrailer() if err != nil { log.Error(ctx, err, "failed to write trailer") @@ -164,6 +183,7 @@ func (m *MP4) closeFile(ctx context.Context) { if err != nil { log.Error(ctx, err, "failed to close mp4 file") } + log.Info(ctx, "mp4 file closed") m.tempFile = nil } } diff --git a/media/streamer/egress/record/webm/handler.go b/media/streamer/egress/record/webm/handler.go index 5c4b09e..acc818d 100644 --- a/media/streamer/egress/record/webm/handler.go +++ b/media/streamer/egress/record/webm/handler.go @@ -99,6 +99,24 @@ func (w *WebM) Start(ctx context.Context, source hub.Source) error { w.onAudio(ctx, data.OPUSAudio) } } + + if w.audioTranscodingProcess != nil { + log.Info(ctx, "draining audio transcoding process for WebM") + packets, err := w.audioTranscodingProcess.Drain() + if err != nil { + log.Error(ctx, err, "failed to drain audio transcoder for WebM") + } + for _, packet := range packets { + w.onAudio(ctx, &hub.OPUSAudio{ + Data: packet.Data, + PTS: packet.PTS, + DTS: packet.DTS, + AudioClockRate: uint32(packet.SampleRate), + }) + } + log.Info(ctx, "audio transcoding process for WebM drained") + } + // Ensure the muxer is finalized w.closeMuxer(ctx) }() diff --git a/media/streamer/egress/whep/whep.go b/media/streamer/egress/whep/whep.go index bfbda51..e7214e3 100644 --- a/media/streamer/egress/whep/whep.go +++ b/media/streamer/egress/whep/whep.go @@ -105,7 +105,24 @@ func (w *WHEP) Start(ctx context.Context, source hub.Source) error { } if audioTranscodingProcess != nil { + log.Info(ctx, "draining audio transcoding process for WHEP") + packets, err := audioTranscodingProcess.Drain() + if err != nil { + log.Error(ctx, err, "failed to drain audio transcoder for WHEP") + } + for _, packet := range packets { + err := w.onAudio(source, &hub.OPUSAudio{ + Data: packet.Data, + PTS: packet.PTS, + DTS: packet.DTS, + AudioClockRate: uint32(packet.SampleRate), + }) + if err != nil { + log.Error(ctx, err, "failed to process drained OPUS audio for WHEP") + } + } audioTranscodingProcess.Close() + log.Info(ctx, "audio transcoding process for WHEP drained and closed") } log.Info(ctx, "end whep") //C.__lsan_do_leak_check() diff --git a/media/streamer/ingress/rtmp/server.go b/media/streamer/ingress/rtmp/server.go index 044a7b2..26a5d51 100644 --- a/media/streamer/ingress/rtmp/server.go +++ b/media/streamer/ingress/rtmp/server.go @@ -67,6 +67,10 @@ func (r *RTMP) Serve(ctx context.Context) error { }, }) log.Info(ctx, "RTMP server started") + go func() { + <-ctx.Done() + srv.Close() + }() if err := srv.Serve(listener); err != nil { log.Errorf(ctx, "Failed: %+v", err) } diff --git a/media/streamer/processes/transcoder.go b/media/streamer/processes/transcoder.go index 2515b2f..bdad1fd 100644 --- a/media/streamer/processes/transcoder.go +++ b/media/streamer/processes/transcoder.go @@ -194,3 +194,63 @@ func (t *AudioTranscodingProcess) Process(data *MediaPacket) ([]*MediaPacket, er return opusAudio, nil } + +func (t *AudioTranscodingProcess) Drain() ([]*MediaPacket, error) { + var packets []*MediaPacket + ctx := context.Background() + + // 1. Flush Decoder + if err := t.decCodecContext.SendPacket(nil); err != nil { + log.Error(ctx, err, "failed to send nil packet to decoder for draining") + // Continue to encoder flushing + } + + frame := astiav.AllocFrame() + defer frame.Free() + + for { + err := t.decCodecContext.ReceiveFrame(frame) + if errors.Is(err, astiav.ErrEof) || errors.Is(err, astiav.ErrEagain) { + break + } + if err != nil { + log.Error(ctx, err, "failed to receive frame from decoder during draining") + continue + } + + if err := t.encCodecContext.SendFrame(frame); err != nil { + log.Error(ctx, err, "failed to send flushed frame to encoder") + } + frame.Unref() + } + + // 2. Flush Encoder + if err := t.encCodecContext.SendFrame(nil); err != nil { + log.Error(ctx, err, "failed to send nil frame to encoder for draining") + return packets, err + } + + pkt := astiav.AllocPacket() + defer pkt.Free() + + for { + err := t.encCodecContext.ReceivePacket(pkt) + if errors.Is(err, astiav.ErrEof) || errors.Is(err, astiav.ErrEagain) { + break + } + if err != nil { + log.Error(ctx, err, "failed to receive packet from encoder during draining") + continue + } + + packets = append(packets, &MediaPacket{ + Data: append([]byte{}, pkt.Data()...), + PTS: pkt.Pts(), + DTS: pkt.Dts(), + SampleRate: t.encCodecContext.SampleRate(), + }) + pkt.Unref() + } + + return packets, nil +}