Updated the device manager pluginwatcher handler

pkg/kubelet/cm/container_manager.go

@@ -95,7 +95,11 @@ type ContainerManager interface {
 
 	// GetPodCgroupRoot returns the cgroup which contains all pods.
 	GetPodCgroupRoot() string
-	GetPluginRegistrationHandlerCallback() pluginwatcher.RegisterCallbackFn
+
+	// GetPluginRegistrationHandler returns a plugin registration handler
+	// The pluginwatcher's Handlers allow to have a single module for handling
+	// registration.
+	GetPluginRegistrationHandler() pluginwatcher.PluginHandler
 }
 
 type NodeConfig struct {

pkg/kubelet/cm/container_manager_linux.go

@@ -605,8 +605,8 @@ func (cm *containerManagerImpl) Start(node *v1.Node,
 	return nil
 }
 
-func (cm *containerManagerImpl) GetPluginRegistrationHandlerCallback() pluginwatcher.RegisterCallbackFn {
-	return cm.deviceManager.GetWatcherCallback()
+func (cm *containerManagerImpl) GetPluginRegistrationHandler() pluginwatcher.PluginHandler {
+	return cm.deviceManager.GetWatcherHandler()
 }
 
 // TODO: move the GetResources logic to PodContainerManager.

pkg/kubelet/cm/container_manager_stub.go

@@ -77,10 +77,8 @@ func (cm *containerManagerStub) GetCapacity() v1.ResourceList {
 	return c
 }
 
-func (cm *containerManagerStub) GetPluginRegistrationHandlerCallback() pluginwatcher.RegisterCallbackFn {
-	return func(name string, endpoint string, versions []string, sockPath string) (chan bool, error) {
-		return nil, nil
-	}
+func (cm *containerManagerStub) GetPluginRegistrationHandler() pluginwatcher.PluginHandler {
+	return nil
 }
 
 func (cm *containerManagerStub) GetDevicePluginResourceCapacity() (v1.ResourceList, v1.ResourceList, []string) {

pkg/kubelet/cm/devicemanager/manager_stub.go

@@ -57,9 +57,7 @@ func (h *ManagerStub) GetCapacity() (v1.ResourceList, v1.ResourceList, []string)
 	return nil, nil, []string{}
 }
 
-// GetWatcherCallback returns plugin watcher callback
-func (h *ManagerStub) GetWatcherCallback() pluginwatcher.RegisterCallbackFn {
-	return func(name string, endpoint string, versions []string, sockPath string) (chan bool, error) {
-		return nil, nil
-	}
+// GetWatcherHandler returns plugin watcher interface
+func (h *ManagerStub) GetWatcherHandler() pluginwatcher.PluginHandler {
+	return nil
 }

pkg/kubelet/cm/devicemanager/manager_test.go

@@ -249,9 +249,10 @@ func setupDevicePlugin(t *testing.T, devs []*pluginapi.Device, pluginSocketName
 
 func setupPluginWatcher(pluginSocketName string, m Manager) *pluginwatcher.Watcher {
 	w := pluginwatcher.NewWatcher(filepath.Dir(pluginSocketName))
-	w.AddHandler(watcherapi.DevicePlugin, m.GetWatcherCallback())
+	w.AddHandler(watcherapi.DevicePlugin, m.GetWatcherHandler())
 	w.Start()
-	return &w
+
+	return w
 }
 
 func setup(t *testing.T, devs []*pluginapi.Device, callback monitorCallback, socketName string, pluginSocketName string) (Manager, <-chan interface{}, *Stub) {
@@ -295,7 +296,7 @@ func TestUpdateCapacityAllocatable(t *testing.T) {
 	// Expects capacity for resource1 to be 2.
 	resourceName1 := "domain1.com/resource1"
 	e1 := &endpointImpl{}
-	testManager.endpoints[resourceName1] = e1
+	testManager.endpoints[resourceName1] = endpointInfo{e: e1, opts: nil}
 	callback(resourceName1, devs)
 	capacity, allocatable, removedResources := testManager.GetCapacity()
 	resource1Capacity, ok := capacity[v1.ResourceName(resourceName1)]
@@ -345,7 +346,7 @@ func TestUpdateCapacityAllocatable(t *testing.T) {
 	// Tests adding another resource.
 	resourceName2 := "resource2"
 	e2 := &endpointImpl{}
-	testManager.endpoints[resourceName2] = e2
+	testManager.endpoints[resourceName2] = endpointInfo{e: e2, opts: nil}
 	callback(resourceName2, devs)
 	capacity, allocatable, removedResources = testManager.GetCapacity()
 	as.Equal(2, len(capacity))
@@ -456,7 +457,7 @@ func TestCheckpoint(t *testing.T) {
 	ckm, err := checkpointmanager.NewCheckpointManager(tmpDir)
 	as.Nil(err)
 	testManager := &ManagerImpl{
-		endpoints:         make(map[string]endpoint),
+		endpoints:         make(map[string]endpointInfo),
 		healthyDevices:    make(map[string]sets.String),
 		unhealthyDevices:  make(map[string]sets.String),
 		allocatedDevices:  make(map[string]sets.String),
@@ -577,7 +578,7 @@ func makePod(limits v1.ResourceList) *v1.Pod {
 	}
 }
 
-func getTestManager(tmpDir string, activePods ActivePodsFunc, testRes []TestResource, opts map[string]*pluginapi.DevicePluginOptions) (*ManagerImpl, error) {
+func getTestManager(tmpDir string, activePods ActivePodsFunc, testRes []TestResource) (*ManagerImpl, error) {
 	monitorCallback := func(resourceName string, devices []pluginapi.Device) {}
 	ckm, err := checkpointmanager.NewCheckpointManager(tmpDir)
 	if err != nil {
@@ -589,41 +590,45 @@ func getTestManager(tmpDir string, activePods ActivePodsFunc, testRes []TestReso
 		healthyDevices:    make(map[string]sets.String),
 		unhealthyDevices:  make(map[string]sets.String),
 		allocatedDevices:  make(map[string]sets.String),
-		endpoints:         make(map[string]endpoint),
-		pluginOpts:        opts,
+		endpoints:         make(map[string]endpointInfo),
 		podDevices:        make(podDevices),
 		activePods:        activePods,
 		sourcesReady:      &sourcesReadyStub{},
 		checkpointManager: ckm,
 	}
+
 	for _, res := range testRes {
 		testManager.healthyDevices[res.resourceName] = sets.NewString()
 		for _, dev := range res.devs {
 			testManager.healthyDevices[res.resourceName].Insert(dev)
 		}
 		if res.resourceName == "domain1.com/resource1" {
-			testManager.endpoints[res.resourceName] = &MockEndpoint{
-				allocateFunc: allocateStubFunc(),
+			testManager.endpoints[res.resourceName] = endpointInfo{
+				e:    &MockEndpoint{allocateFunc: allocateStubFunc()},
+				opts: nil,
 			}
 		}
 		if res.resourceName == "domain2.com/resource2" {
-			testManager.endpoints[res.resourceName] = &MockEndpoint{
-				allocateFunc: func(devs []string) (*pluginapi.AllocateResponse, error) {
-					resp := new(pluginapi.ContainerAllocateResponse)
-					resp.Envs = make(map[string]string)
-					for _, dev := range devs {
-						switch dev {
-						case "dev3":
-							resp.Envs["key2"] = "val2"
-
-						case "dev4":
-							resp.Envs["key2"] = "val3"
+			testManager.endpoints[res.resourceName] = endpointInfo{
+				e: &MockEndpoint{
+					allocateFunc: func(devs []string) (*pluginapi.AllocateResponse, error) {
+						resp := new(pluginapi.ContainerAllocateResponse)
+						resp.Envs = make(map[string]string)
+						for _, dev := range devs {
+							switch dev {
+							case "dev3":
+								resp.Envs["key2"] = "val2"
+
+							case "dev4":
+								resp.Envs["key2"] = "val3"
+							}
 						}
-					}
-					resps := new(pluginapi.AllocateResponse)
-					resps.ContainerResponses = append(resps.ContainerResponses, resp)
-					return resps, nil
+						resps := new(pluginapi.AllocateResponse)
+						resps.ContainerResponses = append(resps.ContainerResponses, resp)
+						return resps, nil
+					},
 				},
+				opts: nil,
 			}
 		}
 	}
@@ -669,10 +674,7 @@ func TestPodContainerDeviceAllocation(t *testing.T) {
 	as.Nil(err)
 	defer os.RemoveAll(tmpDir)
 	nodeInfo := getTestNodeInfo(v1.ResourceList{})
-	pluginOpts := make(map[string]*pluginapi.DevicePluginOptions)
-	pluginOpts[res1.resourceName] = nil
-	pluginOpts[res2.resourceName] = nil
-	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, testResources, pluginOpts)
+	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, testResources)
 	as.Nil(err)
 
 	testPods := []*v1.Pod{
@@ -767,10 +769,8 @@ func TestInitContainerDeviceAllocation(t *testing.T) {
 	tmpDir, err := ioutil.TempDir("", "checkpoint")
 	as.Nil(err)
 	defer os.RemoveAll(tmpDir)
-	pluginOpts := make(map[string]*pluginapi.DevicePluginOptions)
-	pluginOpts[res1.resourceName] = nil
-	pluginOpts[res2.resourceName] = nil
-	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, testResources, pluginOpts)
+
+	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, testResources)
 	as.Nil(err)
 
 	podWithPluginResourcesInInitContainers := &v1.Pod{
@@ -904,18 +904,18 @@ func TestDevicePreStartContainer(t *testing.T) {
 	as.Nil(err)
 	defer os.RemoveAll(tmpDir)
 	nodeInfo := getTestNodeInfo(v1.ResourceList{})
-	pluginOpts := make(map[string]*pluginapi.DevicePluginOptions)
-	pluginOpts[res1.resourceName] = &pluginapi.DevicePluginOptions{PreStartRequired: true}
 
-	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, []TestResource{res1}, pluginOpts)
+	testManager, err := getTestManager(tmpDir, podsStub.getActivePods, []TestResource{res1})
 	as.Nil(err)
 
 	ch := make(chan []string, 1)
-	testManager.endpoints[res1.resourceName] = &MockEndpoint{
-		initChan:     ch,
-		allocateFunc: allocateStubFunc(),
+	testManager.endpoints[res1.resourceName] = endpointInfo{
+		e: &MockEndpoint{
+			initChan:     ch,
+			allocateFunc: allocateStubFunc(),
+		},
+		opts: &pluginapi.DevicePluginOptions{PreStartRequired: true},
 	}
-
 	pod := makePod(v1.ResourceList{
 		v1.ResourceName(res1.resourceName): res1.resourceQuantity})
 	activePods := []*v1.Pod{}

pkg/kubelet/cm/devicemanager/types.go

@@ -53,7 +53,7 @@ type Manager interface {
 	// GetCapacity returns the amount of available device plugin resource capacity, resource allocatable
 	// and inactive device plugin resources previously registered on the node.
 	GetCapacity() (v1.ResourceList, v1.ResourceList, []string)
-	GetWatcherCallback() watcher.RegisterCallbackFn
+	GetWatcherHandler() watcher.PluginHandler
 }
 
 // DeviceRunContainerOptions contains the combined container runtime settings to consume its allocated devices.

pkg/kubelet/kubelet.go

@@ -1367,7 +1367,7 @@ func (kl *Kubelet) initializeRuntimeDependentModules() {
 		// Adding Registration Callback function for CSI Driver
 		kl.pluginWatcher.AddHandler("CSIPlugin", pluginwatcher.PluginHandler(csi.PluginHandler))
 		// Adding Registration Callback function for Device Manager
-		kl.pluginWatcher.AddHandler(pluginwatcherapi.DevicePlugin, kl.containerManager.GetPluginRegistrationHandlerCallback())
+		kl.pluginWatcher.AddHandler(pluginwatcherapi.DevicePlugin, kl.containerManager.GetPluginRegistrationHandler())
 		// Start the plugin watcher
 		glog.V(4).Infof("starting watcher")
 		if err := kl.pluginWatcher.Start(); err != nil {