diff --git a/bittorrent/http/tracker.go b/bittorrent/http/tracker.go index a2edabc..f029de2 100644 --- a/bittorrent/http/tracker.go +++ b/bittorrent/http/tracker.go @@ -29,6 +29,11 @@ import ( "github.com/jzelinskie/trakr/bittorrent" ) +func init() { + prometheus.MustRegister(promResponseDurationMilliseconds) + recordResponseDuration("action", nil, time.Second) +} + var promResponseDurationMilliseconds = prometheus.NewHistogramVec( prometheus.HistogramOpts{ Name: "trakr_http_response_duration_milliseconds", diff --git a/bittorrent/udp/tracker.go b/bittorrent/udp/tracker.go index 4efa640..2ff6ac6 100644 --- a/bittorrent/udp/tracker.go +++ b/bittorrent/udp/tracker.go @@ -31,6 +31,11 @@ import ( "github.com/jzelinskie/trakr/bittorrent/udp/bytepool" ) +func init() { + prometheus.MustRegister(promResponseDurationMilliseconds) + recordResponseDuration("action", nil, time.Second) +} + var promResponseDurationMilliseconds = prometheus.NewHistogramVec( prometheus.HistogramOpts{ Name: "trakr_udp_response_duration_milliseconds", diff --git a/cmd/trakr/main.go b/cmd/trakr/main.go index 65409c5..f250c87 100644 --- a/cmd/trakr/main.go +++ b/cmd/trakr/main.go @@ -2,17 +2,57 @@ package main import ( "errors" + "io/ioutil" "log" + "net/http" "os" "os/signal" "runtime/pprof" "syscall" + "github.com/prometheus/client_golang/prometheus" "github.com/spf13/cobra" + "gopkg.in/yaml.v2" "github.com/jzelinskie/trakr" ) +type ConfigFile struct { + Config struct { + PrometheusAddr string `yaml:"prometheus_addr"` + trakr.MultiTracker + } `yaml:"trakr"` +} + +// ParseConfigFile returns a new ConfigFile given the path to a YAML +// configuration file. +// +// It supports relative and absolute paths and environment variables. +func ParseConfigFile(path string) (*ConfigFile, error) { + if path == "" { + return nil, errors.New("no config path specified") + } + + f, err := os.Open(os.ExpandEnv(path)) + if err != nil { + return nil, err + } + defer f.Close() + + contents, err := ioutil.ReadAll(f) + if err != nil { + return nil, err + } + + var cfgFile ConfigFile + err = yaml.Unmarshal(contents, &cfgFile) + if err != nil { + return nil, err + } + + return &cfgFile, nil +} + func main() { var configFilePath string var cpuProfilePath string @@ -33,19 +73,30 @@ func main() { defer pprof.StopCPUProfile() } - mt, err := trakr.MultiTrackerFromFile(configFilePath) + configFile, err := ParseConfigFile(configFilePath) if err != nil { return errors.New("failed to read config: " + err.Error()) } + go func() { + promServer := http.Server{ + Addr: configFile.Config.PrometheusAddr, + Handler: prometheus.Handler(), + } + log.Println("started serving prometheus stats on", configFile.Config.PrometheusAddr) + if err := promServer.ListenAndServe(); err != nil { + log.Fatal(err) + } + }() + go func() { shutdown := make(chan os.Signal) signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM) <-shutdown - mt.Stop() + configFile.Config.MultiTracker.Stop() }() - if err := mt.ListenAndServe(); err != nil { + if err := configFile.Config.MultiTracker.ListenAndServe(); err != nil { return errors.New("failed to cleanly shutdown: " + err.Error()) } diff --git a/example_config.yaml b/example_config.yaml index 66e8ef1..b6e7760 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -1,7 +1,10 @@ trakr: announce_interval: 15m + gc_interval: 15m + gc_expiration: 15m allow_ip_spoofing: true default_num_want: 50 + prometheus_addr: localhost:6880 http: addr: 0.0.0.0:6881 diff --git a/tracker.go b/tracker.go index 27d2686..a4dab7c 100644 --- a/tracker.go +++ b/tracker.go @@ -18,15 +18,10 @@ package trakr import ( - "errors" - "io" - "io/ioutil" - "os" "time" "github.com/jzelinskie/trakr/bittorrent/http" "github.com/jzelinskie/trakr/bittorrent/udp" - "gopkg.in/yaml.v2" ) // GenericConfig is a block of configuration who's structure is unknown. @@ -49,61 +44,28 @@ type MultiTracker struct { peerStore PeerStore httpTracker http.Tracker udpTracker udp.Tracker -} - -// decodeConfigFile unmarshals an io.Reader into a new MultiTracker. -func decodeConfigFile(r io.Reader) (*MultiTracker, error) { - contents, err := ioutil.ReadAll(r) - if err != nil { - return nil, err - } - - cfgFile := struct { - mt MultiTracker `yaml:"trakr"` - }{} - err = yaml.Unmarshal(contents, cfgFile) - if err != nil { - return nil, err - } - - return &cfgFile.mt, nil -} - -// MultiTrackerFromFile returns a new MultiTracker given the path to a YAML -// configuration file. -// -// It supports relative and absolute paths and environment variables. -func MultiTrackerFromFile(path string) (*MultiTracker, error) { - if path == "" { - return nil, errors.New("no config path specified") - } - - f, err := os.Open(os.ExpandEnv(path)) - if err != nil { - return nil, err - } - defer f.Close() - - cfg, err := decodeConfigFile(f) - if err != nil { - return nil, err - } - - return cfg, nil + closing chan struct{} } // Stop provides a thread-safe way to shutdown a currently running // MultiTracker. func (t *MultiTracker) Stop() { + close(t.closing) } // ListenAndServe listens on the protocols and addresses specified in the // HTTPConfig and UDPConfig then blocks serving BitTorrent requests until // t.Stop() is called or an error is returned. func (t *MultiTracker) ListenAndServe() error { + t.closing = make(chan struct{}) // Build an TrackerFuncs from the PreHooks and PostHooks. // Create a PeerStore instance. // Create a HTTP Tracker instance. // Create a UDP Tracker instance. + select { + case <-t.closing: + return nil + } + return nil }