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 (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"time"

	"bitbucket.org/ausocean/av/rtmp"
	"bitbucket.org/ausocean/utils/ring"
)

// 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(*ring.Chunk) 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

	chunk *ring.Chunk
}

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

func (s *fileSender) load(c *ring.Chunk) error {
	s.chunk = c
	return nil
}

func (s *fileSender) send() error {
	_, err := s.chunk.WriteTo(s.file)
	return err
}

func (s *fileSender) release() {
	s.chunk.Close()
	s.chunk = nil
}

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

// httpSender implements loadSender for an HTTP destination.
type httpSender struct {
	client http.Client
	url    string

	log func(lvl, msg string)

	buf []byte

	chunk *ring.Chunk
}

func newHttpSender(url string, timeout time.Duration, log func(lvl, msg string)) *httpSender {
	return &httpSender{
		client: http.Client{Timeout: timeout},
		url:    url,
		log:    log,
	}
}

func (s *httpSender) load(c *ring.Chunk) error {
	buf := bytes.NewBuffer(s.buf[:0])
	_, err := s.chunk.WriteTo(buf)
	s.buf = buf.Bytes()
	if err != nil {
		return fmt.Errorf("fileSender: %v", err)
	}
	return nil
}

func (s *httpSender) send() error {
	url := s.url + strconv.Itoa(len(s.buf))
	s.log(Debug, fmt.Sprintf("Posting %s (%d bytes)\n", url, len(s.buf)))
	resp, err := s.client.Post(url, "video/mp2t", bytes.NewReader(s.buf))
	if err != nil {
		return fmt.Errorf("Error posting to %s: %s", url, err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err == nil {
		s.log(Debug, fmt.Sprintf("%s\n", body))
	} else {
		s.log(Error, err.Error())
	}
	return err
}

func (s *httpSender) release() {
	s.chunk.Close()
	s.chunk = nil
}

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

// ffmpegSender implements loadSender for an FFMPEG RTMP destination.
type ffmpegSender struct {
	ffmpeg io.WriteCloser

	chunk *ring.Chunk
}

func newFfmpegSender(url, framerate string) (*ffmpegSender, error) {
	cmd := exec.Command(ffmpegPath,
		"-f", "h264",
		"-r", framerate,
		"-i", "-",
		"-f", "lavfi",
		"-i", "aevalsrc=0",
		"-fflags", "nobuffer",
		"-vcodec", "copy",
		"-acodec", "aac",
		"-map", "0:0",
		"-map", "1:0",
		"-strict", "experimental",
		"-f", "flv",
		url,
	)
	w, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	err = cmd.Start()
	if err != nil {
		return nil, err
	}
	return &ffmpegSender{ffmpeg: w}, nil
}

func (s *ffmpegSender) load(c *ring.Chunk) error {
	s.chunk = c
	return nil
}

func (s *ffmpegSender) send() error {
	_, err := s.chunk.WriteTo(s.ffmpeg)
	return err
}

func (s *ffmpegSender) release() {
	s.chunk.Close()
	s.chunk = nil
}

func (s *ffmpegSender) close() error {
	return s.ffmpeg.Close()
}

// rtmpSender implements loadSender for a native RTMP destination.
type rtmpSender struct {
	sess rtmp.Session

	url     string
	timeout uint
	retries int
	log     func(lvl, msg string)

	chunk *ring.Chunk
}

var _ restarter = (*rtmpSender)(nil)

func newRtmpSender(url string, timeout uint, retries int, log func(lvl, msg string)) (*rtmpSender, error) {
	var sess rtmp.Session
	var err error
	for n := 0; n < retries; n++ {
		sess = rtmp.NewSession(url, timeout)
		err = sess.StartSession()
		if err == nil {
			break
		}
		log(Error, err.Error())
		sess.Close()
		if n < retries-1 {
			log(Info, "retry rtmp connection")
		}
	}
	if err != nil {
		return nil, err
	}

	s := &rtmpSender{
		sess:    sess,
		url:     url,
		timeout: timeout,
		retries: retries,
		log:     log,
	}
	return s, nil
}

func (s *rtmpSender) load(c *ring.Chunk) error {
	s.chunk = c
	return nil
}

func (s *rtmpSender) send() error {
	_, err := s.chunk.WriteTo(s.sess)
	return err
}

func (s *rtmpSender) release() {
	s.chunk.Close()
	s.chunk = nil
}

func (s *rtmpSender) restart() error {
	err := s.sess.Close()
	if err != nil {
		return err
	}
	for n := 0; n < s.retries; n++ {
		s.sess = rtmp.NewSession(s.url, s.timeout)
		err = s.sess.StartSession()
		if err == nil {
			break
		}
		s.log(Error, err.Error())
		s.sess.Close()
		if n < s.retries-1 {
			s.log(Info, "retry rtmp connection")
		}
	}
	return err
}

func (s *rtmpSender) close() error {
	return s.sess.Close()
}