tracker/cmd/chihaya/main.go

283 lines
7.1 KiB
Go
Raw Normal View History

2016-08-05 07:47:04 +02:00
package main
import (
"errors"
"os"
"os/signal"
"runtime"
2016-08-05 07:47:04 +02:00
"runtime/pprof"
2018-07-04 05:14:33 +02:00
"runtime/trace"
"strings"
2016-08-05 07:47:04 +02:00
"syscall"
"time"
2016-08-05 07:47:04 +02:00
"github.com/sirupsen/logrus"
2016-08-05 07:47:04 +02:00
"github.com/spf13/cobra"
"github.com/chihaya/chihaya/frontend/http"
"github.com/chihaya/chihaya/frontend/udp"
2016-08-17 03:42:08 +02:00
"github.com/chihaya/chihaya/middleware"
2017-06-20 14:58:44 +02:00
"github.com/chihaya/chihaya/pkg/log"
"github.com/chihaya/chihaya/pkg/prometheus"
"github.com/chihaya/chihaya/pkg/stop"
"github.com/chihaya/chihaya/storage"
2016-08-05 07:47:04 +02:00
)
// Run represents the state of a running instance of Chihaya.
type Run struct {
configFilePath string
peerStore storage.PeerStore
logic *middleware.Logic
sg *stop.Group
}
2016-08-05 09:35:17 +02:00
// NewRun runs an instance of Chihaya.
func NewRun(configFilePath string) (*Run, error) {
r := &Run{
configFilePath: configFilePath,
2016-08-05 09:35:17 +02:00
}
return r, r.Start(nil)
}
// Start begins an instance of Chihaya.
// It is optional to provide an instance of the peer store to avoid the
// creation of a new one.
func (r *Run) Start(ps storage.PeerStore) error {
configFile, err := ParseConfigFile(r.configFilePath)
2016-08-05 09:35:17 +02:00
if err != nil {
return errors.New("failed to read config: " + err.Error())
2016-08-05 09:35:17 +02:00
}
2017-05-02 17:03:49 +02:00
cfg := configFile.Chihaya
r.sg = stop.NewGroup()
2017-05-07 00:48:44 +02:00
2017-06-20 14:58:44 +02:00
log.Info("starting Prometheus server", log.Fields{"addr": cfg.PrometheusAddr})
2017-05-02 17:03:49 +02:00
r.sg.Add(prometheus.NewServer(cfg.PrometheusAddr))
if ps == nil {
log.Info("starting storage", log.Fields{"name": cfg.Storage.Name})
2017-02-21 06:58:57 +01:00
ps, err = storage.NewPeerStore(cfg.Storage.Name, cfg.Storage.Config)
if err != nil {
return errors.New("failed to create storage: " + err.Error())
}
2017-09-19 21:27:52 +02:00
log.Info("started storage", ps)
}
r.peerStore = ps
2017-12-23 20:54:51 +01:00
preHooks, err := middleware.HooksFromHookConfigs(cfg.PreHooks)
2017-05-07 00:48:44 +02:00
if err != nil {
return errors.New("failed to validate hook config: " + err.Error())
}
2017-12-23 20:54:51 +01:00
postHooks, err := middleware.HooksFromHookConfigs(cfg.PostHooks)
if err != nil {
return errors.New("failed to validate hook config: " + err.Error())
}
log.Info("starting tracker logic", log.Fields{
"prehooks": cfg.PreHookNames(),
"posthooks": cfg.PostHookNames(),
2017-06-20 14:58:44 +02:00
})
2017-12-23 20:54:51 +01:00
r.logic = middleware.NewLogic(cfg.ResponseConfig, r.peerStore, preHooks, postHooks)
2016-08-05 09:35:17 +02:00
2017-05-02 17:03:49 +02:00
if cfg.HTTPConfig.Addr != "" {
2017-09-19 21:27:52 +02:00
log.Info("starting HTTP frontend", cfg.HTTPConfig)
2017-05-02 17:03:49 +02:00
httpfe, err := http.NewFrontend(r.logic, cfg.HTTPConfig)
if err != nil {
return err
}
r.sg.Add(httpfe)
}
2017-05-02 17:03:49 +02:00
if cfg.UDPConfig.Addr != "" {
2017-09-19 21:27:52 +02:00
log.Info("starting UDP frontend", cfg.UDPConfig)
2017-05-02 17:03:49 +02:00
udpfe, err := udp.NewFrontend(r.logic, cfg.UDPConfig)
if err != nil {
return err
}
r.sg.Add(udpfe)
}
return nil
}
func combineErrors(prefix string, errs []error) error {
var errStrs []string
for _, err := range errs {
errStrs = append(errStrs, err.Error())
}
return errors.New(prefix + ": " + strings.Join(errStrs, "; "))
}
// Stop shuts down an instance of Chihaya.
func (r *Run) Stop(keepPeerStore bool) (storage.PeerStore, error) {
log.Debug("stopping frontends and prometheus endpoint")
if errs := r.sg.Stop().Wait(); len(errs) != 0 {
return nil, combineErrors("failed while shutting down frontends", errs)
}
log.Debug("stopping logic")
if errs := r.logic.Stop().Wait(); len(errs) != 0 {
return nil, combineErrors("failed while shutting down middleware", errs)
}
if !keepPeerStore {
log.Debug("stopping peer store")
if errs := r.peerStore.Stop().Wait(); len(errs) != 0 {
return nil, combineErrors("failed while shutting down peer store", errs)
}
r.peerStore = nil
}
return r.peerStore, nil
}
// RootRunCmdFunc implements a Cobra command that runs an instance of Chihaya
// and handles reloading and shutdown via process signals.
func RootRunCmdFunc(cmd *cobra.Command, args []string) error {
configFilePath, err := cmd.Flags().GetString("config")
if err != nil {
return err
}
r, err := NewRun(configFilePath)
if err != nil {
return err
}
quit := make(chan os.Signal)
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
reload := makeReloadChan()
for {
select {
case <-reload:
2017-05-07 00:48:44 +02:00
log.Info("reloading; received SIGUSR1")
peerStore, err := r.Stop(true)
if err != nil {
return err
}
if err := r.Start(peerStore); err != nil {
return err
}
case <-quit:
2017-05-07 00:48:44 +02:00
log.Info("shutting down; received SIGINT/SIGTERM")
if _, err := r.Stop(false); err != nil {
return err
}
return nil
}
}
}
// RootPreRunCmdFunc handles command line flags for the Run command.
func RootPreRunCmdFunc(cmd *cobra.Command, args []string) error {
noColors, err := cmd.Flags().GetBool("nocolors")
if err != nil {
return err
}
if noColors {
log.SetFormatter(&logrus.TextFormatter{DisableColors: true})
}
jsonLog, err := cmd.Flags().GetBool("json")
if err != nil {
return err
}
if jsonLog {
log.SetFormatter(&logrus.JSONFormatter{})
log.Info("enabled JSON logging")
}
debugLog, err := cmd.Flags().GetBool("debug")
if err != nil {
return err
}
if debugLog {
log.SetDebug(true)
log.Info("enabled debug logging")
}
2018-07-04 05:14:33 +02:00
tracePath, err := cmd.Flags().GetString("trace")
if err != nil {
return err
}
if tracePath != "" {
f, err := os.Create(tracePath)
if err != nil {
return err
}
trace.Start(f)
log.Info("enabled tracing", log.Fields{"path": tracePath})
}
cpuProfilePath, err := cmd.Flags().GetString("cpuprofile")
if err != nil {
return err
}
if cpuProfilePath != "" {
f, err := os.Create(cpuProfilePath)
if err != nil {
return err
}
pprof.StartCPUProfile(f)
log.Info("enabled CPU profiling", log.Fields{"path": cpuProfilePath})
}
return nil
}
// RootPostRunCmdFunc handles clean up of any state initialized by command line
// flags.
func RootPostRunCmdFunc(cmd *cobra.Command, args []string) error {
2018-07-04 05:14:33 +02:00
// These can be called regardless because it noops when not profiling.
pprof.StopCPUProfile()
2018-07-04 05:14:33 +02:00
trace.Stop()
return nil
}
func main() {
var rootCmd = &cobra.Command{
Use: "chihaya",
Short: "BitTorrent Tracker",
Long: "A customizable, multi-protocol BitTorrent Tracker",
PersistentPreRunE: RootPreRunCmdFunc,
RunE: RootRunCmdFunc,
PersistentPostRunE: RootPostRunCmdFunc,
2016-08-05 07:47:04 +02:00
}
rootCmd.PersistentFlags().String("cpuprofile", "", "location to save a CPU profile")
rootCmd.PersistentFlags().String("trace", "", "location to save a trace")
rootCmd.PersistentFlags().Bool("debug", false, "enable debug logging")
rootCmd.PersistentFlags().Bool("json", false, "enable json logging")
if runtime.GOOS == "windows" {
rootCmd.PersistentFlags().Bool("nocolors", true, "disable log coloring")
} else {
rootCmd.PersistentFlags().Bool("nocolors", false, "disable log coloring")
}
2016-08-05 07:47:04 +02:00
rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file")
var e2eCmd = &cobra.Command{
Use: "e2e",
Short: "exec e2e tests",
Long: "Execute the Chihaya end-to-end test suite",
RunE: EndToEndRunCmdFunc,
}
e2eCmd.Flags().String("httpaddr", "http://127.0.0.1:6969/announce", "address of the HTTP tracker")
e2eCmd.Flags().String("udpaddr", "udp://127.0.0.1:6969", "address of the UDP tracker")
e2eCmd.Flags().Duration("delay", time.Second, "delay between announces")
rootCmd.AddCommand(e2eCmd)
2016-08-05 07:47:04 +02:00
if err := rootCmd.Execute(); err != nil {
log.Fatal("failed when executing root cobra command: " + err.Error())
2016-08-05 07:47:04 +02:00
}
}