From ad6311471ea4f00d2b34a45088e3af7a040733cd Mon Sep 17 00:00:00 2001
From: Etran-H <34226782+Etran-H@users.noreply.github.com>
Date: Thu, 17 Nov 2022 09:56:18 +0800
Subject: [PATCH 1/2] use GracefulStop

---
 grpc_controller.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/grpc_controller.go b/grpc_controller.go
index 1a8a8e70..3ae49e54 100644
--- a/grpc_controller.go
+++ b/grpc_controller.go
@@ -17,7 +17,6 @@ type grpcControllerServer struct {
 func (s *grpcControllerServer) Shutdown(ctx context.Context, _ *plugin.Empty) (*plugin.Empty, error) {
 	resp := &plugin.Empty{}
 
-	// TODO: figure out why GracefullStop doesn't work.
-	s.server.Stop()
+	go s.server.GracefulStop()
 	return resp, nil
 }

From 716a28023cbb6959a6f5278d18f211c64e29906d Mon Sep 17 00:00:00 2001
From: Etran-H <34226782+Etran-H@users.noreply.github.com>
Date: Thu, 17 Nov 2022 10:19:19 +0800
Subject: [PATCH 2/2] Add GracefulStopTimeout to ClientConfig

---
 client.go | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/client.go b/client.go
index 2e86f621..c0571f92 100644
--- a/client.go
+++ b/client.go
@@ -209,6 +209,9 @@ type ClientConfig struct {
 	// to create gRPC connections. This only affects plugins using the gRPC
 	// protocol.
 	GRPCDialOptions []grpc.DialOption
+	
+	// GracefulStopTimeout is the timeout to wait for the plugin stop.
+	GracefulStopTimeout time.Duration
 }
 
 // ReattachConfig is used to configure a client to reattach to an
@@ -326,6 +329,10 @@ func NewClient(config *ClientConfig) (c *Client) {
 	if config.AllowedProtocols == nil {
 		config.AllowedProtocols = []Protocol{ProtocolNetRPC}
 	}
+	
+	if config.GracefulStopTimeout == 0 {
+		config.GracefulStopTimeout = 2 * time.Second
+	}
 
 	if config.Logger == nil {
 		config.Logger = hclog.New(&hclog.LoggerOptions{
@@ -460,7 +467,7 @@ func (c *Client) Kill() {
 		case <-c.doneCtx.Done():
 			c.logger.Debug("plugin exited")
 			return
-		case <-time.After(2 * time.Second):
+		case <-time.After(c.config.GracefulStopTimeout):
 		}
 	}