donut/internal/controllers/probers/srt_mpets.go

package probers

import (
	"context"
	"errors"
	"io"

	astisrt "github.com/asticode/go-astisrt/pkg"
	"github.com/asticode/go-astits"
	"github.com/flavioribeiro/donut/internal/entities"
	"github.com/flavioribeiro/donut/internal/mapper"
	"go.uber.org/fx"
	"go.uber.org/zap"
)

type SrtMpegTs struct {
	c *entities.Config
	l *zap.SugaredLogger
	m *mapper.Mapper
}

type ResultSrtMpegTs struct {
	fx.Out
	SrtMpegTsProber DonutProber `group:"probers"`
}

// NewSrtMpegTs creates a new SrtMpegTs DonutProber
func NewSrtMpegTs(
	c *entities.Config,
	l *zap.SugaredLogger,
	m *mapper.Mapper,
) ResultSrtMpegTs {
	return ResultSrtMpegTs{
		SrtMpegTsProber: &SrtMpegTs{
			c: c,
			l: l,
			m: m,
		},
	}
}

// Match returns true when the request is for an SrtMpegTs prober
func (c *SrtMpegTs) Match(req *entities.RequestParams) bool {
	if req.SRTHost != "" {
		return true
	}
	return false
}

// StreamInfo connects to the SRT stream and probe N packets to discovery the media properties.
func (c *SrtMpegTs) StreamInfo(req *entities.RequestParams) (*entities.StreamInfo, error) {
	streamInfoMap, err := c.streamInfoMap(req)
	if err != nil {
		return nil, err
	}

	si := &entities.StreamInfo{}
	for _, v := range streamInfoMap {
		si.Streams = append(si.Streams, v)
	}
	return si, err
}

func (c *SrtMpegTs) streamInfoMap(req *entities.RequestParams) (map[entities.Codec]entities.Stream, error) {
	r, w := io.Pipe()
	defer r.Close()
	defer w.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	srtConnection, err := c.connect(cancel, req)
	if err != nil {
		return nil, err
	}
	defer srtConnection.Close()

	streamInfoMap := map[entities.Codec]entities.Stream{}

	go c.fromSRTToWriterPipe(srtConnection, w, cancel)

	c.l.Info("probing has starting demuxing")

	mpegTSDemuxer := astits.NewDemuxer(ctx, r)
	for {
		select {
		case <-ctx.Done():
			if errors.Is(ctx.Err(), context.Canceled) {
				c.l.Infow("probing has stopped due cancellation")
				return streamInfoMap, nil
			}
			c.l.Errorw("probing has stopped due errors")
			return streamInfoMap, ctx.Err()
		default:
			stop, err := c.fillStreamInfoFromMpegTS(streamInfoMap, mpegTSDemuxer)
			if stop {
				if err != nil {
					return nil, err
				}
				return streamInfoMap, nil
			}
		}
	}
}

func (c *SrtMpegTs) fromSRTToWriterPipe(srtConnection *astisrt.Connection, w *io.PipeWriter, cancel context.CancelFunc) {
	defer cancel()
	defer w.Close()
	defer srtConnection.Close()

	inboundMpegTsPacket := make([]byte, c.c.SRTReadBufferSizeBytes)
	c.l.Info("probing has started")

	for i := 1; i < c.c.ProbingSize; i++ {
		n, err := srtConnection.Read(inboundMpegTsPacket)
		if err != nil {
			c.l.Errorw("str conn failed to write data to buffer",
				"error", err,
			)
			break
		}

		if _, err := w.Write(inboundMpegTsPacket[:n]); err != nil {
			c.l.Errorw("failed to write mpeg-ts into the pipe",
				"error", err,
			)
			break
		}
	}
	c.l.Info("done probing")
}

func (c *SrtMpegTs) fillStreamInfoFromMpegTS(streamInfo map[entities.Codec]entities.Stream, mpegTSDemuxer *astits.Demuxer) (bool, error) {
	mpegTSDemuxData, err := mpegTSDemuxer.NextData()

	if err != nil {
		if !errors.Is(err, context.Canceled) {
			c.l.Errorw("failed to demux mpeg-ts",
				"error", err,
			)
			return true, err
		}
		return true, nil
	}

	if mpegTSDemuxData.PMT != nil {
		// TODO: add timing information
		// pts https://github.com/asticode/go-astits/blob/b0b19247aa31633650c32638fb55f597fa6e2468/cmd/astits-es-split/main.go#L206
		// https://github.com/asticode/go-astits/blob/master/packet.go#L46
		for _, es := range mpegTSDemuxData.PMT.ElementaryStreams {
			streamInfo[c.m.FromMpegTsStreamTypeToCodec(es.StreamType)] = c.m.FromStreamTypeToEntityStream(es)
		}
	}
	return false, nil
}

// TODO: move to its own component later dup streamer.srt_mpegts, prober.srt_mpegts
func (c *SrtMpegTs) connect(cancel context.CancelFunc, params *entities.RequestParams) (*astisrt.Connection, error) {
	c.l.Info("trying to connect srt")

	if err := params.Valid(); err != nil {
		return nil, err
	}

	c.l.Infow("Connecting to SRT ",
		"offer", params.String(),
	)

	conn, err := astisrt.Dial(astisrt.DialOptions{
		ConnectionOptions: []astisrt.ConnectionOption{
			astisrt.WithLatency(c.c.SRTConnectionLatencyMS),
			astisrt.WithStreamid(params.SRTStreamID),
			astisrt.WithCongestion("live"),
			astisrt.WithTranstype(astisrt.Transtype(astisrt.TranstypeLive)),
		},

		OnDisconnect: func(conn *astisrt.Connection, err error) {
			c.l.Infow("Canceling SRT",
				"error", err,
			)
			cancel()
		},

		Host: params.SRTHost,
		Port: params.SRTPort,
	})
	if err != nil {
		c.l.Errorw("failed to connect srt",
			"error", err,
		)
		return nil, err
	}
	c.l.Infow("Connected to SRT")
	return conn, nil
}