Skip to content

Commit

Permalink
autonat: Clean up after close (#2749)
Browse files Browse the repository at this point in the history
* Clean up after autonat

* Close
  • Loading branch information
MarcoPolo committed Mar 28, 2024
1 parent 6bb53b2 commit 86a6720
Show file tree
Hide file tree
Showing 6 changed files with 41 additions and 25 deletions.
39 changes: 23 additions & 16 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import (
"github.com/libp2p/go-libp2p/p2p/host/autonat"
"github.com/libp2p/go-libp2p/p2p/host/autorelay"
bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
blankhost "github.com/libp2p/go-libp2p/p2p/host/blank"
"github.com/libp2p/go-libp2p/p2p/host/eventbus"
"github.com/libp2p/go-libp2p/p2p/host/peerstore/pstoremem"
rcmgr "github.com/libp2p/go-libp2p/p2p/host/resource-manager"
Expand Down Expand Up @@ -466,8 +465,7 @@ func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
}

// Pull out the pieces of the config that we _actually_ care about.
// Specifically, don't set up things like autorelay, listeners,
// identify, etc.
// Specifically, don't set up things like listeners, identify, etc.
autoNatCfg := Config{
Transports: cfg.Transports,
Muxers: cfg.Muxers,
Expand All @@ -486,30 +484,39 @@ func (cfg *Config) addAutoNAT(h *bhost.BasicHost) error {
},
}

dialer, err := autoNatCfg.makeSwarm(eventbus.NewBus(), false)
if err != nil {
return err
}
dialerHost := blankhost.NewBlankHost(dialer)
fxopts, err := autoNatCfg.addTransports()
if err != nil {
dialerHost.Close()
return err
}
var dialer *swarm.Swarm

fxopts = append(fxopts,
fx.Supply(dialerHost.ID()),
fx.Supply(dialer),
fx.Provide(eventbus.NewBus),
fx.Provide(func(lifecycle fx.Lifecycle, b event.Bus) (*swarm.Swarm, error) {
lifecycle.Append(fx.Hook{
OnStop: func(context.Context) error {
return ps.Close()
}})
var err error
dialer, err = autoNatCfg.makeSwarm(b, false)
return dialer, err

}),
fx.Provide(func() crypto.PrivKey { return autonatPrivKey }),
)
app := fx.New(fxopts...)
if err := app.Err(); err != nil {
dialerHost.Close()
return err
}
// NOTE: We're dropping the blank host here but that's fine. It
// doesn't really _do_ anything and doesn't even need to be
// closed (as long as we close the underlying network).
autonatOpts = append(autonatOpts, autonat.EnableService(dialerHost.Network()))
err = app.Start(context.Background())
if err != nil {
return err
}
go func() {
<-dialer.Done() // The swarm used for autonat has closed, we can cleanup now
app.Stop(context.Background())
}()
autonatOpts = append(autonatOpts, autonat.EnableService(dialer))
}
if cfg.AutoNATConfig.ForceReachability != nil {
autonatOpts = append(autonatOpts, autonat.WithReachability(*cfg.AutoNATConfig.ForceReachability))
Expand Down
7 changes: 0 additions & 7 deletions leaky_tests/leaky_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"testing"

"github.com/libp2p/go-libp2p"
"github.com/stretchr/testify/require"
)

func TestBadTransportConstructor(t *testing.T) {
Expand All @@ -18,9 +17,3 @@ func TestBadTransportConstructor(t *testing.T) {
t.Error("expected error to contain debugging info")
}
}

func TestAutoNATService(t *testing.T) {
h, err := libp2p.New(libp2p.EnableNATService())
require.NoError(t, err)
h.Close()
}
6 changes: 6 additions & 0 deletions libp2p_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,12 @@ func TestRoutedHost(t *testing.T) {
require.Equal(t, []peer.ID{id}, mockRouter.queried)
}

func TestAutoNATService(t *testing.T) {
h, err := New(EnableNATService())
require.NoError(t, err)
h.Close()
}

func TestMain(m *testing.M) {
goleak.VerifyTestMain(
m,
Expand Down
4 changes: 2 additions & 2 deletions p2p/host/autonat/autonat.go
Original file line number Diff line number Diff line change
Expand Up @@ -431,7 +431,7 @@ func (as *AmbientAutoNAT) getPeerToProbe() peer.ID {
func (as *AmbientAutoNAT) Close() error {
as.ctxCancel()
if as.service != nil {
as.service.Disable()
return as.service.Close()
}
<-as.backgroundRunning
return nil
Expand All @@ -444,7 +444,7 @@ func (s *StaticAutoNAT) Status() network.Reachability {

func (s *StaticAutoNAT) Close() error {
if s.service != nil {
s.service.Disable()
return s.service.Close()
}
return nil
}
5 changes: 5 additions & 0 deletions p2p/host/autonat/svc.go
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,11 @@ func (as *autoNATService) Disable() {
}
}

func (as *autoNATService) Close() error {
as.Disable()
return as.config.dialer.Close()
}

func (as *autoNATService) background(ctx context.Context) {
defer close(as.backgroundRunning)

Expand Down
5 changes: 5 additions & 0 deletions p2p/net/swarm/swarm.go
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,11 @@ func (s *Swarm) Close() error {
return nil
}

// Done returns a channel that is closed when the swarm is closed.
func (s *Swarm) Done() <-chan struct{} {
return s.ctx.Done()
}

func (s *Swarm) close() {
s.ctxCancel()

Expand Down

0 comments on commit 86a6720

Please sign in to comment.