Fix ListAndWatch race

cmd/udev-manager/e2e_test.go

@@ -784,6 +784,36 @@ partitions:
 		})
 	})
 
+	Describe("Kubelet drops ListAndWatch stream", func() {
+		It("re-registers when the stream breaks without a kubelet restart", func() {
+			dev := makePartitionDevice("/sys/block/nvme0n1/nvme0n1p1", "/dev/nvme0n1p1", "nvme_disk01")
+			discovery.AddDevice(dev)
+
+			config := mustParseYAML(`
+domain: ydb.tech
+partitions:
+  - matcher: "nvme_(.*)"
+`)
+			startTestApp(ctx, wg, discovery, config, tmpDir, kubeSock)
+			waitForRegistrations(kubelet, 1)
+			sockets := waitForSockets(tmpDir)
+
+			By("opening a ListAndWatch stream, as kubelet would")
+			client, conn := dialPlugin(sockets[0])
+			stream, err := client.ListAndWatch(ctx, &pluginapi.Empty{})
+			Expect(err).NotTo(HaveOccurred())
+			recvWithTimeout(stream, 5*time.Second)
+
+			By("dropping the connection without recreating the kubelet socket")
+			conn.Close()
+
+			By("the plugin re-registers on its own")
+			waitForRegistrations(kubelet, 2)
+			reg := kubelet.Registrations()[1]
+			Expect(reg.ResourceName).To(Equal("ydb.tech/part-disk01"))
+		})
+	})
+
 	Describe("Shutdown", func() {
 		It("terminates ListAndWatch stream when context is cancelled", func() {
 			dev := makePartitionDevice("/sys/block/nvme0n1/nvme0n1p1", "/dev/nvme0n1p1", "nvme_disk01")

internal/plugin/plugin.go

@@ -22,17 +22,25 @@ import (
 type plugin struct {
 	resource  Resource
 	pluginDir string
+	ctx       context.Context
 	cancel    context.CancelFunc
 	stopped   chan struct{} // closed after gRPC server is fully stopped
+	// onStreamBroken is invoked when a ListAndWatch stream ends with an
+	// error while the plugin itself is still running, i.e. the kubelet
+	// dropped the connection. The kubelet never re-opens the stream on its
+	// own — it waits for the plugin to register again.
+	onStreamBroken func(*plugin)
 }
 
-func newPlugin(resource Resource, ctx context.Context, wg *sync.WaitGroup, pluginDir string) (*plugin, error) {
+func newPlugin(resource Resource, ctx context.Context, wg *sync.WaitGroup, pluginDir string, onStreamBroken func(*plugin)) (*plugin, error) {
 	ctx, cancel := context.WithCancel(ctx)
 	plugin := &plugin{
-		resource:  resource,
-		pluginDir: pluginDir,
-		cancel:    cancel,
-		stopped:   make(chan struct{}),
+		resource:       resource,
+		pluginDir:      pluginDir,
+		ctx:            ctx,
+		cancel:         cancel,
+		stopped:        make(chan struct{}),
+		onStreamBroken: onStreamBroken,
 	}
 
 	socketPath := pluginDir + plugin.socketPath()
@@ -96,7 +104,12 @@ func (p *plugin) PreStartContainer(context.Context, *pluginapi.PreStartContainer
 }
 
 func (p *plugin) ListAndWatch(empty *pluginapi.Empty, stream pluginapi.DevicePlugin_ListAndWatchServer) (err error) {
-	defer klog.Infof("%q: closing ListAndWatch connection, err = %v", p.resource.Name(), err)
+	defer func() {
+		klog.Infof("%q: closing ListAndWatch connection, err = %v", p.resource.Name(), err)
+		if err != nil && p.ctx.Err() == nil && p.onStreamBroken != nil {
+			p.onStreamBroken(p)
+		}
+	}()
 
 	ctx := stream.Context()
 	instanceCh := p.resource.ListAndWatch(ctx)

internal/plugin/registry.go

@@ -17,6 +17,10 @@ import (
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
 )
 
+// reregisterDelay spaces out re-registration after a broken ListAndWatch
+// stream, avoiding a tight register/drop loop against an unhealthy kubelet.
+const reregisterDelay = time.Second
+
 // Registry is a lifecycle manager for plugins.
 // It is responsible for (re-)registering plugins with the kubelet.
 type Registry struct {
@@ -77,14 +81,41 @@ func (r *Registry) register(plugin *plugin) error {
 	return nil
 }
 
+// reregister re-registers a plugin with the kubelet after its ListAndWatch
+// stream was broken. The kubelet never re-opens a broken stream on its own;
+// it only opens a new one in response to a Register call. Runs asynchronously
+// because it is invoked from the ListAndWatch handler itself.
+func (r *Registry) reregister(p *plugin) {
+	r.wg.Add(1)
+	go func() {
+		defer r.wg.Done()
+		for {
+			select {
+			case <-time.After(reregisterDelay):
+			case <-p.stopped:
+				// The plugin was stopped (e.g. kubelet restart triggered hup);
+				// the replacement plugin registers itself.
+				return
+			case <-r.ctx.Done():
+				return
+			}
+			klog.Warningf("%s: ListAndWatch stream broken by kubelet; re-registering", p.resource.Name())
+			if err := r.register(p); err == nil {
+				return
+			}
+			// register failed; loop and retry after another delay.
+		}
+	}()
+}
+
 // hup registers all plugins with the freshly kubelet.
 // Newely started kubelet removes all socket files, so we need to re-register
 // all plugins. See https://kubernetes.io/docs/concepts/extend-kubernetes/compute-storage-net/device-plugins/#handling-kubelet-restarts
 func (r *Registry) hup() {
 	r.plugins.Range(func(key, p interface{}) bool {
 		old := p.(*plugin)
 		old.stop()
-		newP, err := newPlugin(old.resource, r.ctx, r.wg, r.pluginDir)
+		newP, err := newPlugin(old.resource, r.ctx, r.wg, r.pluginDir, r.reregister)
 		if err != nil {
 			klog.Errorf("failed to create plugin for %s: %v", old.resource.Name(), err)
 			return true
@@ -145,10 +176,23 @@ func NewRegistry(ctx context.Context, wg *sync.WaitGroup, opts ...RegistryOption
 
 		for {
 			select {
-			case event := <-r.watcher.Events:
+			case event, ok := <-r.watcher.Events:
+				if !ok {
+					return
+				}
 				if event.Op&fsnotify.Create != 0 && event.Name == r.kubeletSocket {
 					r.hup()
 				}
+			case err, ok := <-r.watcher.Errors:
+				if !ok {
+					return
+				}
+				// An error here (typically an inotify queue overflow) means
+				// events were lost — possibly the kubelet socket CREATE.
+				// Leaving this channel undrained would wedge the watcher
+				// entirely. Resync by re-registering everything.
+				klog.Errorf("kubelet socket watcher error, re-registering all plugins: %v", err)
+				r.hup()
 			case <-r.ctx.Done():
 				// Parent context is done, exit the goroutine.
 				return
@@ -189,7 +233,7 @@ func (r *Registry) Healthz(resp http.ResponseWriter, req *http.Request) {
 // kubelet. Attempts to register resource with the same name twice will result
 // in an error.
 func (r *Registry) Add(resource Resource) error {
-	plugin, err := newPlugin(resource, r.ctx, r.wg, r.pluginDir)
+	plugin, err := newPlugin(resource, r.ctx, r.wg, r.pluginDir, r.reregister)
 	if err != nil {
 		klog.Errorf("failed to create plugin for resource %q Cause: %v", resource.Name(), err)
 		return err

internal/udev/udev.go

@@ -414,7 +414,13 @@ func (d *udevDiscovery) monitor(wg *sync.WaitGroup) {
 	// Step 3: process buffered and future events.
 	for {
 		select {
-		case dev := <-devChan:
+		case dev, ok := <-devChan:
+			if !ok {
+				// The monitor goroutine died and closed its channels;
+				// the errChan case performs the reconnect.
+				devChan = nil
+				continue
+			}
 			klog.V(5).Infof("Received device event (%s): %s", dev.Action(), dev.Syspath())
 			switch dev.Action() {
 			case ActionAdd, ActionOnline:
@@ -472,7 +478,10 @@ func (d *udevDiscovery) monitor(wg *sync.WaitGroup) {
 				req.Reply(nil)
 				return
 			}
-		case err := <-errChan:
+		case err, ok := <-errChan:
+			if !ok {
+				err = fmt.Errorf("udev monitor channel closed")
+			}
 			klog.Errorf("Error from udev monitor, will try to retry connecting to udev: %v", err)
 		retry:
 			mon = d.udev.NewMonitorFromNetlink("udev")