Fix kubelet memory leak when device plugin is registered

pkg/kubelet/cm/devicemanager/plugin/v1beta1/client.go

@@ -40,10 +40,14 @@ type DevicePlugin interface {
 // Client interface provides methods for establishing/closing gRPC connection and running the device plugin gRPC client.
 type Client interface {
 	Connect() error
+	SocketPath() string
 	Run()
 	Disconnect() error
 }
 
+var _ Client = &client{}
+var _ DevicePlugin = &client{}
+
 type client struct {
 	mutex    sync.Mutex
 	resource string
@@ -63,16 +67,32 @@ func NewPluginClient(r string, socketPath string, h ClientHandler) Client {
 }
 
 // Connect is for establishing a gRPC connection between device manager and device plugin.
-func (c *client) Connect() error {
+func (c *client) Connect() (rerr error) {
 	client, conn, err := dial(c.socket)
 	if err != nil {
 		klog.ErrorS(err, "Unable to connect to device plugin client with socket path", "path", c.socket)
 		return err
 	}
+
 	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
 	c.grpc = conn
 	c.client = client
-	c.mutex.Unlock()
+
+	defer func() {
+		if rerr == nil {
+			return
+		}
+		klog.ErrorS(err, "Failed to connect to device plugin", "resource", c.resource)
+		if err := conn.Close(); err != nil {
+			klog.V(2).ErrorS(err, "Failed to close gRPC connection", "resource", c.resource)
+			return
+		}
+		c.grpc = nil
+		c.client = nil
+	}()
+
 	return c.handler.PluginConnected(c.resource, c)
 }
 

pkg/kubelet/cm/devicemanager/plugin/v1beta1/handler.go

@@ -66,14 +66,21 @@ func (s *server) ValidatePlugin(pluginName string, endpoint string, versions []s
 }
 
 func (s *server) connectClient(name string, socketPath string) error {
-	c := NewPluginClient(name, socketPath, s.chandler)
+	c := s.getClient(name)
+	if c != nil {
+		if c.SocketPath() != socketPath {
+			return fmt.Errorf("the device plugin %s already registered with a different socket path %s", name, c.SocketPath())
+		}
+		klog.V(2).InfoS("Client already connected with the same socket path", "name", name, "socketPath", socketPath)
+		return nil
+	}
 
-	s.registerClient(name, c)
+	c = NewPluginClient(name, socketPath, s.chandler)
 	if err := c.Connect(); err != nil {
-		s.deregisterClient(name)
 		klog.ErrorS(err, "Failed to connect to new client", "resource", name)
 		return err
 	}
+	s.registerClient(name, c)
 
 	go func() {
 		s.runClient(name, c)
@@ -83,8 +90,11 @@ func (s *server) connectClient(name string, socketPath string) error {
 }
 
 func (s *server) disconnectClient(name string, c Client) error {
+	if err := c.Disconnect(); err != nil {
+		return err
+	}
 	s.deregisterClient(name)
-	return c.Disconnect()
+	return nil
 }
 
 func (s *server) registerClient(name string, c Client) {

pkg/kubelet/cm/devicemanager/plugin/v1beta1/server.go

@@ -52,7 +52,8 @@ type server struct {
 	grpc       *grpc.Server
 	rhandler   RegistrationHandler
 	chandler   ClientHandler
-	clients    map[string]Client
+	// clients is a map of plugin name to connected client
+	clients map[string]Client
 }
 
 // NewServer returns an initialized device plugin registration server.