diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go index 2ff789a..bbc2b0c 100644 --- a/cmd/chihaya/main.go +++ b/cmd/chihaya/main.go @@ -119,7 +119,11 @@ func rootCmdRun(cmd *cobra.Command, args []string) error { } } - // TODO(jzelinskie): stop hooks here + // Stop hooks. + errs := logic.Stop() + for _, err := range errs { + errChan <- err + } close(errChan) }() diff --git a/middleware/jwt/jwt.go b/middleware/jwt/jwt.go index 7fcfe3d..c97a567 100644 --- a/middleware/jwt/jwt.go +++ b/middleware/jwt/jwt.go @@ -23,6 +23,7 @@ import ( "github.com/chihaya/chihaya/bittorrent" "github.com/chihaya/chihaya/middleware" + "github.com/chihaya/chihaya/stopper" ) var ( @@ -94,8 +95,18 @@ func NewHook(cfg Config) middleware.Hook { return h } -func (h *hook) Stop() { - close(h.closing) +func (h *hook) Stop() <-chan error { + select { + case <-h.closing: + return stopper.AlreadyStopped + default: + } + c := make(chan error) + go func() { + close(h.closing) + close(c) + }() + return c } func (h *hook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (context.Context, error) { diff --git a/middleware/middleware.go b/middleware/middleware.go index 77b246f..5321a62 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -9,6 +9,7 @@ import ( log "github.com/Sirupsen/logrus" "github.com/chihaya/chihaya/bittorrent" "github.com/chihaya/chihaya/frontend" + "github.com/chihaya/chihaya/stopper" "github.com/chihaya/chihaya/storage" ) @@ -94,3 +95,25 @@ func (l *Logic) AfterScrape(ctx context.Context, req *bittorrent.ScrapeRequest, } } } + +// Stop stops the Logic. +// +// This stops any hooks that implement stopper.Stopper. +func (l *Logic) Stop() []error { + stopGroup := stopper.NewStopGroup() + for _, hook := range l.preHooks { + stoppable, ok := hook.(stopper.Stopper) + if ok { + stopGroup.Add(stoppable) + } + } + + for _, hook := range l.postHooks { + stoppable, ok := hook.(stopper.Stopper) + if ok { + stopGroup.Add(stoppable) + } + } + + return stopGroup.Stop() +}