av/revid/senders.go

/*
NAME
  senders.go

DESCRIPTION
  See Readme.md

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>
  Alan Noble <alan@ausocean.org>

LICENSE
  revid is Copyright (C) 2017-2018 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
 for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

package revid

import (
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"

	"github.com/Comcast/gots/packet"

	"bitbucket.org/ausocean/av/container/mts"
	"bitbucket.org/ausocean/av/protocol/rtmp"
	"bitbucket.org/ausocean/av/protocol/rtp"
	"bitbucket.org/ausocean/iot/pi/netsender"
	"bitbucket.org/ausocean/utils/logger"
)

// Log is used by the multiSender.
type Log func(level int8, message string, params ...interface{})

// multiSender implements io.Writer. It provides the capacity to send to multiple
// senders from a single Write call.
type multiSender struct {
	dst []io.Writer
	log Log
}

// newMultiSender returns a pointer to a new multiSender.
func newMultiSender(senders []io.Writer, log Log) *multiSender {
	return &multiSender{
		dst: senders,
		log: log,
	}
}

// Write implements io.Writer. This will call load (with the passed slice), send
// and release on all senders of multiSender.
func (s *multiSender) Write(d []byte) (int, error) {
	for i, sender := range s.dst {
		_, err := sender.Write(d)
		if err != nil {
			s.log(logger.Warning, pkg+"send failed", "sender", i, "error", err)
		}
	}
	for _, sender := range s.dst {
		s, ok := sender.(loadSender)
		if !ok {
			panic("sender is not a loadSender")
		}
		s.release()
	}
	return len(d), nil
}

// minimalHttpSender implements Sender for posting HTTP to netreceiver or vidgrind.
type httpSender struct {
	client *netsender.Sender
	log    func(lvl int8, msg string, args ...interface{})
}

// newMinimalHttpSender returns a pointer to a new minimalHttpSender.
func newHttpSender(ns *netsender.Sender, log func(lvl int8, msg string, args ...interface{})) *httpSender {
	return &httpSender{
		client: ns,
		log:    log,
	}
}

// send takes a bytes slice d and sends to http using s' http client.
func (s *httpSender) Write(d []byte) (int, error) {
	return len(d), httpSend(d, s.client, s.log)
}

func httpSend(d []byte, client *netsender.Sender, log func(lvl int8, msg string, args ...interface{})) error {
	// Only send if "V0" is configured as an input.
	send := false
	ip := client.Param("ip")
	pins := netsender.MakePins(ip, "V")
	for i, pin := range pins {
		if pin.Name == "V0" {
			send = true
			pins[i].Value = len(d)
			pins[i].Data = d
			pins[i].MimeType = "video/mp2t"
			break
		}
	}

	if !send {
		return nil
	}
	var err error
	var reply string
	reply, _, err = client.Send(netsender.RequestRecv, pins)
	if err != nil {
		return err
	}
	return extractMeta(reply, log)
}

// extractMeta looks at a reply at extracts any time or location data - then used
// to update time and location information in the mpegts encoder.
func extractMeta(r string, log func(lvl int8, msg string, args ...interface{})) error {
	dec, err := netsender.NewJSONDecoder(r)
	if err != nil {
		return nil
	}
	// Extract time from reply
	t, err := dec.Int("ts")
	if err != nil {
		log(logger.Warning, pkg+"No timestamp in reply")
	} else {
		log(logger.Debug, fmt.Sprintf("%v got timestamp: %v", pkg, t))
		mts.Meta.Add("ts", strconv.Itoa(t))
	}

	// Extract location from reply
	g, err := dec.String("ll")
	if err != nil {
		log(logger.Warning, pkg+"No location in reply")
	} else {
		log(logger.Debug, fmt.Sprintf("%v got location: %v", pkg, g))
		mts.Meta.Add("loc", g)
	}

	return nil
}

// loadSender is a destination to send a *ring.Chunk to.
// When a loadSender has finished using the *ring.Chunk
// it must be Closed.
type loadSender interface {
	// load assigns the *ring.Chunk to the loadSender.
	// The load call may fail, but must not mutate the
	// the chunk.
	load(d []byte) error

	// send performs a destination-specific send
	// operation. It must not mutate the chunk.
	send() error

	// release releases the *ring.Chunk.
	release()

	// close cleans up after use of the loadSender.
	close() error
}

// restart is an optional interface for loadSenders that
// can restart their connection.
type restarter interface {
	restart() error
}

// fileSender implements loadSender for a local file destination.
type fileSender struct {
	file *os.File
	data []byte
}

func newFileSender(path string) (io.Writer, error) {
	f, err := os.Create(path)
	if err != nil {
		return nil, err
	}
	return &fileSender{file: f}, nil
}

// Write implements io.Writer.
func (s *fileSender) Write(d []byte) (int, error) {
	return s.file.Write(d)
}

func (s *fileSender) close() error { return s.file.Close() }

// mtsSender implements loadSender and provides sending capability specifically
// for use with MPEGTS packetization. It handles the construction of appropriately
// lengthed clips based on PSI. It also fixes accounts for discontinuities by
// setting the discontinuity indicator for the first packet of a clip.
type mtsSender struct {
	dst      io.Writer
	buf      []byte
	next     []byte
	pkt      packet.Packet
	repairer *mts.DiscontinuityRepairer
	curPid   int
}

// newMtsSender returns a new mtsSender.
func newMtsSender(dst io.Writer, log func(lvl int8, msg string, args ...interface{})) *mtsSender {
	return &mtsSender{
		dst:      dst,
		repairer: mts.NewDiscontinuityRepairer(),
	}
}

// Write implements io.Writer.
func (s *mtsSender) Write(d []byte) (int, error) {
	if s.next != nil {
		s.buf = append(s.buf, s.next...)
	}
	bytes := make([]byte, len(d))
	copy(bytes, d)
	s.next = bytes
	copy(s.pkt[:], bytes)
	s.curPid = s.pkt.PID()
	if s.curPid == mts.PatPid && len(s.buf) > 0 {
		err := s.repairer.Repair(s.buf)
		if err == nil {
			_, err = s.dst.Write(s.buf)
			if err == nil {
				goto done
			}
		}
		s.repairer.Failed()
	done:
		s.buf = s.buf[:0]
	}
	return len(d), nil
}

func (s *mtsSender) close() error { return nil }

// rtmpSender implements loadSender for a native RTMP destination.
type rtmpSender struct {
	conn *rtmp.Conn

	url     string
	timeout uint
	retries int
	log     func(lvl int8, msg string, args ...interface{})

	data []byte
}

var _ restarter = (*rtmpSender)(nil)

func newRtmpSender(url string, timeout uint, retries int, log func(lvl int8, msg string, args ...interface{})) (*rtmpSender, error) {
	var conn *rtmp.Conn
	var err error
	for n := 0; n < retries; n++ {
		conn, err = rtmp.Dial(url, timeout, log)
		if err == nil {
			break
		}
		log(logger.Error, err.Error())
		if n < retries-1 {
			log(logger.Info, pkg+"retry rtmp connection")
		}
	}
	s := &rtmpSender{
		conn:    conn,
		url:     url,
		timeout: timeout,
		retries: retries,
		log:     log,
	}
	return s, err
}

// Write implements io.Writer.
func (s *rtmpSender) Write(d []byte) (int, error) {
	return write(s, d)
}

func (s *rtmpSender) load(d []byte) error {
	s.data = d
	return nil
}

func (s *rtmpSender) send() error {
	if s.conn == nil {
		return errors.New("no rtmp connection, cannot write")
	}
	_, err := s.conn.Write(s.data)
	if err != nil {
		err = s.restart()
	}
	return err
}

func (s *rtmpSender) release() {}

func (s *rtmpSender) restart() error {
	s.close()
	var err error
	for n := 0; n < s.retries; n++ {
		s.conn, err = rtmp.Dial(s.url, s.timeout, s.log)
		if err == nil {
			break
		}
		s.log(logger.Error, err.Error())
		if n < s.retries-1 {
			s.log(logger.Info, pkg+"retry rtmp connection")
		}
	}
	return err
}

func (s *rtmpSender) close() error {
	if s.conn != nil {
		return s.conn.Close()
	}
	return nil
}

// TODO: Write restart func for rtpSender
// rtpSender implements loadSender for a native udp destination with rtp packetization.
type rtpSender struct {
	log     func(lvl int8, msg string, args ...interface{})
	encoder *rtp.Encoder
	data    []byte
}

// Write implements io.Writer.
func (s *rtpSender) Write(d []byte) (int, error) {
	return write(s, d)
}

func newRtpSender(addr string, log func(lvl int8, msg string, args ...interface{}), fps uint) (*rtpSender, error) {
	conn, err := net.Dial("udp", addr)
	if err != nil {
		return nil, err
	}
	s := &rtpSender{
		log:     log,
		encoder: rtp.NewEncoder(conn, int(fps)),
	}
	return s, nil
}

func (s *rtpSender) load(d []byte) error {
	s.data = make([]byte, len(d))
	copy(s.data, d)
	return nil
}

func (s *rtpSender) close() error { return nil }

func (s *rtpSender) release() {}

func (s *rtpSender) send() error {
	_, err := s.encoder.Write(s.data)
	return err
}

// write wraps the load and send method for loadSenders.
func write(s loadSender, d []byte) (int, error) {
	err := s.load(d)
	if err != nil {
		return 0, err
	}
	err = s.send()
	if err != nil {
		return len(d), err
	}
	return len(d), nil
}