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 (
|
||||
"errors"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
|
|
@ -12,7 +11,6 @@ import (
|
|||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/spf13/cobra"
|
||||
"gopkg.in/yaml.v2"
|
||||
|
||||
httpfrontend "github.com/chihaya/chihaya/frontend/http"
|
||||
udpfrontend "github.com/chihaya/chihaya/frontend/udp"
|
||||
|
|
@ -20,170 +18,131 @@ import (
|
|||
"github.com/chihaya/chihaya/storage/memory"
|
||||
)
|
||||
|
||||
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")
|
||||
func rootCmdRun(cmd *cobra.Command, args []string) error {
|
||||
cpuProfilePath, _ := cmd.Flags().GetString("cpuprofile")
|
||||
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()
|
||||
}
|
||||
|
||||
f, err := os.Open(os.ExpandEnv(path))
|
||||
configFilePath, _ := cmd.Flags().GetString("config")
|
||||
configFile, err := ParseConfigFile(configFilePath)
|
||||
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 {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
|
||||
var cfgFile ConfigFile
|
||||
err = yaml.Unmarshal(contents, &cfgFile)
|
||||
// TODO create Hooks
|
||||
logic := middleware.NewLogic(cfg.Config, peerStore, nil, nil, nil, 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() {
|
||||
var configFilePath string
|
||||
var cpuProfilePath string
|
||||
|
||||
var rootCmd = &cobra.Command{
|
||||
Use: "chihaya",
|
||||
Short: "BitTorrent Tracker",
|
||||
Long: "A customizible, multi-protocol BitTorrent Tracker",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if err := func() error {
|
||||
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 {
|
||||
if err := rootCmdRun(cmd, args); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
rootCmd.Flags().StringVar(&configFilePath, "config", "/etc/chihaya.yaml", "location of configuration file")
|
||||
rootCmd.Flags().StringVarP(&cpuProfilePath, "cpuprofile", "", "", "location to save a CPU profile")
|
||||
rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file")
|
||||
rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile")
|
||||
|
||||
if err := rootCmd.Execute(); err != nil {
|
||||
log.Fatal(err)
|
||||
|
|
|
|||
Loading…
Reference in a new issue