Feature/drain (#18)

* feat: mp4 split issue

* feat: drain process for HLS, webm, mp4, WHEP

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

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,

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 {
+		if audioTranscodingProcess != nil {
-			log.Error(ctx, err, "failed to write trailer")
+			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
 	}
 }

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)
 	}()

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()

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)
 	}

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
+}