fix(device-plugin): return error on failed Send in ListAndWatch

pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go

@@ -454,7 +454,11 @@ func (plugin *NvidiaDevicePlugin) GetDevicePluginOptions(context.Context, *kubel
 
 // ListAndWatch lists devices and update that list according to the health status
 func (plugin *NvidiaDevicePlugin) ListAndWatch(e *kubeletdevicepluginv1beta1.Empty, s kubeletdevicepluginv1beta1.DevicePlugin_ListAndWatchServer) error {
-	s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()})
+	err := s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()})
+	if err != nil {
+		klog.Errorf("Failed to send ListAndWatch response: %v", err)
+		return err
+	}
 
 	for {
 		select {
@@ -464,7 +468,10 @@ func (plugin *NvidiaDevicePlugin) ListAndWatch(e *kubeletdevicepluginv1beta1.Emp
 			// FIXME: there is no way to recover from the Unhealthy state.
 			d.Health = kubeletdevicepluginv1beta1.Unhealthy
 			klog.Infof("'%s' device marked unhealthy: %s", plugin.rm.Resource(), d.ID)
-			s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()})
+			if err := s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()}); err != nil {
+				klog.Errorf("Failed to send health-update ListAndWatch response: %v", err)
+				return nil
+			}
 		}
 	}
 }

pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go

@@ -48,6 +48,7 @@ import (
 	"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
 	"github.com/Project-HAMi/HAMi/pkg/util"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
@@ -1308,3 +1309,66 @@ func TestMigCurrentConfigsNeverNil(t *testing.T) {
 		require.False(t, cfg.MigEnabled)
 	}
 }
+
+type mockListAndWatchServer struct {
+	grpc.ServerStream
+	sendErrs []error
+	sent     []*kubeletdevicepluginv1beta1.ListAndWatchResponse
+}
+
+func (m *mockListAndWatchServer) Send(resp *kubeletdevicepluginv1beta1.ListAndWatchResponse) error {
+	m.sent = append(m.sent, resp)
+	if len(m.sendErrs) > 0 {
+		err := m.sendErrs[0]
+		m.sendErrs = m.sendErrs[1:]
+		return err
+	}
+	return nil
+}
+
+func TestListAndWatch_SendError(t *testing.T) {
+	mockRM := &rm.ResourceManagerMock{
+		DevicesFunc: func() rm.Devices {
+			return rm.Devices{}
+		},
+		ResourceFunc: func() v1.ResourceName {
+			return v1.ResourceName("nvidia.com/gpu")
+		},
+	}
+
+	t.Run("InitialSendFails", func(t *testing.T) {
+		expectedErr := fmt.Errorf("initial send failed")
+		server := &mockListAndWatchServer{
+			sendErrs: []error{expectedErr},
+		}
+		plugin := &NvidiaDevicePlugin{
+			rm:              mockRM,
+			stop:            make(chan any),
+			health:          make(chan *rm.Device, 1),
+			schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}},
+		}
+
+		err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server)
+		require.ErrorIs(t, err, expectedErr)
+		require.Len(t, server.sent, 1)
+	})
+
+	t.Run("UpdateSendFails", func(t *testing.T) {
+		expectedErr := fmt.Errorf("update send failed")
+		server := &mockListAndWatchServer{
+			sendErrs: []error{nil, expectedErr},
+		}
+		plugin := &NvidiaDevicePlugin{
+			rm:              mockRM,
+			stop:            make(chan any),
+			health:          make(chan *rm.Device, 1),
+			schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}},
+		}
+
+		plugin.health <- &rm.Device{Device: kubeletdevicepluginv1beta1.Device{ID: "gpu-1"}}
+		err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server)
+		require.NoError(t, err)
+		require.Len(t, server.sent, 2)
+	})
+}
+