From 65bd8dc8b032ac2697e738e1d73567f71a1a3906 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeff@hashicorp.com>
Date: Wed, 7 Mar 2018 09:09:37 -0500
Subject: [PATCH] Make grpc plugin client use an atomic server value to fix a
 data race. (#4089)

Also add some coordination to ensure we don't try to clean up the grpc
server before it's created/started
---
 logical/plugin/backend.go             | 12 ++++++++++--
 logical/plugin/grpc_backend_client.go | 23 +++++++++++++++++++----
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/logical/plugin/backend.go b/logical/plugin/backend.go
index b79798fb5d..713f5d3c39 100644
--- a/logical/plugin/backend.go
+++ b/logical/plugin/backend.go
@@ -3,6 +3,7 @@ package plugin
 import (
 	"context"
 	"net/rpc"
+	"sync/atomic"
 
 	"google.golang.org/grpc"
 
@@ -42,10 +43,17 @@ func (b BackendPlugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) err
 }
 
 func (p *BackendPlugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
-	return &backendGRPCPluginClient{
+	ret := &backendGRPCPluginClient{
 		client:     pb.NewBackendClient(c),
 		clientConn: c,
 		broker:     broker,
+		cleanupCh:  make(chan struct{}),
 		doneCtx:    ctx,
-	}, nil
+	}
+
+	// Create the value and set the type
+	ret.server = new(atomic.Value)
+	ret.server.Store((*grpc.Server)(nil))
+
+	return ret, nil
 }
diff --git a/logical/plugin/grpc_backend_client.go b/logical/plugin/grpc_backend_client.go
index 7df90837d1..c980f795be 100644
--- a/logical/plugin/grpc_backend_client.go
+++ b/logical/plugin/grpc_backend_client.go
@@ -3,6 +3,7 @@ package plugin
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 
 	"google.golang.org/grpc"
 
@@ -28,8 +29,13 @@ type backendGRPCPluginClient struct {
 	system logical.SystemView
 	logger log.Logger
 
+	// This is used to signal to the Cleanup function that it can proceed
+	// because we have a defined server
+	cleanupCh chan struct{}
+
 	// server is the grpc server used for serving storage and sysview requests.
-	server *grpc.Server
+	server *atomic.Value
+
 	// clientConn is the underlying grpc connection to the server, we store it
 	// so it can be cleaned up.
 	clientConn *grpc.ClientConn
@@ -139,8 +145,16 @@ func (b *backendGRPCPluginClient) Cleanup(ctx context.Context) {
 	defer cancel()
 
 	b.client.Cleanup(ctx, &pb.Empty{})
-	if b.server != nil {
-		b.server.GracefulStop()
+
+	// This will block until Setup has run the function to create a new server
+	// in b.server. If we stop here before it has a chance to actually start
+	// listening, when it starts listening it will immediatley error out and
+	// exit, which is fine. Overall this ensures that we do not miss stopping
+	// the server if it ends up being created after Cleanup is called.
+	<-b.cleanupCh
+	server := b.server.Load()
+	if server != nil {
+		server.(*grpc.Server).GracefulStop()
 	}
 	b.clientConn.Close()
 }
@@ -184,7 +198,8 @@ func (b *backendGRPCPluginClient) Setup(ctx context.Context, config *logical.Bac
 		s := grpc.NewServer(opts...)
 		pb.RegisterSystemViewServer(s, sysView)
 		pb.RegisterStorageServer(s, storage)
-		b.server = s
+		b.server.Store(s)
+		close(b.cleanupCh)
 		return s
 	}
 	brokerID := b.broker.NextId()