Code change:
"context" "errors" "fmt" <ins> "net" "sync" </ins> "time" <del> v1 "k8s.io/api/core/v1" resourceapi "k8s.io/api/resource/v1alpha3" apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/fields" </del> <ins> "google.golang.org/grpc" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials/insecure" </ins> utilversion "k8s.io/apimachinery/pkg/util/version" <del> "k8s.io/apimachinery/pkg/util/wait" "k8s.io/client-go/kubernetes" </del> "k8s.io/klog/v2" <del> "k8s.io/kubernetes/pkg/kubelet/pluginmanager/cache" </del> <ins> drapb "k8s.io/kubelet/pkg/apis/dra/v1alpha4" </ins> ) <del> // defaultClientCallTimeout is the default amount of time that a DRA driver has // to respond to any of the gRPC calls. kubelet uses this value by passing nil // to RegisterPlugin. Some tests use a different, usually shorter timeout to // speed up testing. // // This is half of the kubelet retry period (according to // https://github.com/kubernetes/kubernetes/commit/0449cef8fd5217d394c5cd331d852bd50983e6b3). const defaultClientCallTimeout = 45 * time.Second // RegistrationHandler is the handler which is fed to the pluginwatcher API. type RegistrationHandler struct { // backgroundCtx is used for all future activities of the handler. // This is necessary because it implements APIs which don't // provide a context. backgroundCtx context.Context kubeClient    kubernetes.Interface getNode       func() (*v1.Node, error) } </del> <ins> // NewDRAPluginClient returns a wrapper around those gRPC methods of a DRA // driver kubelet plugin which need to be called by kubelet. The wrapper // handles gRPC connection management and logging. Connections are reused // across different NewDRAPluginClient calls. func NewDRAPluginClient(pluginName string) (*Plugin, error) { if pluginName == "" { return nil, fmt.Errorf("plugin name is empty") } </ins> <del> var _ cache.PluginHandler = &RegistrationHandler{} // NewPluginHandler returns new registration handler. // // Must only be called once per process because it manages global state. // If a kubeClient is provided, then it synchronizes ResourceSlices // with the resource information provided by plugins. func NewRegistrationHandler(kubeClient kubernetes.Interface, getNode func() (*v1.Node, error)) *RegistrationHandler { handler := &RegistrationHandler{ // The context and thus logger should come from the caller. backgroundCtx: klog.NewContext(context.TODO(), klog.LoggerWithName(klog.TODO(), "DRA registration handler")), kubeClient:    kubeClient, getNode:       getNode, </del> <ins> existingPlugin := draPlugins.get(pluginName) if existingPlugin == nil { return nil, fmt.Errorf("plugin name %s not found in the list of registered DRA plugins", pluginName) </ins> } <del> // When kubelet starts up, no DRA driver has registered yet. None of // the drivers are usable until they come back, which might not happen // at all. Therefore it is better to not advertise any local resources // because pods could get stuck on the node waiting for the driver // to start up. // // This has to run in the background. go handler.wipeResourceSlices("") </del> <ins> return existingPlugin, nil } type Plugin struct { backgroundCtx context.Context cancel        func(cause error) </ins> <del> return handler </del> <ins> mutex                   sync.Mutex conn                    *grpc.ClientConn endpoint                string highestSupportedVersion *utilversion.Version clientCallTimeout       time.Duration </ins> } <del> // wipeResourceSlices deletes ResourceSlices of the node, optionally just for a specific driver. func (h *RegistrationHandler) wipeResourceSlices(driver string) { if h.kubeClient == nil { return } ctx := h.backgroundCtx logger := klog.FromContext(ctx) </del> <ins> func (p *Plugin) getOrCreateGRPCConn() (*grpc.ClientConn, error) { p.mutex.Lock() defer p.mutex.Unlock() </ins> <del> backoff := wait.Backoff{ Duration: time.Second, Factor:   2, Jitter:   0.2, Cap:      5 * time.Minute, Steps:    100, </del> <ins> if p.conn != nil { return p.conn, nil </ins> } <del> // Error logging is done inside the loop. Context cancellation doesn't get logged. _ = wait.ExponentialBackoffWithContext(ctx, backoff, func(ctx context.Context) (bool, error) { node, err := h.getNode() if apierrors.IsNotFound(err) { return false, nil } if err != nil { logger.Error(err, "Unexpected error checking for node") return false, nil } fieldSelector := fields.Set{resourceapi.ResourceSliceSelectorNodeName: node.Name} if driver != "" { fieldSelector[resourceapi.ResourceSliceSelectorDriver] = driver } err = h.kubeClient.ResourceV1alpha3().ResourceSlices().DeleteCollection(ctx, metav1.DeleteOptions{}, metav1.ListOptions{FieldSelector: fieldSelector.String()}) switch { case err == nil: logger.V(3).Info("Deleted ResourceSlices", "fieldSelector", fieldSelector) return true, nil case apierrors.IsUnauthorized(err): // This can happen while kubelet is still figuring out // its credentials. logger.V(5).Info("Deleting ResourceSlice failed, retrying", "fieldSelector", fieldSelector, "err", err) return false, nil default: // Log and retry for other errors. logger.V(3).Info("Deleting ResourceSlice failed, retrying", "fieldSelector", fieldSelector, "err", err) return false, nil } }) } // RegisterPlugin is called when a plugin can be registered. func (h *RegistrationHandler) RegisterPlugin(pluginName string, endpoint string, versions []string, pluginClientTimeout *time.Duration) error { // Prepare a context with its own logger for the plugin. // // The lifecycle of the plugin's background activities is tied to our // root context, so canceling that will also cancel the plugin. // // The logger injects the plugin name as additional value // into all log output related to the plugin. ctx := h.backgroundCtx </del> <ins> ctx := p.backgroundCtx </ins> logger := klog.FromContext(ctx) <del> logger = klog.LoggerWithValues(logger, "pluginName", pluginName) ctx = klog.NewContext(ctx, logger) logger.V(3).Info("Register new DRA plugin", "endpoint", endpoint) highestSupportedVersion, err := h.validateVersions(pluginName, versions) </del> <ins> network := "unix" logger.V(4).Info("Creating new gRPC connection", "protocol", network, "endpoint", p.endpoint) // grpc.Dial is deprecated. grpc.NewClient should be used instead. // For now this gets ignored because this function is meant to establish // the connection, with the one second timeout below. Perhaps that // approach should be reconsidered? //nolint:staticcheck conn, err := grpc.Dial( p.endpoint, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(func(ctx context.Context, target string) (net.Conn, error) { return (&net.Dialer{}).DialContext(ctx, network, target) }), ) </ins> if err != nil { <del> return fmt.Errorf("version check of plugin %s failed: %w", pluginName, err) } var timeout time.Duration if pluginClientTimeout == nil { timeout = defaultClientCallTimeout } else { timeout = *pluginClientTimeout </del> <ins> return nil, err </ins> } <del> ctx, cancel := context.WithCancelCause(ctx) </del> <ins> ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() </ins> <del> pluginInstance := &Plugin{ backgroundCtx:           ctx, cancel:                  cancel, conn:                    nil, endpoint:                endpoint, highestSupportedVersion: highestSupportedVersion, clientCallTimeout:       timeout, </del> <ins> if ok := conn.WaitForStateChange(ctx, connectivity.Connecting); !ok { return nil, errors.New("timed out waiting for gRPC connection to be ready") </ins> } <del> // Storing endpoint of newly registered DRA Plugin into the map, where plugin name will be the key // all other DRA components will be able to get the actual socket of DRA plugins by its name. if draPlugins.add(pluginName, pluginInstance) { logger.V(1).Info("Already registered, previous plugin was replaced") } return nil </del> <ins> p.conn = conn return p.conn, nil </ins> } <del> func (h *RegistrationHandler) validateVersions( pluginName string, versions []string, ) (*utilversion.Version, error) { if len(versions) == 0 { return nil, errors.New("empty list for supported versions") } </del> <ins> func (p *Plugin) NodePrepareResources( ctx context.Context, req *drapb.NodePrepareResourcesRequest, opts ...grpc.CallOption, ) (*drapb.NodePrepareResourcesResponse, error) { logger := klog.FromContext(ctx) logger.V(4).Info("Calling NodePrepareResources rpc", "request", req) </ins> <del> // Validate version newPluginHighestVersion, err := utilversion.HighestSupportedVersion(versions) </del> <ins> conn, err := p.getOrCreateGRPCConn() </ins> if err != nil { <del> // HighestSupportedVersion includes the list of versions in its error // if relevant, no need to repeat it here. return nil, fmt.Errorf("none of the versions are supported: %w", err) } existingPlugin := draPlugins.get(pluginName) if existingPlugin == nil { return newPluginHighestVersion, nil </del> <ins> return nil, err </ins> } <del> if existingPlugin.highestSupportedVersion.LessThan(newPluginHighestVersion) { return newPluginHighestVersion, nil } return nil, fmt.Errorf("another plugin instance is already registered with a higher supported version: %q < %q", newPluginHighestVersion, existingPlugin.highestSupportedVersion) } // DeRegisterPlugin is called when a plugin has removed its socket, // signaling it is no longer available. func (h *RegistrationHandler) DeRegisterPlugin(pluginName string) { if p := draPlugins.delete(pluginName); p != nil { logger := klog.FromContext(p.backgroundCtx) logger.V(3).Info("Deregister DRA plugin", "endpoint", p.endpoint) // Clean up the ResourceSlices for the deleted Plugin since it // may have died without doing so itself and might never come // back. go h.wipeResourceSlices(pluginName) return } </del> <ins> ctx, cancel := context.WithTimeout(ctx, p.clientCallTimeout) defer cancel() </ins> <del> logger := klog.FromContext(h.backgroundCtx) logger.V(3).Info("Deregister DRA plugin not necessary, was already removed") </del> <ins> nodeClient := drapb.NewNodeClient(conn) response, err := nodeClient.NodePrepareResources(ctx, req) logger.V(4).Info("Done calling NodePrepareResources rpc", "response", response, "err", err) return response, err </ins> } <del> // ValidatePlugin is called by kubelet's plugin watcher upon detection // of a new registration socket opened by DRA plugin. func (h *RegistrationHandler) ValidatePlugin(pluginName string, endpoint string, versions []string) error { _, err := h.validateVersions(pluginName, versions) </del> <ins> func (p *Plugin) NodeUnprepareResources( ctx context.Context, req *drapb.NodeUnprepareResourcesRequest, opts ...grpc.CallOption, ) (*drapb.NodeUnprepareResourcesResponse, error) { logger := klog.FromContext(ctx) logger.V(4).Info("Calling NodeUnprepareResource rpc", "request", req) conn, err := p.getOrCreateGRPCConn() </ins> if err != nil { <del> return fmt.Errorf("invalid versions of plugin %s: %w", pluginName, err) </del> <ins> return nil, err </ins> } <del> return err </del> <ins> ctx, cancel := context.WithTimeout(ctx, p.clientCallTimeout) defer cancel() nodeClient := drapb.NewNodeClient(conn) response, err := nodeClient.NodeUnprepareResources(ctx, req) logger.V(4).Info("Done calling NodeUnprepareResources rpc", "response", response, "err", err) return response, err </ins> }

Related discussion:
text: The PR above moved everything related to into because that was where most of the code already was, but didn't renamed the files to avoid a larger diff and merge conflicts. Once is merged, we can do the renaming. Code clarity. /sig node
/triage accepted /priority important-soon
/assign
a PR can be created for this now.
will do