segmento/pkg/media/media.go

// Package media contains logic related to parsing and segmenting media
// streams, such as MP3 streams.
//
// It is depended upon by the playlist package.
package media

import (
	"bytes"
	"io"
	"log"
	"time"

	"github.com/tcolgate/mp3"
)

const DefaultTargetDuration = 3 * time.Second

type Segmenter interface {
	Segment(r io.Reader) (chan *Segment, error)
}

type SegmentPublisher interface {
	Publish(s *Segment) (string, error)
}

type Segment struct {
	duration       time.Duration
	targetDuration time.Duration
	data           *bytes.Buffer
}

// Initialize a new Segment struct. The capacity is the initial maximum capacity of the internal
// buffer, in bytes. It should be initialized with a value greater than the expected maximum buffer size,
// depending on the implementation.
func NewSegment(targetDuration time.Duration, capacity int) *Segment {
	return &Segment{
		data:           bytes.NewBuffer(make([]byte, 0, capacity)),
		duration:       0,
		targetDuration: targetDuration,
	}
}

func (s *Segment) ReadFrom(r io.Reader) (n int64, err error) {
	return s.data.ReadFrom(r)
}

func (s *Segment) IncrementDuration(d time.Duration) time.Duration {
	s.duration += d
	return s.duration
}

func (s *Segment) CanWrite(d time.Duration) bool {
	return s.targetDuration-s.duration >= d
}

func (s *Segment) Duration() time.Duration {
	return s.duration
}

func (s *Segment) Len() int {
	return s.data.Len()
}

type MP3Segmenter struct {
	decoder *mp3.Decoder
}

func (s *MP3Segmenter) Segment(r io.Reader) (chan *Segment, error) {
	c := make(chan *Segment)

	go func() {
		d := mp3.NewDecoder(r)

		var (
			v       mp3.Frame
			skipped int
		)

		var (
			s *Segment
		)

		for {
			if err := d.Decode(&v, &skipped); err != nil {
				if err == io.EOF {
					break
				}
				log.Fatal(err)
			}

			if s != nil && !s.CanWrite(v.Duration()) {
				// publish the current segment, and initialize a new one.
				c <- s
				s = nil
			}

			if s == nil {
				// TODO what is a good initial capacity?
				s = NewSegment(DefaultTargetDuration, 1024)
			}

			n, err := s.ReadFrom(v.Reader())
			if err != nil {
				log.Fatal(err) // TODO: some proper error handling
			}

			if n != int64(v.Size()) {
				// TODO would this ever happen? What should IncrementDuration be called with?
				log.Fatal("unexpeced frame size, want = ", v.Size(), "got = ", n)
			}

			s.IncrementDuration(v.Duration())
		}
	}()

	return c, nil
}

func NewMP3Segmenter() *MP3Segmenter {
	return &MP3Segmenter{}
}

type FakePublisher struct{}

func (p *FakePublisher) Publish(s *Segment) (string, error) {
	return "https://www.example.com/segment.mp3", nil
}