From 2d0ebcedd7de3b995cbc8dd93fadcd4e6cfd6eaf Mon Sep 17 00:00:00 2001
From: Cole Brown <bigswim@gmail.com>
Date: Mon, 11 Jun 2018 19:54:17 -0400
Subject: [PATCH] Accept parent ctx in constructor, simplify TTL loop

---
 addr_manager_ds.go   | 14 +++++---------
 addr_manager_test.go |  3 ++-
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/addr_manager_ds.go b/addr_manager_ds.go
index b3b3a19..ba51652 100644
--- a/addr_manager_ds.go
+++ b/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
diff --git a/addr_manager_test.go b/addr_manager_test.go
index a9519f5..d3eb0c0 100644
--- a/addr_manager_test.go
+++ b/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()