vine: Fix context cancellation leaks.

Change-Id: Id4290ad32e5263d0d5b9e136b936745ebfb0a268
diff --git a/runtime/protocols/vine/vine.go b/runtime/protocols/vine/vine.go
index c694d56..a69b195 100644
--- a/runtime/protocols/vine/vine.go
+++ b/runtime/protocols/vine/vine.go
@@ -40,18 +40,25 @@
 // The ctx returned from Init:
 // (1) has localTag as the default localTag for dialers and acceptors.
 // (2) has all addresses in the listenspec altered to listen on the vine protocol.
-func Init(ctx *context.T, name string, auth security.Authorizer, localTag string, discoveryTTL time.Duration) (*context.T, error) {
+func Init(ctx *context.T, name string, auth security.Authorizer, localTag string, discoveryTTL time.Duration) (*context.T, func(), error) {
 	protocol, _ := flow.RegisteredProtocol("vine")
 	v := protocol.(*vine)
-	_, _, err := v23.WithNewServer(ctx, name, VineServer(v), auth)
+	ctx, cancel := context.WithCancel(ctx)
+	_, server, err := v23.WithNewServer(ctx, name, VineServer(v), auth)
+	serverShutdown := func() {
+		cancel()
+		<-server.Closed()
+	}
 	// Nodes are not discoverable until the test controller sets nodes as discoverable.
 	plugin, err := vineplugin.NewWithTTL(ctx, discoveryServerName(localTag), v.discPeers, discoveryTTL)
 	if err != nil {
-		return nil, err
+		serverShutdown()
+		return nil, func() {}, err
 	}
 	df, err := discovery.NewFactory(ctx, plugin)
 	if err != nil {
-		return nil, err
+		serverShutdown()
+		return nil, func() {}, err
 	}
 	factory.InjectFactory(df)
 	lspec := v23.GetListenSpec(ctx).Copy()
@@ -61,7 +68,11 @@
 	}
 	ctx = v23.WithListenSpec(ctx, lspec)
 	ctx = WithLocalTag(ctx, localTag)
-	return ctx, err
+	shutdown := func() {
+		df.Shutdown()
+		serverShutdown()
+	}
+	return ctx, shutdown, nil
 }
 
 // WithLocalTag returns a ctx that will have localTag as the default localTag for
diff --git a/runtime/protocols/vine/vine_test.go b/runtime/protocols/vine/vine_test.go
index 3d2f6b3..46a01f0 100644
--- a/runtime/protocols/vine/vine_test.go
+++ b/runtime/protocols/vine/vine_test.go
@@ -24,10 +24,11 @@
 	ctx, shutdown := test.V23InitWithMounttable()
 	defer shutdown()
 
-	ctx, err := vine.Init(ctx, "vineserver", security.AllowEveryone(), "client", 0)
+	ctx, shutdown, err := vine.Init(ctx, "vineserver", security.AllowEveryone(), "client", 0)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer shutdown()
 	// Create reachable and unreachable server, ensure they have corresponding tags set.
 	ctx, cancel := context.WithCancel(ctx)
 	rctx := vine.WithLocalTag(ctx, "reachable")
@@ -97,10 +98,11 @@
 	ctx, shutdown := test.V23InitWithMounttable()
 	defer shutdown()
 
-	ctx, err := vine.Init(ctx, "vineserver", security.AllowEveryone(), "client", 0)
+	ctx, shutdown, err := vine.Init(ctx, "vineserver", security.AllowEveryone(), "client", 0)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer shutdown()
 	denyCtx := vine.WithLocalTag(ctx, "denyClient")
 	if denyCtx, _, err = v23.WithNewClient(denyCtx); err != nil {
 		t.Fatal(err)
@@ -168,14 +170,16 @@
 	ctx, shutdown := test.V23InitWithMounttable()
 	defer shutdown()
 
-	actx, err := vine.Init(ctx, "advertiser", security.AllowEveryone(), "advertiser", 20*time.Millisecond)
+	actx, shutdown, err := vine.Init(ctx, "advertiser", security.AllowEveryone(), "advertiser", 20*time.Millisecond)
 	if err != nil {
 		t.Fatal(err)
 	}
-	dctx, err := vine.Init(ctx, "scanner", security.AllowEveryone(), "scanner", 20*time.Millisecond)
+	defer shutdown()
+	dctx, shutdown, err := vine.Init(ctx, "scanner", security.AllowEveryone(), "scanner", 20*time.Millisecond)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer shutdown()
 	a, err := v23.NewDiscovery(actx)
 	if err != nil {
 		t.Fatal(err)
diff --git a/services/syncbase/longevity_tests/control/control.go b/services/syncbase/longevity_tests/control/control.go
index 083732a..453c669 100644
--- a/services/syncbase/longevity_tests/control/control.go
+++ b/services/syncbase/longevity_tests/control/control.go
@@ -95,6 +95,9 @@
 
 	// Universe model that the controller is currently simulating.
 	universe *model.Universe
+
+	// vineShutdown shutdown the vine server and vine plugin.
+	vineShutdown func()
 }
 
 // Opts is used to configure the Controller.
@@ -150,20 +153,32 @@
 		return err
 	}
 
+	// Configure context for controller.
+	ctx, err := c.setContextBlessings(ctx, "controller")
+	if err != nil {
+		return err
+	}
+
+	// Set proper namespace root.
+	ctx, _, err = v23.WithNewNamespace(ctx, c.namespaceRoot)
+	if err != nil {
+		return err
+	}
+
+	// Modify ctx to use vine protocol.
+	ctx, c.vineShutdown, err = vine.Init(ctx, c.vineName, security.AllowEveryone(), c.vineName, 0)
+	if err != nil {
+		return err
+	}
+	c.ctx = ctx
+
 	// Configure context for checker.
-	if checkerCtx, err := c.configureContext(ctx, "checker"); err != nil {
+	if checkerCtx, err := c.setContextBlessings(ctx, "checker"); err != nil {
 		return err
 	} else {
 		c.checkerCtx = checkerCtx
 	}
 
-	// Configure context for controller.
-	if ctx, err := c.configureContext(ctx, "controller"); err != nil {
-		return err
-	} else {
-		c.ctx = ctx
-	}
-
 	return nil
 }
 
@@ -193,6 +208,7 @@
 	}
 	c.instancesMu.Unlock()
 
+	c.vineShutdown()
 	c.sh.Cleanup()
 	return retErr
 }
@@ -437,10 +453,9 @@
 	return nil
 }
 
-// configureContext returns a new context based off the given one, which is
-// blessed by the controller's identity provider, configured with the
-// controller's namespace root, and associated with a vine server.
-func (c *Controller) configureContext(ctx *context.T, blessingName string) (*context.T, error) {
+// setContextBlessings returns a new context based off the given one, with a new
+// principal that is blessed by the controller's identity provider.
+func (c *Controller) setContextBlessings(ctx *context.T, blessingName string) (*context.T, error) {
 	// Create a principal blessed by the identity provider.
 	p := testutil.NewPrincipal()
 	if err := c.identityProvider.Bless(p, blessingName); err != nil {
@@ -450,16 +465,6 @@
 	if err != nil {
 		return nil, err
 	}
-	// Set proper namespace root.
-	newCtx, _, err = v23.WithNewNamespace(newCtx, c.namespaceRoot)
-	if err != nil {
-		return nil, err
-	}
-	// Modify ctx to use vine protocol.
-	newCtx, err = vine.Init(newCtx, c.vineName, security.AllowEveryone(), c.vineName, 0)
-	if err != nil {
-		return nil, nil
-	}
 
 	return newCtx, nil
 }
diff --git a/services/syncbase/longevity_tests/control/control_internal_test.go b/services/syncbase/longevity_tests/control/control_internal_test.go
index 0c2e4d6..123ab6f 100644
--- a/services/syncbase/longevity_tests/control/control_internal_test.go
+++ b/services/syncbase/longevity_tests/control/control_internal_test.go
@@ -49,7 +49,7 @@
 	clientRegistry = make(map[string]ClientGenerator)
 }
 
-// InternalConfigureContext exposes controller.configureContext to tests.
-func (c *Controller) InternalConfigureContext(ctx *context.T, blessingName string) (*context.T, error) {
-	return c.configureContext(ctx, blessingName)
+// InternalSetContextBlessings exposes controller.configureContext to tests.
+func (c *Controller) InternalSetContextBlessings(ctx *context.T, blessingName string) (*context.T, error) {
+	return c.setContextBlessings(ctx, blessingName)
 }
diff --git a/services/syncbase/longevity_tests/control/control_test.go b/services/syncbase/longevity_tests/control/control_test.go
index c32c3f4..9a10020 100644
--- a/services/syncbase/longevity_tests/control/control_test.go
+++ b/services/syncbase/longevity_tests/control/control_test.go
@@ -568,7 +568,7 @@
 
 	// Check that Bob gets ErrNoAccess when writing to the collection on his
 	// own device because he does not have write permissions.
-	bobCtx, err := c.InternalConfigureContext(c.InternalCtx(), "u"+security.ChainSeparator+userBob.Name)
+	bobCtx, err := c.InternalSetContextBlessings(c.InternalCtx(), "u"+security.ChainSeparator+userBob.Name)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/services/syncbase/longevity_tests/syncbased_vine/main.go b/services/syncbase/longevity_tests/syncbased_vine/main.go
index 557f7eb..9e3a25a 100644
--- a/services/syncbase/longevity_tests/syncbased_vine/main.go
+++ b/services/syncbase/longevity_tests/syncbased_vine/main.go
@@ -48,10 +48,11 @@
 	ctx = v23.WithListenSpec(ctx, rpc.ListenSpec{
 		Addrs: rpc.ListenAddrs{{"tcp", "127.0.0.1:0"}},
 	})
-	ctx, err := vine.Init(ctx, vineServerName, security.AllowEveryone(), vineTag, 0)
+	ctx, shutdown, err := vine.Init(ctx, vineServerName, security.AllowEveryone(), vineTag, 0)
 	if err != nil {
 		panic(err)
 	}
+	defer shutdown()
 
 	// Syncbase now uses a modified rpc.ListenSpec which has been set to use
 	// the VINE protocol.