use the new grpc controller

grpc_client.go

@@ -55,12 +55,15 @@ func newGRPCClient(doneCtx context.Context, c *Client) (*GRPCClient, error) {
 	go broker.Run()
 	go brokerGRPCClient.StartStream()
 
-	return &GRPCClient{
-		Conn:    conn,
-		Plugins: c.config.Plugins,
-		doneCtx: doneCtx,
-		broker:  broker,
-	}, nil
+	cl := &GRPCClient{
+		Conn:       conn,
+		Plugins:    c.config.Plugins,
+		doneCtx:    doneCtx,
+		broker:     broker,
+		controller: NewGRPCControllerClient(conn),
+	}
+
+	return cl, nil
 }
 
 // GRPCClient connects to a GRPCServer over gRPC to dispense plugin types.
@@ -70,11 +73,14 @@ type GRPCClient struct {
 
 	doneCtx context.Context
 	broker  *GRPCBroker
+
+	controller GRPCControllerClient
 }
 
 // ClientProtocol impl.
 func (c *GRPCClient) Close() error {
 	c.broker.Close()
+	c.controller.Shutdown(c.doneCtx, &Empty{})
 	return c.Conn.Close()
 }
 

grpc_server.go

@@ -76,6 +76,12 @@ func (s *GRPCServer) Init() error {
 	s.broker = newGRPCBroker(brokerServer, s.TLS)
 	go s.broker.Run()
 
+	// Register the controller
+	controllerServer := &grpcControllerServer{
+		server: s,
+	}
+	RegisterGRPCControllerServer(s.server, controllerServer)
+
 	// Register all our plugins onto the gRPC server.
 	for k, raw := range s.Plugins {
 		p, ok := raw.(GRPCPlugin)

testing.go

@@ -140,6 +140,7 @@ func TestPluginGRPCConn(t testing.T, ps map[string]Plugin) (*GRPCClient, *GRPCSe
 	// Start up the server
 	server := &GRPCServer{
 		Plugins: ps,
+		DoneCh:  make(chan struct{}),
 		Server:  DefaultGRPCServer,
 		Stdout:  new(bytes.Buffer),
 		Stderr:  new(bytes.Buffer),
@@ -165,10 +166,11 @@ func TestPluginGRPCConn(t testing.T, ps map[string]Plugin) (*GRPCClient, *GRPCSe
 
 	// Create the client
 	client := &GRPCClient{
-		Conn:    conn,
-		Plugins: ps,
-		broker:  broker,
-		doneCtx: context.Background(),
+		Conn:       conn,
+		Plugins:    ps,
+		broker:     broker,
+		doneCtx:    context.Background(),
+		controller: NewGRPCControllerClient(conn),
 	}
 
 	return client, server