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()