From ea0dba3a3d9c84d41bc9d4565deb8d8d82a52c7d Mon Sep 17 00:00:00 2001 From: Jimmy Zelinskie Date: Sat, 29 Apr 2017 22:29:27 -0400 Subject: [PATCH] cmd/chihaya: refactor root run command This change refactors a bunch of the state of execution into its own object. It also attempts to simplify stopping and adjusts some other packages to integrate with the stopper interface. Fixes #309. --- cmd/chihaya/config.go | 91 ++++++------ cmd/chihaya/main.go | 281 ++++++++++++++++------------------- frontend/http/frontend.go | 130 ++++++++-------- frontend/udp/frontend.go | 52 +++++-- storage/memory/peer_store.go | 88 ++++++----- 5 files changed, 330 insertions(+), 312 deletions(-) diff --git a/cmd/chihaya/config.go b/cmd/chihaya/config.go index bc50f1e..4fb60e2 100644 --- a/cmd/chihaya/config.go +++ b/cmd/chihaya/config.go @@ -21,52 +21,21 @@ type hookConfig struct { Config interface{} `yaml:"config"` } -// ConfigFile represents a namespaced YAML configation file. -type ConfigFile struct { - MainConfigBlock struct { - middleware.Config `yaml:",inline"` - PrometheusAddr string `yaml:"prometheus_addr"` - HTTPConfig httpfrontend.Config `yaml:"http"` - UDPConfig udpfrontend.Config `yaml:"udp"` - Storage memory.Config `yaml:"storage"` - PreHooks []hookConfig `yaml:"prehooks"` - PostHooks []hookConfig `yaml:"posthooks"` - } `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 +// Config represents the configuration used for executing Chihaya. +type Config struct { + middleware.Config `yaml:",inline"` + PrometheusAddr string `yaml:"prometheus_addr"` + HTTPConfig httpfrontend.Config `yaml:"http"` + UDPConfig udpfrontend.Config `yaml:"udp"` + Storage memory.Config `yaml:"storage"` + PreHooks []hookConfig `yaml:"prehooks"` + PostHooks []hookConfig `yaml:"posthooks"` } // CreateHooks creates instances of Hooks for all of the PreHooks and PostHooks -// configured in a ConfigFile. -func (cfg ConfigFile) CreateHooks() (preHooks, postHooks []middleware.Hook, err error) { - for _, hookCfg := range cfg.MainConfigBlock.PreHooks { +// configured in a Config. +func (cfg Config) CreateHooks() (preHooks, postHooks []middleware.Hook, err error) { + for _, hookCfg := range cfg.PreHooks { cfgBytes, err := yaml.Marshal(hookCfg.Config) if err != nil { panic("failed to remarshal valid YAML") @@ -109,10 +78,44 @@ func (cfg ConfigFile) CreateHooks() (preHooks, postHooks []middleware.Hook, err } } - for _, hookCfg := range cfg.MainConfigBlock.PostHooks { + for _, hookCfg := range cfg.PostHooks { switch hookCfg.Name { } } return } + +// ConfigFile represents a namespaced YAML configation file. +type ConfigFile struct { + Chihaya Config `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 +} diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go index 7536a6b..ab2414a 100644 --- a/cmd/chihaya/main.go +++ b/cmd/chihaya/main.go @@ -2,24 +2,125 @@ package main import ( "errors" - "net/http" "os" "os/signal" "runtime/pprof" "syscall" log "github.com/Sirupsen/logrus" - "github.com/prometheus/client_golang/prometheus" "github.com/spf13/cobra" - httpfrontend "github.com/chihaya/chihaya/frontend/http" - udpfrontend "github.com/chihaya/chihaya/frontend/udp" + "github.com/chihaya/chihaya/frontend/http" + "github.com/chihaya/chihaya/frontend/udp" "github.com/chihaya/chihaya/middleware" + "github.com/chihaya/chihaya/pkg/prometheus" + "github.com/chihaya/chihaya/pkg/stop" "github.com/chihaya/chihaya/storage" "github.com/chihaya/chihaya/storage/memory" ) -func rootCmdRun(cmd *cobra.Command, args []string) error { +// Run represents the state of a running instance of Chihaya. +type Run struct { + configFilePath string + peerStore storage.PeerStore + logic *middleware.Logic + sg *stop.Group +} + +// NewRun runs an instance of Chihaya. +func NewRun(configFilePath string) (*Run, error) { + r := &Run{ + configFilePath: configFilePath, + } + + return r, r.Start() +} + +// Start begins an instance of Chihaya. +func (r *Run) Start() error { + configFile, err := ParseConfigFile(r.configFilePath) + if err != nil { + return errors.New("failed to read config: " + err.Error()) + } + + chihayaCfg := configFile.Chihaya + preHooks, postHooks, err := chihayaCfg.CreateHooks() + if err != nil { + return errors.New("failed to validate hook config: " + err.Error()) + } + + r.sg = stop.NewGroup() + r.sg.Add(prometheus.NewServer(chihayaCfg.PrometheusAddr)) + + r.peerStore, err = memory.New(chihayaCfg.Storage) + if err != nil { + return errors.New("failed to create memory storage: " + err.Error()) + } + + r.logic = middleware.NewLogic(chihayaCfg.Config, r.peerStore, preHooks, postHooks) + + if chihayaCfg.HTTPConfig.Addr != "" { + httpfe, err := http.NewFrontend(r.logic, chihayaCfg.HTTPConfig) + if err != nil { + return err + } + r.sg.Add(httpfe) + } + + if chihayaCfg.UDPConfig.Addr != "" { + udpfe, err := udp.NewFrontend(r.logic, chihayaCfg.UDPConfig) + if err != nil { + return err + } + r.sg.Add(udpfe) + } + + return nil +} + +// Stop shuts down an instance of Chihaya. +func (r *Run) Stop() error { + log.Debug("stopping frontends and prometheus endpoint") + if errs := r.sg.Stop(); len(errs) != 0 { + errDelimiter := "; " + errStr := "failed while shutting down frontends: " + + for _, err := range errs { + errStr += err.Error() + errDelimiter + } + + // Remove the last delimiter. + errStr = errStr[0 : len(errStr)-len(errDelimiter)] + + return errors.New(errStr) + } + + log.Debug("stopping logic") + if errs := r.logic.Stop(); len(errs) != 0 { + errDelimiter := "; " + errStr := "failed while shutting down middleware: " + + for _, err := range errs { + errStr += err.Error() + errDelimiter + } + + // Remove the last delimiter. + errStr = errStr[0 : len(errStr)-len(errDelimiter)] + + return errors.New(errStr) + } + + log.Debug("stopping peer store") + if err, closed := <-r.peerStore.Stop(); !closed { + return err + } + + return nil +} + +// RunCmdFunc implements a Cobra command that runs an instance of Chihaya and +// handles reloading and shutdown via process signals. +func RunCmdFunc(cmd *cobra.Command, args []string) error { cpuProfilePath, _ := cmd.Flags().GetString("cpuprofile") if cpuProfilePath != "" { log.Infoln("enabled CPU profiling to", cpuProfilePath) @@ -31,161 +132,41 @@ func rootCmdRun(cmd *cobra.Command, args []string) error { defer pprof.StopCPUProfile() } - configFilePath, _ := cmd.Flags().GetString("config") - configFile, err := ParseConfigFile(configFilePath) + configFilePath, err := cmd.Flags().GetString("config") 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.Infoln("started serving prometheus stats on", cfg.PrometheusAddr) - if err := promServer.ListenAndServe(); err != nil { - log.Fatalln("failed to start prometheus server:", err.Error()) - } - }() - - peerStore, err := memory.New(cfg.Storage) - if err != nil { - return errors.New("failed to create memory storage: " + err.Error()) + return err } - preHooks, postHooks, err := configFile.CreateHooks() + r, err := NewRun(configFilePath) if err != nil { - return errors.New("failed to create hooks: " + err.Error()) + return err } - logic := middleware.NewLogic(cfg.Config, peerStore, preHooks, postHooks) - - errChan := make(chan error) - - httpFrontend, udpFrontend := startFrontends(cfg.HTTPConfig, cfg.UDPConfig, logic, errChan) - - shutdown := make(chan struct{}) quit := make(chan os.Signal) - restart := make(chan os.Signal) signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) - signal.Notify(restart, syscall.SIGUSR1) - go func() { - for { - select { - case <-restart: - log.Info("Got signal to restart") + reload := make(chan os.Signal) + signal.Notify(reload, syscall.SIGUSR1) - // Reload config - configFile, err = ParseConfigFile(configFilePath) - if err != nil { - log.Error("failed to read config: " + err.Error()) - } - cfg = configFile.MainConfigBlock - - preHooks, postHooks, err = configFile.CreateHooks() - if err != nil { - log.Error("failed to create hooks: " + err.Error()) - } - - // Stop frontends and logic - stopFrontends(udpFrontend, httpFrontend) - - stopLogic(logic, errChan) - - // Restart - log.Debug("Restarting logic") - logic = middleware.NewLogic(cfg.Config, peerStore, preHooks, postHooks) - - log.Debug("Restarting frontends") - httpFrontend, udpFrontend = startFrontends(cfg.HTTPConfig, cfg.UDPConfig, logic, errChan) - - log.Debug("Successfully restarted") - - case <-quit: - stop(udpFrontend, httpFrontend, logic, errChan, peerStore) - case <-shutdown: - stop(udpFrontend, httpFrontend, logic, errChan, peerStore) + for { + select { + case <-reload: + log.Info("received SIGUSR1") + if err := r.Stop(); err != nil { + return err } - } - }() - - closed := false - var bufErr error - for err = range errChan { - if err != nil { - if !closed { - close(shutdown) - closed = true - } else { - log.Errorln(bufErr) + if err := r.Start(); err != nil { + return err } - bufErr = err + case <-quit: + log.Info("received SIGINT/SIGTERM") + if err := r.Stop(); err != nil { + return err + } + + return nil } } - - return bufErr -} - -func stopFrontends(udpFrontend *udpfrontend.Frontend, httpFrontend *httpfrontend.Frontend) { - log.Debug("Stopping frontends") - if udpFrontend != nil { - udpFrontend.Stop() - } - - if httpFrontend != nil { - httpFrontend.Stop() - } -} - -func stopLogic(logic *middleware.Logic, errChan chan error) { - log.Debug("Stopping logic") - errs := logic.Stop() - for _, err := range errs { - errChan <- err - } -} - -func stop(udpFrontend *udpfrontend.Frontend, httpFrontend *httpfrontend.Frontend, logic *middleware.Logic, errChan chan error, peerStore storage.PeerStore) { - stopFrontends(udpFrontend, httpFrontend) - stopLogic(logic, errChan) - - // Stop storage - log.Debug("Stopping storage") - for err := range peerStore.Stop() { - if err != nil { - errChan <- err - } - } - - close(errChan) -} - -func startFrontends(httpConfig httpfrontend.Config, udpConfig udpfrontend.Config, logic *middleware.Logic, errChan chan<- error) (httpFrontend *httpfrontend.Frontend, udpFrontend *udpfrontend.Frontend) { - if httpConfig.Addr != "" { - httpFrontend = httpfrontend.NewFrontend(logic, httpConfig) - - go func() { - log.Infoln("started serving HTTP on", httpConfig.Addr) - if err := httpFrontend.ListenAndServe(); err != nil { - errChan <- err - } - }() - } - - if udpConfig.Addr != "" { - udpFrontend = udpfrontend.NewFrontend(logic, udpConfig) - - go func() { - log.Infoln("started serving UDP on", udpConfig.Addr) - if err := udpFrontend.ListenAndServe(); err != nil { - errChan <- err - } - }() - } - - return } func main() { @@ -200,17 +181,13 @@ func main() { log.Debugln("debug logging enabled") } }, - Run: func(cmd *cobra.Command, args []string) { - if err := rootCmdRun(cmd, args); err != nil { - log.Fatal(err) - } - }, + RunE: RunCmdFunc, } rootCmd.Flags().String("config", "/etc/chihaya.yaml", "location of configuration file") rootCmd.Flags().String("cpuprofile", "", "location to save a CPU profile") rootCmd.Flags().Bool("debug", false, "enable debug logging") if err := rootCmd.Execute(); err != nil { - log.Fatal(err) + log.Fatal("failed when executing root cobra command: " + err.Error()) } } diff --git a/frontend/http/frontend.go b/frontend/http/frontend.go index 86a4a62..c2dc3e4 100644 --- a/frontend/http/frontend.go +++ b/frontend/http/frontend.go @@ -5,7 +5,6 @@ package http import ( "context" "crypto/tls" - "errors" "net" "net/http" "time" @@ -72,94 +71,95 @@ type Config struct { TLSKeyPath string `yaml:"tls_key_path"` } -// Frontend holds the state of an HTTP BitTorrent Frontend. +// Frontend represents the state of an HTTP BitTorrent Frontend. type Frontend struct { - s *http.Server + srv *http.Server + tlsCfg *tls.Config logic frontend.TrackerLogic Config } -// NewFrontend allocates a new instance of a Frontend. -func NewFrontend(logic frontend.TrackerLogic, cfg Config) *Frontend { - return &Frontend{ +// NewFrontend creates a new instance of an HTTP Frontend that asynchronously +// serves requests. +func NewFrontend(logic frontend.TrackerLogic, cfg Config) (*Frontend, error) { + f := &Frontend{ logic: logic, Config: cfg, } + + // If TLS is enabled, create a key pair. + if cfg.TLSCertPath != "" && cfg.TLSKeyPath != "" { + var err error + f.tlsCfg = &tls.Config{ + Certificates: make([]tls.Certificate, 1), + } + f.tlsCfg.Certificates[0], err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath) + if err != nil { + return nil, err + } + } + + go func() { + if err := f.listenAndServe(); err != nil { + log.Fatal("failed while serving http: " + err.Error()) + } + }() + + return f, nil } // Stop provides a thread-safe way to shutdown a currently running Frontend. -func (t *Frontend) Stop() { - if err := t.s.Shutdown(context.Background()); err != nil { - log.Warn("Error shutting down HTTP frontend:", err) - } +func (f *Frontend) Stop() <-chan error { + c := make(chan error) + go func() { + if err := f.srv.Shutdown(context.Background()); err != nil { + c <- err + } else { + close(c) + } + }() + + return c } -func (t *Frontend) handler() http.Handler { +func (f *Frontend) handler() http.Handler { router := httprouter.New() - router.GET("/announce", t.announceRoute) - router.GET("/scrape", t.scrapeRoute) + router.GET("/announce", f.announceRoute) + router.GET("/scrape", f.scrapeRoute) return router } -// ListenAndServe listens on the TCP network address t.Addr and blocks serving -// BitTorrent requests until t.Stop() is called or an error is returned. -func (t *Frontend) ListenAndServe() error { - t.s = &http.Server{ - Addr: t.Addr, - Handler: t.handler(), - ReadTimeout: t.ReadTimeout, - WriteTimeout: t.WriteTimeout, - ConnState: func(conn net.Conn, state http.ConnState) { - switch state { - case http.StateNew: - //stats.RecordEvent(stats.AcceptedConnection) - - case http.StateClosed: - //stats.RecordEvent(stats.ClosedConnection) - - case http.StateHijacked: - panic("http: connection impossibly hijacked") - - // Ignore the following cases. - case http.StateActive, http.StateIdle: - - default: - panic("http: connection transitioned to unknown state") - } - }, - } - t.s.SetKeepAlivesEnabled(false) - - // If TLS is enabled, create a key pair and add it to the HTTP server. - if t.Config.TLSCertPath != "" && t.Config.TLSKeyPath != "" { - var err error - tlsCfg := &tls.Config{ - Certificates: make([]tls.Certificate, 1), - } - tlsCfg.Certificates[0], err = tls.LoadX509KeyPair(t.Config.TLSCertPath, t.Config.TLSKeyPath) - if err != nil { - return err - } - t.s.TLSConfig = tlsCfg +// listenAndServe blocks while listening and serving HTTP BitTorrent requests +// until Stop() is called or an error is returned. +func (f *Frontend) listenAndServe() error { + f.srv = &http.Server{ + Addr: f.Addr, + TLSConfig: f.tlsCfg, + Handler: f.handler(), + ReadTimeout: f.ReadTimeout, + WriteTimeout: f.WriteTimeout, } - // Start the HTTP server and gracefully handle any network errors. - if err := t.s.ListenAndServe(); err != nil && err != http.ErrServerClosed { - return errors.New("http: failed to run HTTP server: " + err.Error()) + // Disable KeepAlives. + f.srv.SetKeepAlivesEnabled(false) + + // Start the HTTP server. + if err := f.srv.ListenAndServe(); err != http.ErrServerClosed { + return err } return nil } -// announceRoute parses and responds to an Announce by using t.TrackerLogic. -func (t *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { +// announceRoute parses and responds to an Announce. +func (f *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { var err error start := time.Now() var af *bittorrent.AddressFamily defer func() { recordResponseDuration("announce", af, err, time.Since(start)) }() - req, err := ParseAnnounce(r, t.RealIPHeader, t.AllowIPSpoofing) + req, err := ParseAnnounce(r, f.RealIPHeader, f.AllowIPSpoofing) if err != nil { WriteError(w, err) return @@ -167,7 +167,7 @@ func (t *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httpr af = new(bittorrent.AddressFamily) *af = req.IP.AddressFamily - resp, err := t.logic.HandleAnnounce(context.Background(), req) + resp, err := f.logic.HandleAnnounce(context.Background(), req) if err != nil { WriteError(w, err) return @@ -179,11 +179,11 @@ func (t *Frontend) announceRoute(w http.ResponseWriter, r *http.Request, _ httpr return } - go t.logic.AfterAnnounce(context.Background(), req, resp) + go f.logic.AfterAnnounce(context.Background(), req, resp) } -// scrapeRoute parses and responds to a Scrape by using t.TrackerLogic. -func (t *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { +// scrapeRoute parses and responds to a Scrape. +func (f *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { var err error start := time.Now() var af *bittorrent.AddressFamily @@ -215,7 +215,7 @@ func (t *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprou af = new(bittorrent.AddressFamily) *af = req.AddressFamily - resp, err := t.logic.HandleScrape(context.Background(), req) + resp, err := f.logic.HandleScrape(context.Background(), req) if err != nil { WriteError(w, err) return @@ -227,5 +227,5 @@ func (t *Frontend) scrapeRoute(w http.ResponseWriter, r *http.Request, _ httprou return } - go t.logic.AfterScrape(context.Background(), req, resp) + go f.logic.AfterScrape(context.Background(), req, resp) } diff --git a/frontend/udp/frontend.go b/frontend/udp/frontend.go index 441af64..6678310 100644 --- a/frontend/udp/frontend.go +++ b/frontend/udp/frontend.go @@ -17,6 +17,7 @@ import ( "github.com/chihaya/chihaya/bittorrent" "github.com/chihaya/chihaya/frontend" "github.com/chihaya/chihaya/frontend/udp/bytepool" + "github.com/chihaya/chihaya/pkg/stop" ) var allowedGeneratedPrivateKeyRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890") @@ -82,8 +83,9 @@ type Frontend struct { Config } -// NewFrontend allocates a new instance of a Frontend. -func NewFrontend(logic frontend.TrackerLogic, cfg Config) *Frontend { +// NewFrontend creates a new instance of an UDP Frontend that asynchronously +// serves requests. +func NewFrontend(logic frontend.TrackerLogic, cfg Config) (*Frontend, error) { // Generate a private key if one isn't provided by the user. if cfg.PrivateKey == "" { rand.Seed(time.Now().UnixNano()) @@ -96,40 +98,68 @@ func NewFrontend(logic frontend.TrackerLogic, cfg Config) *Frontend { log.Warn("UDP private key was not provided, using generated key: ", cfg.PrivateKey) } - return &Frontend{ + f := &Frontend{ closing: make(chan struct{}), logic: logic, Config: cfg, } + + go func() { + if err := f.listenAndServe(); err != nil { + log.Fatal("failed while serving udp: " + err.Error()) + } + }() + + return f, nil } // Stop provides a thread-safe way to shutdown a currently running Frontend. -func (t *Frontend) Stop() { - close(t.closing) - t.socket.SetReadDeadline(time.Now()) - t.wg.Wait() +func (t *Frontend) Stop() <-chan error { + select { + case <-t.closing: + return stop.AlreadyStopped + default: + } + + c := make(chan error) + go func() { + close(t.closing) + t.socket.SetReadDeadline(time.Now()) + t.wg.Wait() + if err := t.socket.Close(); err != nil { + c <- err + } else { + close(c) + } + }() + + return c } -// ListenAndServe listens on the UDP network address t.Addr and blocks serving -// BitTorrent requests until t.Stop() is called or an error is returned. -func (t *Frontend) ListenAndServe() error { +// listenAndServe blocks while listening and serving UDP BitTorrent requests +// until Stop() is called or an error is returned. +func (t *Frontend) listenAndServe() error { udpAddr, err := net.ResolveUDPAddr("udp", t.Addr) if err != nil { return err } + log.Debugf("listening on udp socket") t.socket, err = net.ListenUDP("udp", udpAddr) if err != nil { return err } - defer t.socket.Close() pool := bytepool.New(2048) + t.wg.Add(1) + defer t.wg.Done() + for { // Check to see if we need to shutdown. select { case <-t.closing: + log.Debugf("returning from udp listen&serve") return nil default: } diff --git a/storage/memory/peer_store.go b/storage/memory/peer_store.go index d5b7482..1b300eb 100644 --- a/storage/memory/peer_store.go +++ b/storage/memory/peer_store.go @@ -64,18 +64,20 @@ func New(cfg Config) (storage.PeerStore, error) { } ps := &peerStore{ - shards: make([]*peerShard, shardCount*2), - closed: make(chan struct{}), + shards: make([]*peerShard, shardCount*2), + closing: make(chan struct{}), } for i := 0; i < shardCount*2; i++ { ps.shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)} } + ps.wg.Add(1) go func() { + defer ps.wg.Done() for { select { - case <-ps.closed: + case <-ps.closing: return case <-time.After(cfg.GarbageCollectionInterval): before := time.Now().Add(-cfg.PeerLifetime) @@ -102,19 +104,20 @@ type swarm struct { } type peerStore struct { - shards []*peerShard - closed chan struct{} + shards []*peerShard + closing chan struct{} + wg sync.WaitGroup } var _ storage.PeerStore = &peerStore{} -func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, af bittorrent.AddressFamily) uint32 { +func (ps *peerStore) shardIndex(infoHash bittorrent.InfoHash, af bittorrent.AddressFamily) uint32 { // There are twice the amount of shards specified by the user, the first // half is dedicated to IPv4 swarms and the second half is dedicated to // IPv6 swarms. - idx := binary.BigEndian.Uint32(infoHash[:4]) % (uint32(len(s.shards)) / 2) + idx := binary.BigEndian.Uint32(infoHash[:4]) % (uint32(len(ps.shards)) / 2) if af == bittorrent.IPv6 { - idx += uint32(len(s.shards) / 2) + idx += uint32(len(ps.shards) / 2) } return idx } @@ -146,16 +149,16 @@ func decodePeerKey(pk serializedPeer) bittorrent.Peer { return peer } -func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error { +func (ps *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } pk := newPeerKey(p) - shard := s.shards[s.shardIndex(ih, p.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, p.IP.AddressFamily)] shard.Lock() if _, ok := shard.swarms[ih]; !ok { @@ -172,16 +175,16 @@ func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error { return nil } -func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error { +func (ps *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } pk := newPeerKey(p) - shard := s.shards[s.shardIndex(ih, p.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, p.IP.AddressFamily)] shard.Lock() if _, ok := shard.swarms[ih]; !ok { @@ -205,16 +208,16 @@ func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) erro return nil } -func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { +func (ps *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } pk := newPeerKey(p) - shard := s.shards[s.shardIndex(ih, p.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, p.IP.AddressFamily)] shard.Lock() if _, ok := shard.swarms[ih]; !ok { @@ -231,16 +234,16 @@ func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error return nil } -func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { +func (ps *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } pk := newPeerKey(p) - shard := s.shards[s.shardIndex(ih, p.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, p.IP.AddressFamily)] shard.Lock() if _, ok := shard.swarms[ih]; !ok { @@ -264,16 +267,16 @@ func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) err return nil } -func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { +func (ps *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } pk := newPeerKey(p) - shard := s.shards[s.shardIndex(ih, p.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, p.IP.AddressFamily)] shard.Lock() if _, ok := shard.swarms[ih]; !ok { @@ -292,14 +295,14 @@ func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) e return nil } -func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant int, announcer bittorrent.Peer) (peers []bittorrent.Peer, err error) { +func (ps *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant int, announcer bittorrent.Peer) (peers []bittorrent.Peer, err error) { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } - shard := s.shards[s.shardIndex(ih, announcer.IP.AddressFamily)] + shard := ps.shards[ps.shardIndex(ih, announcer.IP.AddressFamily)] shard.RLock() if _, ok := shard.swarms[ih]; !ok { @@ -354,15 +357,15 @@ func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant i return } -func (s *peerStore) ScrapeSwarm(ih bittorrent.InfoHash, addressFamily bittorrent.AddressFamily) (resp bittorrent.Scrape) { +func (ps *peerStore) ScrapeSwarm(ih bittorrent.InfoHash, addressFamily bittorrent.AddressFamily) (resp bittorrent.Scrape) { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } resp.InfoHash = ih - shard := s.shards[s.shardIndex(ih, addressFamily)] + shard := ps.shards[ps.shardIndex(ih, addressFamily)] shard.RLock() if _, ok := shard.swarms[ih]; !ok { @@ -382,9 +385,9 @@ func (s *peerStore) ScrapeSwarm(ih bittorrent.InfoHash, addressFamily bittorrent // // This function must be able to execute while other methods on this interface // are being executed in parallel. -func (s *peerStore) collectGarbage(cutoff time.Time) error { +func (ps *peerStore) collectGarbage(cutoff time.Time) error { select { - case <-s.closed: + case <-ps.closing: panic("attempted to interact with stopped memory store") default: } @@ -392,7 +395,7 @@ func (s *peerStore) collectGarbage(cutoff time.Time) error { var ihDelta float64 cutoffUnix := cutoff.UnixNano() start := time.Now() - for _, shard := range s.shards { + for _, shard := range ps.shards { shard.RLock() var infohashes []bittorrent.InfoHash for ih := range shard.swarms { @@ -440,16 +443,21 @@ func (s *peerStore) collectGarbage(cutoff time.Time) error { return nil } -func (s *peerStore) Stop() <-chan error { - toReturn := make(chan error) +func (ps *peerStore) Stop() <-chan error { + c := make(chan error) go func() { - shards := make([]*peerShard, len(s.shards)) - for i := 0; i < len(s.shards); i++ { + close(ps.closing) + ps.wg.Wait() + + // Explicitly deallocate our storage. + shards := make([]*peerShard, len(ps.shards)) + for i := 0; i < len(ps.shards); i++ { shards[i] = &peerShard{swarms: make(map[bittorrent.InfoHash]swarm)} } - s.shards = shards - close(s.closed) - close(toReturn) + ps.shards = shards + + close(c) }() - return toReturn + + return c }