refactor code for context fallthrough

frontend/frontend.go

@@ -4,12 +4,12 @@ package frontend
 
 import (
 	"fmt"
+	"io"
 	"sync"
 
 	"github.com/sot-tech/mochi/middleware"
 	"github.com/sot-tech/mochi/pkg/conf"
 	"github.com/sot-tech/mochi/pkg/log"
-	"github.com/sot-tech/mochi/pkg/stop"
 )
 
 var (
@@ -46,7 +46,7 @@ func RegisterBuilder(name string, b Builder) {
 
 // Frontend dummy interface for bittorrent frontends
 type Frontend interface {
-	stop.Stopper
+	io.Closer
 }
 
 // NewFrontends is a utility function for initializing Frontend-s in bulk.

frontend/http/frontend.go

@@ -18,7 +18,6 @@ import (
 	"github.com/sot-tech/mochi/pkg/conf"
 	"github.com/sot-tech/mochi/pkg/log"
 	"github.com/sot-tech/mochi/pkg/metrics"
-	"github.com/sot-tech/mochi/pkg/stop"
 )
 
 var (
@@ -81,17 +80,17 @@ type httpFE struct {
 }
 
 // NewFrontend builds and starts http bittorrent frontend from provided configuration
-func NewFrontend(c conf.MapConfig, logic *middleware.Logic) (_ frontend.Frontend, err error) {
+func NewFrontend(c conf.MapConfig, logic *middleware.Logic) (frontend.Frontend, error) {
 	var cfg Config
+	var err error
 	if err = c.Unmarshal(&cfg); err != nil {
-		return
+		return nil, err
 	}
 	if cfg, err = cfg.Validate(); err != nil {
-		return
+		return nil, err
 	}
 	if len(cfg.AnnounceRoutes) < 1 || len(cfg.ScrapeRoutes) < 1 {
-		err = errRoutesNotProvided
-		return
+		return nil, errRoutesNotProvided
 	}
 
 	f := &httpFE{
@@ -110,7 +109,7 @@ func NewFrontend(c conf.MapConfig, logic *middleware.Logic) (_ frontend.Frontend
 	if cfg.UseTLS {
 		var cert tls.Certificate
 		if cert, err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath); err != nil {
-			return
+			return nil, err
 		}
 		f.srv.TLSConfig = &tls.Config{
 			Certificates: []tls.Certificate{cert},
@@ -150,34 +149,17 @@ func NewFrontend(c conf.MapConfig, logic *middleware.Logic) (_ frontend.Frontend
 	return f, nil
 }
 
-// Stop provides a thread-safe way to shut down a currently running Frontend.
-func (f *httpFE) Stop() stop.Result {
-	c := make(stop.Channel)
-	if f.srv != nil {
-		go func() {
-			c.Done(f.srv.Shutdown(context.Background()))
-		}()
-	}
-	return c.Result()
+// Close provides a thread-safe way to shut down a currently running Frontend.
+func (f *httpFE) Close() error {
+	return f.srv.Shutdown(context.Background())
 }
 
-func injectRouteParamsToContext(ctx context.Context, ps httprouter.Params) context.Context {
-	rp := bittorrent.RouteParams{}
-	for _, p := range ps {
-		rp = append(rp, bittorrent.RouteParam{Key: p.Key, Value: p.Value})
+func httpParamsToRouteParams(in httprouter.Params) (out bittorrent.RouteParams) {
+	out = make([]bittorrent.RouteParam, 0, len(in))
+	for _, p := range in {
+		out = append(out, bittorrent.RouteParam{Key: p.Key, Value: p.Value})
 	}
-	return context.WithValue(ctx, bittorrent.RouteParamsKey, rp)
-}
-
-func remapRouteParamsToBgContext(inCtx context.Context) context.Context {
-	rp, isOk := inCtx.Value(bittorrent.RouteParamsKey).(bittorrent.RouteParams)
-	if !isOk {
-		rp = bittorrent.RouteParams{}
-	} else {
-		logger.Warn().Msg("unable to fetch route parameters, probably jammed context")
-	}
-	// FIXME: cancelable context
-	return context.WithValue(context.TODO(), bittorrent.RouteParamsKey, rp)
+	return
 }
 
 // announceRoute parses and responds to an Announce.
@@ -200,7 +182,7 @@ func (f *httpFE) announceRoute(w http.ResponseWriter, r *http.Request, ps httpro
 	}
 	addr = req.GetFirst()
 
-	ctx := injectRouteParamsToContext(r.Context(), ps)
+	ctx := bittorrent.InjectRouteParamsToContext(r.Context(), httpParamsToRouteParams(ps))
 	ctx, resp, err := f.logic.HandleAnnounce(ctx, req)
 	if err != nil {
 		WriteError(w, err)
@@ -215,8 +197,7 @@ func (f *httpFE) announceRoute(w http.ResponseWriter, r *http.Request, ps httpro
 	}
 
 	// next actions are background and should not be canceled after http writer closed
-	ctx = remapRouteParamsToBgContext(ctx)
-
+	ctx = bittorrent.RemapRouteParamsToBgContext(ctx)
 	go f.logic.AfterAnnounce(ctx, req, resp)
 }
 
@@ -239,7 +220,7 @@ func (f *httpFE) scrapeRoute(w http.ResponseWriter, r *http.Request, ps httprout
 	}
 	addr = req.GetFirst()
 
-	ctx := injectRouteParamsToContext(r.Context(), ps)
+	ctx := bittorrent.InjectRouteParamsToContext(r.Context(), httpParamsToRouteParams(ps))
 	ctx, resp, err := f.logic.HandleScrape(ctx, req)
 	if err != nil {
 		WriteError(w, err)
@@ -254,8 +235,7 @@ func (f *httpFE) scrapeRoute(w http.ResponseWriter, r *http.Request, ps httprout
 	}
 
 	// next actions are background and should not be canceled after http writer closed
-	ctx = remapRouteParamsToBgContext(ctx)
-
+	ctx = bittorrent.RemapRouteParamsToBgContext(ctx)
 	go f.logic.AfterScrape(ctx, req, resp)
 }
 

frontend/udp/frontend.go

@@ -20,7 +20,6 @@ import (
 	"github.com/sot-tech/mochi/pkg/conf"
 	"github.com/sot-tech/mochi/pkg/log"
 	"github.com/sot-tech/mochi/pkg/metrics"
-	"github.com/sot-tech/mochi/pkg/stop"
 	"github.com/sot-tech/mochi/pkg/timecache"
 )
 
@@ -81,6 +80,8 @@ type udpFE struct {
 	logic          *middleware.Logic
 	maxClockSkew   time.Duration
 	collectTimings bool
+	ctxCancel      context.CancelFunc
+	onceCloser     sync.Once
 	frontend.ParseOptions
 }
 
@@ -109,43 +110,37 @@ func NewFrontend(c conf.MapConfig, logic *middleware.Logic) (frontend.Frontend,
 	}
 
 	if f.socket, err = cfg.ListenUDP(); err == nil {
+		var ctx context.Context
+		ctx, f.ctxCancel = context.WithCancel(context.Background())
 		f.wg.Add(1)
-		go func() {
-			if err := f.serve(); err != nil {
+		go func(ctx context.Context) {
+			if err := f.serve(ctx); err != nil {
 				logger.Fatal().Err(err).Msg("server failed")
 			}
-		}()
+		}(ctx)
 	}
 
 	return f, err
 }
 
-// Stop provides a thread-safe way to shut down a currently running Frontend.
-func (t *udpFE) Stop() stop.Result {
-	select {
-	case <-t.closing:
-		return stop.AlreadyStopped
-	default:
-	}
-
-	c := make(stop.Channel)
-	go func() {
+// Close provides a thread-safe way to shut down a currently running Frontend.
+func (t *udpFE) Close() (err error) {
+	t.onceCloser.Do(func() {
 		close(t.closing)
-		var err error
 		if t.socket != nil {
+			t.ctxCancel()
 			_ = t.socket.SetReadDeadline(time.Now())
 			t.wg.Wait()
 			err = t.socket.Close()
 		}
-		c.Done(err)
-	}()
+	})
 
-	return c.Result()
+	return
 }
 
 // serve blocks while listening and serving UDP BitTorrent requests
 // until Stop() is called or an error is returned.
-func (t *udpFE) serve() error {
+func (t *udpFE) serve(ctx context.Context) error {
 	pool := bytepool.NewBytePool(2048)
 	defer t.wg.Done()
 
@@ -188,7 +183,7 @@ func (t *udpFE) serve() error {
 			if t.collectTimings && metrics.Enabled() {
 				start = time.Now()
 			}
-			action, err := t.handleRequest(
+			action, err := t.handleRequest(ctx,
 				Request{(*buffer)[:n], addr},
 				ResponseWriter{t.socket, addrPort},
 			)
@@ -218,7 +213,7 @@ func (w ResponseWriter) Write(b []byte) (int, error) {
 }
 
 // handleRequest parses and responds to a UDP Request.
-func (t *udpFE) handleRequest(r Request, w ResponseWriter) (actionName string, err error) {
+func (t *udpFE) handleRequest(ctx context.Context, r Request, w ResponseWriter) (actionName string, err error) {
 	if len(r.Packet) < 16 {
 		// Malformed, no client packets are less than 16 bytes.
 		// We explicitly return nothing in case this is a DoS attempt.
@@ -265,9 +260,9 @@ func (t *udpFE) handleRequest(r Request, w ResponseWriter) (actionName string, e
 			return
 		}
 
-		var ctx context.Context
 		var resp *bittorrent.AnnounceResponse
-		ctx, resp, err = t.logic.HandleAnnounce(context.Background(), req)
+		ctx := bittorrent.InjectRouteParamsToContext(ctx, bittorrent.RouteParams{})
+		ctx, resp, err = t.logic.HandleAnnounce(ctx, req)
 		if err != nil {
 			WriteError(w, txID, err)
 			return
@@ -275,6 +270,7 @@ func (t *udpFE) handleRequest(r Request, w ResponseWriter) (actionName string, e
 
 		WriteAnnounce(w, txID, resp, actionID == announceV6ActionID, r.IP.Is6())
 
+		ctx = bittorrent.RemapRouteParamsToBgContext(ctx)
 		go t.logic.AfterAnnounce(ctx, req, resp)
 
 	case scrapeActionID:
@@ -287,9 +283,9 @@ func (t *udpFE) handleRequest(r Request, w ResponseWriter) (actionName string, e
 			return
 		}
 
-		var ctx context.Context
 		var resp *bittorrent.ScrapeResponse
-		ctx, resp, err = t.logic.HandleScrape(context.Background(), req)
+		ctx := bittorrent.InjectRouteParamsToContext(ctx, bittorrent.RouteParams{})
+		ctx, resp, err = t.logic.HandleScrape(ctx, req)
 		if err != nil {
 			WriteError(w, txID, err)
 			return
@@ -297,6 +293,7 @@ func (t *udpFE) handleRequest(r Request, w ResponseWriter) (actionName string, e
 
 		WriteScrape(w, txID, resp)
 
+		ctx = bittorrent.RemapRouteParamsToBgContext(ctx)
 		go t.logic.AfterScrape(ctx, req, resp)
 
 	default:

frontend/udp/frontend_test.go

@@ -29,8 +29,7 @@ func TestStartStopRaceIssue437(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	errC := fe.Stop()
-	if errs := <-errC; len(errs) != 0 {
-		t.Fatal(errs)
+	if err = fe.Close(); err != nil {
+		t.Fatal(err)
 	}
 }

frontend/udp/writer.go

@@ -3,7 +3,6 @@ package udp
 import (
 	"encoding/binary"
 	"errors"
-	"fmt"
 	"io"
 	"time"
 
@@ -12,17 +11,15 @@ import (
 
 // WriteError writes the failure reason as a null-terminated string.
 func WriteError(w io.Writer, txID []byte, err error) {
-	// If the client wasn't at fault, acknowledge it.
-	var clientErr bittorrent.ClientError
-	if !errors.Is(err, &clientErr) {
-		err = fmt.Errorf("internal error occurred: %w", err)
-	}
-
 	buf := reqRespBufferPool.Get()
 	defer reqRespBufferPool.Put(buf)
 	writeHeader(buf, txID, errorActionID)
-	_, _ = buf.WriteString(err.Error())
-	_, _ = buf.WriteRune('\000')
+	// If the client wasn't at fault, acknowledge it.
+	if !errors.As(err, new(bittorrent.ClientError)) {
+		buf.WriteString("internal error occurred: ")
+	}
+	buf.WriteString(err.Error())
+	buf.WriteByte('\000')
 	_, _ = w.Write(buf.Bytes())
 }
 
@@ -50,7 +47,7 @@ func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse,
 	}
 
 	for _, peer := range peers {
-		_, _ = buf.Write(peer.Addr().AsSlice())
+		buf.Write(peer.Addr().AsSlice())
 		_ = binary.Write(buf, binary.BigEndian, peer.Port())
 	}
 
@@ -79,8 +76,7 @@ func WriteConnectionID(w io.Writer, txID, connID []byte) {
 	defer reqRespBufferPool.Put(buf)
 
 	writeHeader(buf, txID, connectActionID)
-	_, _ = buf.Write(connID)
-
+	buf.Write(connID)
 	_, _ = w.Write(buf.Bytes())
 }