Accept parent ctx in constructor, simplify TTL loop

addr_manager_ds.go

@@ -18,10 +18,10 @@ type DatastoreAddrManager struct {
 	addrSubs   map[peer.ID][]*addrSub
 }
 
-func NewDatastoreAddrManager(ds ds.Datastore, ttlInterval time.Duration) *DatastoreAddrManager {
+func NewDatastoreAddrManager(ctx context.Context, ds ds.Datastore, ttlInterval time.Duration) *DatastoreAddrManager {
 	mgr := &DatastoreAddrManager{
 		ds:         ds,
-		ttlManager: newTTLManager(context.Background(), ds, ttlInterval),
+		ttlManager: newTTLManager(ctx, ds, ttlInterval),
 		addrSubs:   make(map[peer.ID][]*addrSub),
 	}
 	return mgr
@@ -158,20 +158,16 @@ func newTTLManager(parent context.Context, d ds.Datastore, tick time.Duration) *
 	}
 
 	go func() {
-		stop := false
 		for {
 			select {
 			case <-mgr.ctx.Done():
-				stop = true
+				mgr.ticker.Stop()
+				mgr.done <- struct{}{}
+				return
 			case <-mgr.ticker.C:
 				mgr.tick()
 			}
-
-			if stop {
-				break
-			}
 		}
-		mgr.done <- struct{}{}
 	}()
 
 	return mgr

addr_manager_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/ipfs/go-ds-badger"
 	"github.com/libp2p/go-libp2p-peer"
 	ma "github.com/multiformats/go-multiaddr"
+	"context"
 )
 
 func IDS(t *testing.T, ids string) peer.ID {
@@ -77,7 +78,7 @@ func setupDatastoreAddrManager(t *testing.T) (*DatastoreAddrManager, func()) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	mgr := NewDatastoreAddrManager(ds, 100*time.Microsecond)
+	mgr := NewDatastoreAddrManager(context.Background(), ds, 100*time.Microsecond)
 	closer := func() {
 		mgr.Stop()
 		ds.Close()