Merge pull request #201 from jzelinskie/main-split
cmd: pull closure out into its own func and add config.go
This commit is contained in:
mrd0ll4r
GitHub
commit
b5bf6b8acd
2 changed files with 156 additions and 143 deletions
54
cmd/chihaya/config.go
Normal file
54
cmd/chihaya/config.go
Normal file
|
|
@ -0,0 +1,54 @@
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"io/ioutil"
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"gopkg.in/yaml.v2"
|
||||||
|
|
||||||
|
httpfrontend "github.com/chihaya/chihaya/frontend/http"
|
||||||
|
udpfrontend "github.com/chihaya/chihaya/frontend/udp"
|
||||||
|
"github.com/chihaya/chihaya/middleware"
|
||||||
|
"github.com/chihaya/chihaya/storage/memory"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ConfigFile represents a namespaced YAML configation file.
|
||||||
|
type ConfigFile struct {
|
||||||
|
MainConfigBlock struct {
|
||||||
|
middleware.Config
|
||||||
|
PrometheusAddr string `yaml:"prometheus_addr"`
|
||||||
|
HTTPConfig httpfrontend.Config `yaml:"http"`
|
||||||
|
UDPConfig udpfrontend.Config `yaml:"udp"`
|
||||||
|
Storage memory.Config `yaml:"storage"`
|
||||||
|
} `yaml:"chihaya"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
@ -2,7 +2,6 @@ package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"errors"
|
"errors"
|
||||||
"io/ioutil"
|
|
||||||
"log"
|
"log"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
|
|
@ -12,7 +11,6 @@ import (
|
||||||
|
|
||||||
"github.com/prometheus/client_golang/prometheus"
|
"github.com/prometheus/client_golang/prometheus"
|
||||||
"github.com/spf13/cobra"
|
"github.com/spf13/cobra"
|
||||||
"gopkg.in/yaml.v2"
|
|
||||||
|
|
||||||
httpfrontend "github.com/chihaya/chihaya/frontend/http"
|
httpfrontend "github.com/chihaya/chihaya/frontend/http"
|
||||||
udpfrontend "github.com/chihaya/chihaya/frontend/udp"
|
udpfrontend "github.com/chihaya/chihaya/frontend/udp"
|
||||||
|
|
@ -20,170 +18,131 @@ import (
|
||||||
"github.com/chihaya/chihaya/storage/memory"
|
"github.com/chihaya/chihaya/storage/memory"
|
||||||
)
|
)
|
||||||
|
|
||||||
type ConfigFile struct {
|
func rootCmdRun(cmd *cobra.Command, args []string) error {
|
||||||
MainConfigBlock struct {
|
cpuProfilePath, _ := cmd.Flags().GetString("cpuprofile")
|
||||||
middleware.Config
|
if cpuProfilePath != "" {
|
||||||
PrometheusAddr string `yaml:"prometheus_addr"`
|
log.Println("enabled CPU profiling to " + cpuProfilePath)
|
||||||
HTTPConfig httpfrontend.Config `yaml:"http"`
|
f, err := os.Create(cpuProfilePath)
|
||||||
UDPConfig udpfrontend.Config `yaml:"udp"`
|
if err != nil {
|
||||||
Storage memory.Config `yaml:"storage"`
|
return err
|
||||||
} `yaml:"chihaya"`
|
}
|
||||||
}
|
pprof.StartCPUProfile(f)
|
||||||
|
defer pprof.StopCPUProfile()
|
||||||
// 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))
|
configFilePath, _ := cmd.Flags().GetString("config")
|
||||||
|
configFile, err := ParseConfigFile(configFilePath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return errors.New("failed to read config: " + err.Error())
|
||||||
}
|
}
|
||||||
defer f.Close()
|
cfg := configFile.MainConfigBlock
|
||||||
|
|
||||||
contents, err := ioutil.ReadAll(f)
|
go func() {
|
||||||
|
promServer := http.Server{
|
||||||
|
Addr: cfg.PrometheusAddr,
|
||||||
|
Handler: prometheus.Handler(),
|
||||||
|
}
|
||||||
|
log.Println("started serving prometheus stats on", cfg.PrometheusAddr)
|
||||||
|
if err := promServer.ListenAndServe(); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Force the compiler to enforce memory against the storage interface.
|
||||||
|
peerStore, err := memory.New(cfg.Storage)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
var cfgFile ConfigFile
|
// TODO create Hooks
|
||||||
err = yaml.Unmarshal(contents, &cfgFile)
|
logic := middleware.NewLogic(cfg.Config, peerStore, nil, nil, nil, nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
return &cfgFile, nil
|
shutdown := make(chan struct{})
|
||||||
|
errChan := make(chan error)
|
||||||
|
|
||||||
|
var httpFrontend *httpfrontend.Frontend
|
||||||
|
var udpFrontend *udpfrontend.Frontend
|
||||||
|
|
||||||
|
if cfg.HTTPConfig.Addr != "" {
|
||||||
|
httpFrontend = httpfrontend.NewFrontend(logic, cfg.HTTPConfig)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
log.Println("started serving HTTP on", cfg.HTTPConfig.Addr)
|
||||||
|
if err := httpFrontend.ListenAndServe(); err != nil {
|
||||||
|
errChan <- errors.New("failed to cleanly shutdown HTTP frontend: " + err.Error())
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
if cfg.UDPConfig.Addr != "" {
|
||||||
|
udpFrontend = udpfrontend.NewFrontend(logic, cfg.UDPConfig)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
log.Println("started serving UDP on", cfg.UDPConfig.Addr)
|
||||||
|
if err := udpFrontend.ListenAndServe(); err != nil {
|
||||||
|
errChan <- errors.New("failed to cleanly shutdown UDP frontend: " + err.Error())
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
sigChan := make(chan os.Signal)
|
||||||
|
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
||||||
|
go func() {
|
||||||
|
select {
|
||||||
|
case <-sigChan:
|
||||||
|
case <-shutdown:
|
||||||
|
}
|
||||||
|
|
||||||
|
if udpFrontend != nil {
|
||||||
|
udpFrontend.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
if httpFrontend != nil {
|
||||||
|
httpFrontend.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
for err := range peerStore.Stop() {
|
||||||
|
if err != nil {
|
||||||
|
errChan <- err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
close(errChan)
|
||||||
|
}()
|
||||||
|
|
||||||
|
closed := false
|
||||||
|
var bufErr error
|
||||||
|
for err = range errChan {
|
||||||
|
if err != nil {
|
||||||
|
if !closed {
|
||||||
|
close(shutdown)
|
||||||
|
closed = true
|
||||||
|
} else {
|
||||||
|
log.Println(bufErr)
|
||||||
|
}
|
||||||
|
bufErr = err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return bufErr
|
||||||
}
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
var configFilePath string
|
|
||||||
var cpuProfilePath string
|
|
||||||
|
|
||||||
var rootCmd = &cobra.Command{
|
var rootCmd = &cobra.Command{
|
||||||
Use: "chihaya",
|
Use: "chihaya",
|
||||||
Short: "BitTorrent Tracker",
|
Short: "BitTorrent Tracker",
|
||||||
Long: "A customizible, multi-protocol BitTorrent Tracker",
|
Long: "A customizible, multi-protocol BitTorrent Tracker",
|
||||||
Run: func(cmd *cobra.Command, args []string) {
|
Run: func(cmd *cobra.Command, args []string) {
|
||||||
if err := func() error {
|
if err := rootCmdRun(cmd, args); err != nil {
|
||||||
if cpuProfilePath != "" {
|
|
||||||
log.Println("enabled CPU profiling to " + cpuProfilePath)
|
|
||||||
f, err := os.Create(cpuProfilePath)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
pprof.StartCPUProfile(f)
|
|
||||||
defer pprof.StopCPUProfile()
|
|
||||||
}
|
|
||||||
|
|
||||||
configFile, err := ParseConfigFile(configFilePath)
|
|
||||||
if err != nil {
|
|
||||||
return errors.New("failed to read config: " + err.Error())
|
|
||||||
}
|
|
||||||
cfg := configFile.MainConfigBlock
|
|
||||||
|
|
||||||
go func() {
|
|
||||||
promServer := http.Server{
|
|
||||||
Addr: cfg.PrometheusAddr,
|
|
||||||
Handler: prometheus.Handler(),
|
|
||||||
}
|
|
||||||
log.Println("started serving prometheus stats on", cfg.PrometheusAddr)
|
|
||||||
if err := promServer.ListenAndServe(); err != nil {
|
|
||||||
log.Fatal(err)
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
|
|
||||||
// Force the compiler to enforce memory against the storage interface.
|
|
||||||
peerStore, err := memory.New(cfg.Storage)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
// TODO create Hooks
|
|
||||||
logic := middleware.NewLogic(cfg.Config, peerStore, nil, nil, nil, nil)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
shutdown := make(chan struct{})
|
|
||||||
errChan := make(chan error)
|
|
||||||
|
|
||||||
var httpFrontend *httpfrontend.Frontend
|
|
||||||
var udpFrontend *udpfrontend.Frontend
|
|
||||||
|
|
||||||
if cfg.HTTPConfig.Addr != "" {
|
|
||||||
httpFrontend = httpfrontend.NewFrontend(logic, cfg.HTTPConfig)
|
|
||||||
|
|
||||||
go func() {
|
|
||||||
log.Println("started serving HTTP on", cfg.HTTPConfig.Addr)
|
|
||||||
if err := httpFrontend.ListenAndServe(); err != nil {
|
|
||||||
errChan <- errors.New("failed to cleanly shutdown HTTP frontend: " + err.Error())
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
if cfg.UDPConfig.Addr != "" {
|
|
||||||
udpFrontend = udpfrontend.NewFrontend(logic, cfg.UDPConfig)
|
|
||||||
|
|
||||||
go func() {
|
|
||||||
log.Println("started serving UDP on", cfg.UDPConfig.Addr)
|
|
||||||
if err := udpFrontend.ListenAndServe(); err != nil {
|
|
||||||
errChan <- errors.New("failed to cleanly shutdown UDP frontend: " + err.Error())
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
sigChan := make(chan os.Signal)
|
|
||||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
|
||||||
go func() {
|
|
||||||
select {
|
|
||||||
case <-sigChan:
|
|
||||||
case <-shutdown:
|
|
||||||
}
|
|
||||||
|
|
||||||
if udpFrontend != nil {
|
|
||||||
udpFrontend.Stop()
|
|
||||||
}
|
|
||||||
|
|
||||||
if httpFrontend != nil {
|
|
||||||
httpFrontend.Stop()
|
|
||||||
}
|
|
||||||
|
|
||||||
for err := range peerStore.Stop() {
|
|
||||||
if err != nil {
|
|
||||||
errChan <- err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
close(errChan)
|
|
||||||
}()
|
|
||||||
|
|
||||||
closed := false
|
|
||||||
var bufErr error
|
|
||||||
for err = range errChan {
|
|
||||||
if err != nil {
|
|
||||||
if !closed {
|
|
||||||
close(shutdown)
|
|
||||||
closed = true
|
|
||||||
} else {
|
|
||||||
log.Println(bufErr)
|
|
||||||
}
|
|
||||||
bufErr = err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return bufErr
|
|
||||||
}(); err != nil {
|
|
||||||
log.Fatal(err)
|
log.Fatal(err)
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file")
|
||||||
rootCmd.Flags().StringVar(&configFilePath, "config", "/etc/chihaya.yaml", "location of configuration file")
|
rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile")
|
||||||
rootCmd.Flags().StringVarP(&cpuProfilePath, "cpuprofile", "", "", "location to save a CPU profile")
|
|
||||||
|
|
||||||
if err := rootCmd.Execute(); err != nil {
|
if err := rootCmd.Execute(); err != nil {
|
||||||
log.Fatal(err)
|
log.Fatal(err)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue