1 files changed, 14 insertions, 10 deletions

diff --git a/vendor/github.com/hashicorp/go-plugin/grpc_client.go b/vendor/github.com/hashicorp/go-plugin/grpc_client.go
index 44294d0..294518e 100644
--- a/vendor/github.com/hashicorp/go-plugin/grpc_client.go
+++ b/vendor/github.com/hashicorp/go-plugin/grpc_client.go
@@ -6,6 +6,7 @@ import (
         "net"
         "time"
 
+        "github.com/hashicorp/go-plugin/internal/plugin"
         "golang.org/x/net/context"
         "google.golang.org/grpc"
         "google.golang.org/grpc/credentials"
@@ -16,12 +17,9 @@ func dialGRPCConn(tls *tls.Config, dialer func(string, time.Duration) (net.Conn,
         // Build dialing options.
         opts := make([]grpc.DialOption, 0, 5)
 
-        // We use a custom dialer so that we can connect over unix domain sockets
+        // We use a custom dialer so that we can connect over unix domain sockets.
         opts = append(opts, grpc.WithDialer(dialer))
 
-        // go-plugin expects to block the connection
-        opts = append(opts, grpc.WithBlock())
-
         // Fail right away
         opts = append(opts, grpc.FailOnNonTempDialError(true))
 
@@ -58,12 +56,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: plugin.NewGRPCControllerClient(conn),
+        }
+
+        return cl, nil
 }
 
 // GRPCClient connects to a GRPCServer over gRPC to dispense plugin types.
@@ -73,11 +74,14 @@ type GRPCClient struct {
 
         doneCtx context.Context
         broker  *GRPCBroker
+
+        controller plugin.GRPCControllerClient
 }
 
 // ClientProtocol impl.
 func (c *GRPCClient) Close() error {
         c.broker.Close()
+        c.controller.Shutdown(c.doneCtx, &plugin.Empty{})
         return c.Conn.Close()
 }