diff --git a/p2p/host/relaysvc/relay_test.go b/p2p/host/relaysvc/relay_test.go index d995cac200..4117e11966 100644 --- a/p2p/host/relaysvc/relay_test.go +++ b/p2p/host/relaysvc/relay_test.go @@ -34,7 +34,7 @@ func TestReachabilityChangeEvent(t *testing.T) { require.Eventually( t, func() bool { rmgr.mutex.Lock(); defer rmgr.mutex.Unlock(); return rmgr.relay == nil }, - 1*time.Second, + 3*time.Second, 100*time.Millisecond, "relay should be nil on private reachability") @@ -45,7 +45,7 @@ func TestReachabilityChangeEvent(t *testing.T) { require.Eventually( t, func() bool { rmgr.mutex.Lock(); defer rmgr.mutex.Unlock(); return rmgr.relay == nil }, - 1*time.Second, + 3*time.Second, 100*time.Millisecond, "relay should be nil on unknown reachability") diff --git a/p2p/protocol/circuitv2/relay/relay.go b/p2p/protocol/circuitv2/relay/relay.go index 5a5e97bd75..77c592e8b7 100644 --- a/p2p/protocol/circuitv2/relay/relay.go +++ b/p2p/protocol/circuitv2/relay/relay.go @@ -51,6 +51,7 @@ type Relay struct { constraints *constraints scope network.ResourceScopeSpan notifiee network.Notifiee + wg sync.WaitGroup mx sync.Mutex rsvp map[peer.ID]time.Time @@ -98,6 +99,8 @@ func New(h host.Host, opts ...Option) (*Relay, error) { h.SetStreamHandler(proto.ProtoIDv2Hop, r.handleStream) r.notifiee = &network.NotifyBundle{DisconnectedF: r.disconnected} h.Network().Notify(r.notifiee) + + r.wg.Add(1) go r.background() return r, nil @@ -105,17 +108,18 @@ func New(h host.Host, opts ...Option) (*Relay, error) { func (r *Relay) Close() error { r.mx.Lock() - defer r.mx.Unlock() if !r.closed { r.closed = true + r.mx.Unlock() + r.host.RemoveStreamHandler(proto.ProtoIDv2Hop) r.host.Network().StopNotify(r.notifiee) r.scope.Done() r.cancel() - for p := range r.rsvp { - r.host.ConnManager().UntagPeer(p, "relay-reservation") - } + r.wg.Wait() + return nil } + r.mx.Unlock() return nil } @@ -564,6 +568,8 @@ func (r *Relay) background() { case <-ticker.C: r.gc() case <-r.ctx.Done(): + r.gc() + r.wg.Done() return } } @@ -576,7 +582,7 @@ func (r *Relay) gc() { now := time.Now() for p, expire := range r.rsvp { - if expire.Before(now) { + if r.closed || expire.Before(now) { delete(r.rsvp, p) r.host.ConnManager().UntagPeer(p, "relay-reservation") }