dogtamer/cmd/dogtamer.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"

	"github.com/urfave/cli/v2"
	"github.uio.no/torjus/dogtamer"
	"github.uio.no/torjus/dogtamer/config"
	"github.uio.no/torjus/dogtamer/server"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

const (
	ExitSuccess      = 0
	ExitGenericError = 1
)

func main() {
	cli.VersionPrinter = func(c *cli.Context) {
		fmt.Println(dogtamer.GetVersion())
	}

	app := &cli.App{
		Name:    "dogtamer",
		Version: dogtamer.Version,
		Authors: []*cli.Author{
			{
				Name:  "Torjus Håkestad",
				Email: "torjus@usit.uio.no",
			},
		},
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "config",
				Usage:   "Path to config-file",
				Aliases: []string{"c"},
			},
			&cli.StringFlag{
				Name:  "cpu-profile",
				Usage: "Profile cpu usage and write to file",
			},
		},
		Commands: []*cli.Command{
			{
				Name:   "serve",
				Usage:  "Start dogtamer server",
				Action: ActionServe,
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		fmt.Printf("Error: %s\n", err)
		os.Exit(ExitGenericError)
	}
}

func ActionServe(c *cli.Context) error {
	// Setup temporary logger, incase config could not be read
	logger := setupServerLogger("INFO")
	cfg, err := config.FromDefaultLocations()
	if err != nil {
		logger.Warn("No config file found.")
	}
	// Setup proper logger
	logger = setupServerLogger(cfg.LogLevel)
	cfg.DebugLog(logger)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	// Setup RTMP-server
	rtmpServer := server.NewRTMPServer(ctx, cfg.RTMPListenAddr)
	rtmpServer.Logger = logger
	rtmpServer.Hostname = cfg.Hostname

	if c.IsSet("cpu-profile") {
		filename := c.String("cpu-profile")
		logger.Infow("CPU-profiling enabled.", "filename", filename)
		f, err := os.Create(filename)
		if err != nil {
			return cli.Exit(fmt.Sprintf("Unable to create cpu-profile file '%s': %s", filename, err), 1)
		}
		defer f.Close()

		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Setup web-server
	webDone := make(chan struct{})
	if cfg.HTTPServerEnable {
		ws := server.NewWebServer(ctx, rtmpServer)
		ws.ListenAddr = cfg.HTTPListenAddr
		ws.Logger = logger
		go func() {
			rtmpServer.Logger.Infow("Starting HTTP server.", "listen_addr", ws.ListenAddr)
			err := ws.Serve()
			if err != nil && err != http.ErrServerClosed {
				rtmpServer.Logger.Infow("HTTP server shut down with error.", "err", err)
			}
			close(webDone)
		}()
	} else {
		close(webDone)
	}
	// Listen for SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		<-sigChan
		rtmpServer.Logger.Debug("Got shutdown signal.")
		cancel()
	}()

	// Start RTMP-server
	rtmpServer.Logger.Infow("Starting RTMP server.", "listen_addr", rtmpServer.ListenAddr)
	if err := rtmpServer.Listen(); err != nil {
		return cli.Exit(err, ExitGenericError)
	}

	// Wait for webserver to exit, if started
	<-webDone
	rtmpServer.Logger.Info("Server shut down.")
	return nil
}

func setupServerLogger(loglevel string) *zap.SugaredLogger {
	logEncoderConfig := zap.NewProductionEncoderConfig()
	logEncoderConfig.EncodeCaller = zapcore.ShortCallerEncoder
	logEncoderConfig.EncodeLevel = zapcore.CapitalColorLevelEncoder
	logEncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
	logEncoderConfig.EncodeDuration = zapcore.StringDurationEncoder

	rootLoggerConfig := &zap.Config{
		Level:            zap.NewAtomicLevelAt(zap.DebugLevel),
		OutputPaths:      []string{"stdout"},
		ErrorOutputPaths: []string{"stdout"},
		Encoding:         "console",
		EncoderConfig:    logEncoderConfig,
		DisableCaller:    true,
	}

	switch strings.ToUpper(loglevel) {
	case "DEBUG":
		rootLoggerConfig.DisableCaller = false
		rootLoggerConfig.Level = zap.NewAtomicLevelAt(zap.DebugLevel)
	case "INFO":
		rootLoggerConfig.Level = zap.NewAtomicLevelAt(zap.InfoLevel)
	case "WARN", "WARNING":
		rootLoggerConfig.Level = zap.NewAtomicLevelAt(zap.WarnLevel)
	case "ERR", "ERROR":
		rootLoggerConfig.Level = zap.NewAtomicLevelAt(zap.ErrorLevel)
	}

	rootLogger, err := rootLoggerConfig.Build()
	if err != nil {
		panic(err)
	}

	return rootLogger.Named("RTMP").WithOptions().Sugar()
}