v.io/x/ref/lib/discovery: fix some leaked context.CancelFuncs

Clients should call context.CancelFunc functions to avoid memory leaks.

Change-Id: I7cbdb2275654f84af5d6a02079ca0627f5fb715e
diff --git a/lib/discovery/advertise_server.go b/lib/discovery/advertise_server.go
index 75e70b6..9e9a73a 100644
--- a/lib/discovery/advertise_server.go
+++ b/lib/discovery/advertise_server.go
@@ -37,6 +37,7 @@
 	curAddrs := sortedNames(status.Endpoints)
 	stop, err := advertiseServer(ctx, d, ad, curAddrs, suffix, visibility)
 	if err != nil {
+		stop()
 		return nil, err
 	}
 
@@ -49,6 +50,7 @@
 			case <-status.Dirty:
 				status = server.Status()
 				if status.State != rpc.ServerActive {
+					stop()
 					return
 				}
 				newAddrs := sortedNames(status.Endpoints)
@@ -59,6 +61,7 @@
 				stop() // Stop the previous advertisement.
 				stop, err = advertiseServer(ctx, d, ad, newAddrs, suffix, visibility)
 				if err != nil {
+					stop()
 					ctx.Error(err)
 					return
 				}
diff --git a/lib/discovery/directory.go b/lib/discovery/directory.go
index 399e36e..02e6ca8 100644
--- a/lib/discovery/directory.go
+++ b/lib/discovery/directory.go
@@ -102,6 +102,7 @@
 	s := &dirServer{d: d, adMap: make(map[discovery.AdId]*AdInfo), cancel: cancel}
 	ctx, server, err := v23.WithNewServer(ctx, "", DirectoryServer(s), security.AllowEveryone(), opts...)
 	if err != nil {
+		cancel()
 		return nil, err
 	}
 
diff --git a/lib/discovery/discovery_test.go b/lib/discovery/discovery_test.go
index fb32e5a..a6e2dca 100644
--- a/lib/discovery/discovery_test.go
+++ b/lib/discovery/discovery_test.go
@@ -196,10 +196,18 @@
 	}
 
 	d, _ := df.New(ctx)
-	if _, err := testutil.Advertise(ctx, d, nil, &ad); err != nil {
+	stop, err := testutil.Advertise(ctx, d, nil, &ad)
+	if stop != nil {
+		defer stop()
+	}
+	if err != nil {
 		t.Fatal(err)
 	}
-	if _, err := testutil.Advertise(ctx, d, nil, &ad); err == nil {
+	stop, err = testutil.Advertise(ctx, d, nil, &ad)
+	if stop != nil {
+		defer stop()
+	}
+	if err == nil {
 		t.Error("expect an error; but got none")
 	}
 }
@@ -527,7 +535,11 @@
 	}
 
 	d1, _ := df.New(ctx)
-	if _, err := testutil.Advertise(ctx, d1, nil, &ad); err != nil {
+	stop, err := testutil.Advertise(ctx, d1, nil, &ad)
+	if stop != nil {
+		defer stop()
+	}
+	if err != nil {
 		t.Error(err)
 	}
 	d2, _ := df.New(ctx)
@@ -541,7 +553,11 @@
 
 	// Make sure advertise and scan do not work after closed.
 	ad.Id = discovery.AdId{} // To avoid dup error.
-	if _, err := testutil.Advertise(ctx, d1, nil, &ad); err == nil {
+	stop, err = testutil.Advertise(ctx, d1, nil, &ad)
+	if stop != nil {
+		defer stop()
+	}
+	if err == nil {
 		t.Error("expect an error; but got none")
 	}
 	if err := testutil.ScanAndMatch(ctx, d2, ``, ad); err == nil {
diff --git a/lib/discovery/plugins/ble/ble_test.go b/lib/discovery/plugins/ble/ble_test.go
index 2bbcea7..e8dff79 100644
--- a/lib/discovery/plugins/ble/ble_test.go
+++ b/lib/discovery/plugins/ble/ble_test.go
@@ -515,6 +515,7 @@
 	adinfo.Hash = idiscovery.AdHash{7, 8, 9}
 	adinfo.TimestampNs = 1002
 	stop, err = testutil.Advertise(ctx, p1, &adinfo)
+	defer stop()
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/lib/discovery/plugins/testutil/util.go b/lib/discovery/plugins/testutil/util.go
index 5178542..0433729 100644
--- a/lib/discovery/plugins/testutil/util.go
+++ b/lib/discovery/plugins/testutil/util.go
@@ -22,6 +22,7 @@
 	for _, adinfo := range adinfos {
 		wg.Add(1)
 		if err := p.Advertise(ctx, adinfo, wg.Done); err != nil {
+			cancel()
 			return nil, fmt.Errorf("Advertise failed: %v", err)
 		}
 	}
diff --git a/lib/discovery/plugins/vine/vine.go b/lib/discovery/plugins/vine/vine.go
index 4f71bc1..3795321 100644
--- a/lib/discovery/plugins/vine/vine.go
+++ b/lib/discovery/plugins/vine/vine.go
@@ -60,7 +60,8 @@
 
 		// We need a context that is detached from the deadlines and cancellation of ctx
 		// since we have to unpublish the advertisement from peers after ctx is canceled.
-		rctx, _ := context.WithRootCancel(ctx)
+		rctx, cancel := context.WithRootCancel(ctx)
+		defer cancel()
 		defer p.unpublish(rctx, adinfo.Ad.Id)
 
 		for {
diff --git a/lib/discovery/testutil/util.go b/lib/discovery/testutil/util.go
index 738e495..deb5676 100644
--- a/lib/discovery/testutil/util.go
+++ b/lib/discovery/testutil/util.go
@@ -79,6 +79,7 @@
 	ctx, cancel := context.WithCancel(ctx)
 	scanCh, err := d.Scan(ctx, query)
 	if err != nil {
+		cancel()
 		return nil, nil, fmt.Errorf("Scan failed: %v", err)
 	}
 	stop := func() {