From c0dbc6123efff2562f78674298a256534c76c263 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 20 Feb 2023 19:23:23 +1300
Subject: [PATCH] config: use Fx lifecycle hooks to start AutoRelay

---
 config/config.go                         | 37 ++++++++++++------------
 config/host.go                           | 19 ++++++++++++
 p2p/host/autorelay/autorelay_test.go     |  3 +-
 p2p/host/autorelay/host.go               | 23 ---------------
 p2p/protocol/holepunch/holepunch_test.go |  3 ++
 5 files changed, 41 insertions(+), 44 deletions(-)
 create mode 100644 config/host.go
 delete mode 100644 p2p/host/autorelay/host.go

diff --git a/config/config.go b/config/config.go
index 55e5303eca..a24030812a 100644
--- a/config/config.go
+++ b/config/config.go
@@ -305,6 +305,9 @@ func (cfg *Config) newBasicHost(swrm *swarm.Swarm, eventBus event.Bus) (*bhost.B
 //
 // This function consumes the config. Do not reuse it (really!).
 func (cfg *Config) NewNode() (host.Host, error) {
+	if cfg.EnableAutoRelay && !cfg.Relay {
+		return nil, fmt.Errorf("cannot enable autorelay; relay is not enabled")
+	}
 	fxopts := []fx.Option{
 		fx.Provide(func() event.Bus {
 			return eventbus.NewBus(eventbus.WithMetricsTracer(eventbus.NewMetricsTracer(eventbus.WithRegisterer(cfg.PrometheusRegisterer))))
@@ -337,10 +340,6 @@ func (cfg *Config) NewNode() (host.Host, error) {
 
 	var h *bhost.BasicHost
 	fxopts = append(fxopts, fx.Invoke(func(ho *bhost.BasicHost) { h = ho }))
-	app := fx.New(fxopts...)
-	if err := app.Start(context.Background()); err != nil {
-		return nil, err
-	}
 
 	// Configure routing and autorelay
 	var router routing.PeerRouting
@@ -354,17 +353,22 @@ func (cfg *Config) NewNode() (host.Host, error) {
 
 	// Note: h.AddrsFactory may be changed by relayFinder, but non-relay version is
 	// used by AutoNAT below.
-	var ar *autorelay.AutoRelay
 	if cfg.EnableAutoRelay {
-		if !cfg.Relay {
-			h.Close()
-			return nil, fmt.Errorf("cannot enable autorelay; relay is not enabled")
-		}
+		fxopts = append(fxopts,
+			fx.Invoke(func(h *bhost.BasicHost, lifecycle fx.Lifecycle) (*autorelay.AutoRelay, error) {
+				ar, err := autorelay.NewAutoRelay(h, cfg.AutoRelayOpts...)
+				if err != nil {
+					return nil, err
+				}
+				lifecycle.Append(fx.StartStopHook(ar.Start, ar.Close))
+				return ar, nil
+			}),
+		)
+	}
 
-		ar, err = autorelay.NewAutoRelay(h, cfg.AutoRelayOpts...)
-		if err != nil {
-			return nil, err
-		}
+	app := fx.New(fxopts...)
+	if err := app.Start(context.Background()); err != nil {
+		return nil, err
 	}
 
 	if err := cfg.addAutoNAT(h); err != nil {
@@ -377,12 +381,7 @@ func (cfg *Config) NewNode() (host.Host, error) {
 	if router != nil {
 		ho = routed.Wrap(h, router)
 	}
-	if ar != nil {
-		arh := autorelay.NewAutoRelayHost(ho, ar)
-		arh.Start()
-		return arh, nil
-	}
-	return ho, nil
+	return &closableHost{App: app, Host: ho}, nil
 }
 
 func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
diff --git a/config/host.go b/config/host.go
new file mode 100644
index 0000000000..dd8930da39
--- /dev/null
+++ b/config/host.go
@@ -0,0 +1,19 @@
+package config
+
+import (
+	"context"
+
+	"github.com/libp2p/go-libp2p/core/host"
+
+	"go.uber.org/fx"
+)
+
+type closableHost struct {
+	*fx.App
+	host.Host
+}
+
+func (h *closableHost) Close() error {
+	_ = h.App.Stop(context.Background())
+	return h.Host.Close()
+}
diff --git a/p2p/host/autorelay/autorelay_test.go b/p2p/host/autorelay/autorelay_test.go
index b31b465b68..fceb8e7e49 100644
--- a/p2p/host/autorelay/autorelay_test.go
+++ b/p2p/host/autorelay/autorelay_test.go
@@ -444,9 +444,8 @@ func TestMinInterval(t *testing.T) {
 	)
 	defer h.Close()
 
-	cl.Add(500 * time.Millisecond)
 	// The second call to peerSource should happen after 1 second
 	require.Never(t, func() bool { return numRelays(h) > 0 }, 500*time.Millisecond, 100*time.Millisecond)
-	cl.Add(500 * time.Millisecond)
+	cl.Add(time.Second)
 	require.Eventually(t, func() bool { return numRelays(h) > 0 }, 10*time.Second, 100*time.Millisecond)
 }
diff --git a/p2p/host/autorelay/host.go b/p2p/host/autorelay/host.go
deleted file mode 100644
index c6bd9c5706..0000000000
--- a/p2p/host/autorelay/host.go
+++ /dev/null
@@ -1,23 +0,0 @@
-package autorelay
-
-import (
-	"github.com/libp2p/go-libp2p/core/host"
-)
-
-type AutoRelayHost struct {
-	host.Host
-	ar *AutoRelay
-}
-
-func (h *AutoRelayHost) Close() error {
-	_ = h.ar.Close()
-	return h.Host.Close()
-}
-
-func (h *AutoRelayHost) Start() {
-	h.ar.Start()
-}
-
-func NewAutoRelayHost(h host.Host, ar *AutoRelay) *AutoRelayHost {
-	return &AutoRelayHost{Host: h, ar: ar}
-}
diff --git a/p2p/protocol/holepunch/holepunch_test.go b/p2p/protocol/holepunch/holepunch_test.go
index 9cf57ff65c..1c210d9360 100644
--- a/p2p/protocol/holepunch/holepunch_test.go
+++ b/p2p/protocol/holepunch/holepunch_test.go
@@ -102,6 +102,9 @@ func TestNoHolePunchIfDirectConnExists(t *testing.T) {
 }
 
 func TestDirectDialWorks(t *testing.T) {
+	if race.WithRace() {
+		t.Skip("modifying manet.Private4 is racy")
+	}
 	// mark all addresses as public
 	cpy := manet.Private4
 	manet.Private4 = []*net.IPNet{}