Repository: koordinator-sh/koordinator

Issue:
[question] NRI Plugin Conflict: Duplicate CPU Pinning Attempt Detected
<!-- Please only use this template for submitting question -->

**What happened:**
To enable Koordinator's [NRI mode resource management](https://koordinator.sh/zh-Hans/docs/designs/nri-mode-resource-management), we have enabled containerd NRI. However, intermittent errors occur with the following message :
"failed to get NRI adjustment for container: plugins "00-nri/nri.sock" and "00-nri/nri.sock" both tried to set CPU pinning"

**What you expected to happen:**
No error above should occur after this operation(enable containerd NRI).

**Environment:**

<!-- Please add environment information, if you know what is necessary, just provide some -->

- Koordinator version: v1.5.0
- Kubernetes version (use kubectl version):  v1.28.15+rke2r1
- docker/containerd version: containerd 1.7.22-k3s1
- OS (e.g: cat /etc/os-release): Ubuntu 22.04.3 LTS
- Kernel (e.g. uname -a): Linux  5.15.0-126-generic # 136-Ubuntu SMP Wed Nov 6 10:38:22 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux

**Anything else we need to know:**
The issue occurs intermittently. Restarting koordlet temporarily resolves the problem. Upon reviewing the containerd logs, we observed two consecutive NRI socket disconnections.

Hints from issue comments:
@songtao98: PTAL @saintube @saintube: PTAL /cc @kangclzjc 

link: #1972 @Klay4: any news?@qinfustu: We have encountered this issue multiple times in our production environment. The scenario is likely to occur when node load is high or when containerd restarts, although the exact reproduction scenario is not yet confirmed. By starting two koordlet pods on the same node, we were able to reproduce the problem, confirming that it is caused by duplicate registration of NRI plugins by koordlet. The root cause may lie in the on **(func (p *NriServer) onClose())**.

If the root cause cannot be identified through investigation, could we consider an alternative solution: when containerd restarts or is upgraded, have koordlet exit abnormally and then restart automatically to restore the connection with containerd?@qinfustu: Before successfully creating a new newStub assigning it to p.stub, the old p.stub is not explicitly closed or cleaned up. If p.stub holds resources such as sock connections, this may cause resource leaks. Before replacing p.stub, safely close and clean up the old p.stub first. However, it is difficult to reproduce and verify this issue.​

```go
...
newStub, err := stub.New(p, append(p.stubOpts, stub.WithOnClose(p.onClose))...)
		if err != nil {
			klog.Errorf("failed to create plugin stub: %v", err)
			return false, nil
		}
		p.stub.Stop() // safely close the old p.stub first
		p.stub = newStub
		err = p.Start()
		if err != nil {
...
```@ZiMengSheng: > We have encountered this issue multiple times in our production environment. The scenario is likely to occur when node load is high or when containerd restarts, although the exact reproduction scenario is not yet confirmed. By starting two koordlet pods on the same node, we were able to reproduce the problem, confirming that it is caused by duplicate registration of NRI plugins by koordlet. The root cause may lie in the on **(func (p *NriServer) onClose())**.
> 
> If the root cause cannot be identified through investigation, could we consider an alternative solution: when containerd restarts or is upgraded, have koordlet exit abnormally and then restart automatically to restore the connection with containerd?

Is there no better way to close it? @saintube @saintube: > Before successfully creating a new newStub assigning it to p.stub, the old p.stub is not explicitly closed or cleaned up. If p.stub holds resources such as sock connections, this may cause resource leaks. Before replacing p.stub, safely close and clean up the old p.stub first. However, it is difficult to reproduce and verify this issue.​
> 
> ...
> newStub, err := stub.New(p, append(p.stubOpts, stub.WithOnClose(p.onClose))...)
> 		if err != nil {
> 			klog.Errorf("failed to create plugin stub: %v", err)
> 			return false, nil
> 		}
> 		p.stub.Stop() // safely close the old p.stub first
> 		p.stub = newStub
> 		err = p.Start()
> 		if err != nil {
> ...

@qinfustu This logic should be triggered during the `onClose` of the old stub, which should not be invoked more than once for a stub's lifecycle, so the abnormal might be the invocation of the NRI Stub's onClose. Shall we check the NRI-related logs of containerd and koordlet about when a stub is closed and new?@qinfustu: > > Before successfully creating a new newStub assigning it to p.stub, the old p.stub is not explicitly closed or cleaned up. If p.stub holds resources such as sock connections, this may cause resource leaks. Before replacing p.stub, safely close and clean up the old p.stub first. However, it is difficult to reproduce and verify this issue.​
> > ...
> > newStub, err := stub.New(p, append(p.stubOpts, stub.WithOnClose(p.onClose))...)
> > if err != nil {
> > klog.Errorf("failed to create plugin stub: %v", err)
> > return false, nil
> > }
> > p.stub.Stop() // safely close the old p.stub first
> > p.stub = newStub
> > err = p.Start()
> > if err != nil {
> > ...
> 
> [@qinfustu](https://github.com/qinfustu) This logic should be triggered during the `onClose` of the old stub, which should not be invoked more than once for a stub's lifecycle, so the abnormal might be the invocation of the NRI Stub's onClose. Shall we check the NRI-related logs of containerd and koordlet about when a stub is closed and new?

[@saintube](https://github.com/saintube) No relevant useful information was found in the logs. 
It is confirmed that stub.close()was called before onClose was triggered, but there is a backoff and retry mechanism in onClose(). During the retry process, the newly created stubs did not have stub.close() called? Is it possible that the stub was created successfully, but the start failed, while the plugin was actually registered successfully? Subsequent retries would then create new stubs, leading to duplicate registrations. Would adding a p.stub.Stop() here also have no side effects?@yyrdl: 也遇到了同样的问题，这个事情 koordlet 实现有锅， containerd 的实现也不牢靠,对同一个 plugin 应该只保持一个 connection , [这个地方](https://github.com/containerd/containerd/blob/main/vendor/github.com/containerd/nri/pkg/adaptation/adaptation.go#L612) 应该对同一个plugin 的 conn 只保留一个的，后面的操作只是将 closed 的 conn 去掉了。 

下面是 containerd 的日志，确实是 koordlet 在重复注册，为啥重复注册，从代码还没找出来，有进展我同步一下

<img width="1528" height="465" alt="Image" src="https://github.com/user-attachments/assets/d338b4b0-9c7b-4823-92c1-c30c0e2a6ecc" />
@qinfustu: @yyrdl 请问containerd打印了重复注册的日志后，启动pod时是否出现了"failed to get NRI adjustment for container: plugins "00-nri/nri.sock" and "00-nri/nri.sock" both tried to set CPU pinning"的报错？如果是，请问是怎么复现的。我看[其它的nri plugin中](https://github.com/search?q=repo%3Acontainers%2Fnri-plugins+onclose&type=code)，onClose的逻辑是直接exit(0)，如果查不出根因，我觉得这种实现改成让koordlet重启也能接受，虽然不是很优雅，但是确实能解决线上问题。@yyrdl: > [@yyrdl](https://github.com/yyrdl) 请问containerd打印了重复注册的日志后，启动pod时是否出现了"failed to get NRI adjustment for container: plugins "00-nri/nri.sock" and "00-nri/nri.sock" both tried to set CPU pinning"的报错？如果是，请问是怎么复现的。我看[其它的nri plugin中](https://github.com/search?q=repo%3Acontainers%2Fnri-plugins+onclose&type=code)，onClose的逻辑是直接exit(0)，如果查不出根因，我觉得这种实现改成让koordlet重启也能接受，虽然不是很优雅，但是确实能解决线上问题。

类似的错，重启的话， containerd 内部持有的 conn 确实会断开， 能解决问题，但太丑了  ，你可以先不走 nri ， 走 reconcile ，待问题解决了，再用 nri  @yyrdl: @saintube @qinfustu  @kangclzjc  bug 找到了，  koordinator 的 pipeline 开一下 race 检测吧 ，这个 [stub](https://github.com/koordinator-sh/koordinator/blob/main/pkg/koordlet/runtimehooks/nri/server.go#L276) 的数据同步没做, p.Start 其实是使用的老的 stub ，稳定返回 error ,进一步导致重建新的 stub .. . 所以出现了这种奇怪的日志

<img width="769" height="410" alt="Image" src="https://github.com/user-attachments/assets/cd3d207f-6c6f-4018-82eb-c0dbac061866" />@qinfustu: @yyrdl  🐮，请问数据同步怎么做？p.stub = newStub之前要加锁？@yyrdl: c

> [@yyrdl](https://github.com/yyrdl) 🐮，请问数据同步怎么做？p.stub = newStub之前要加锁？

还有另外一个原因是 onClose 存在递归调用@qinfustu: > c
> 
> > [@yyrdl](https://github.com/yyrdl) 🐮，请问数据同步怎么做？p.stub = newStub之前要加锁？
> 
> 还有另外一个原因是 onClose 存在递归调用

@yyrdl 大佬能给一个完整的解决方案吗？我这边不能稳定复现，修改后不能确定问题是否修复。@yyrdl: > > c
> > > [@yyrdl](https://github.com/yyrdl) 🐮，请问数据同步怎么做？p.stub = newStub之前要加锁？
> > 
> > 
> > 还有另外一个原因是 onClose 存在递归调用
> 
> [@yyrdl](https://github.com/yyrdl) 大佬能给一个完整的解决方案吗？我这边不能稳定复现，修改后不能确定问题是否修复。

在修@yyrdl: @qinfustu  你们的 koordlet 是用容器部署的么 ？ 是不是也出现了 OOM ？@qinfustu: > [@qinfustu](https://github.com/qinfustu) 你们的 koordlet 是用容器部署的么 ？ 是不是也出现了 OOM ？

遇到过，koordlet会把metrics保存到内存中，默认配置的memory比较小，运行一段时间后会OOM，可以把koordlet的limit调大一点。@yyrdl: @qinfustu  mr 提交了@qinfustu: > [@qinfustu](https://github.com/qinfustu) mr 提交了

nice!

Solution patch:
diff --git a/pkg/koordlet/runtimehooks/nri/handlers.go b/pkg/koordlet/runtimehooks/nri/handlers.go
new file mode 100644
index 00000000..f5645be7
--- /dev/null
+++ b/pkg/koordlet/runtimehooks/nri/handlers.go
@@ -0,0 +1,137 @@
+/*
+Copyright 2022 The Koordinator Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package nri
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/containerd/nri/pkg/api"
+	"github.com/containerd/nri/pkg/stub"
+	"k8s.io/klog/v2"
+	"sigs.k8s.io/yaml"
+
+	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/hooks"
+	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/protocol"
+	rmconfig "github.com/koordinator-sh/koordinator/pkg/runtimeproxy/config"
+)
+
+func (p *NriServer) Configure(_ context.Context, config, runtime, version string) (stub.EventMask, error) {
+	klog.V(4).Infof("got configuration data: %q from runtime %s %s", config, runtime, version)
+	if config == "" {
+		return p.mask, nil
+	}
+
+	err := yaml.Unmarshal([]byte(config), &p.cfg)
+	if err != nil {
+		return 0, fmt.Errorf("failed to parse provided configuration: %w", err)
+	}
+
+	p.mask, err = api.ParseEventMask(p.cfg.Events...)
+	if err != nil {
+		return 0, fmt.Errorf("failed to parse events in configuration: %w", err)
+	}
+
+	klog.V(6).Infof("handle NRI Configure successfully, config %s, runtime %s, version %s",
+		config, runtime, version)
+	return p.mask, nil
+}
+
+func (p *NriServer) Synchronize(_ context.Context, pods []*api.PodSandbox, containers []*api.Container) ([]*api.ContainerUpdate, error) {
+	// todo: update existed containers configure
+	return nil, nil
+}
+
+func (p *NriServer) RunPodSandbox(_ context.Context, pod *api.PodSandbox) error {
+	podCtx := &protocol.PodContext{}
+	podCtx.FromNri(pod)
+	// todo: return error or bypass error based on PluginFailurePolicy
+	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreRunPodSandbox, podCtx)
+	if err != nil {
+		klog.Errorf("nri hooks run error: %v", err)
+		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
+			return err
+		}
+	}
+	podCtx.NriDone(p.options.Executor)
+
+	klog.V(6).Infof("handle NRI RunPodSandbox successfully, pod %s/%s", pod.GetNamespace(), pod.GetName())
+	return nil
+}
+
+func (p *NriServer) CreateContainer(_ context.Context, pod *api.PodSandbox, container *api.Container) (*api.ContainerAdjustment, []*api.ContainerUpdate, error) {
+	containerCtx := &protocol.ContainerContext{}
+	containerCtx.FromNri(pod, container)
+	// todo: return error or bypass error based on PluginFailurePolicy
+	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreCreateContainer, containerCtx)
+	if err != nil {
+		klog.Errorf("nri run hooks error: %v", err)
+		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
+			return nil, nil, err
+		}
+	}
+
+	adjust, _, err := containerCtx.NriDone(p.options.Executor)
+	if err != nil {
+		klog.Errorf("containerCtx nri done failed: %v", err)
+		return nil, nil, nil
+	}
+
+	klog.V(6).Infof("handle NRI CreateContainer successfully, container %s/%s/%s",
+		pod.GetNamespace(), pod.GetName(), container.GetName())
+	return adjust, nil, nil
+}
+
+func (p *NriServer) UpdateContainer(_ context.Context, pod *api.PodSandbox, container *api.Container, r *api.LinuxResources) ([]*api.ContainerUpdate, error) {
+	containerCtx := &protocol.ContainerContext{}
+	containerCtx.FromNri(pod, container)
+	// todo: return error or bypass error based on PluginFailurePolicy
+	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreUpdateContainerResources, containerCtx)
+	if err != nil {
+		klog.Errorf("nri run hooks error: %v", err)
+		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
+			return nil, err
+		}
+	}
+
+	_, update, err := containerCtx.NriDone(p.options.Executor)
+	if err != nil {
+		klog.Errorf("containerCtx nri done failed: %v", err)
+		return nil, nil
+	}
+
+	klog.V(6).Infof("handle NRI UpdateContainer successfully, container %s/%s/%s",
+		pod.GetNamespace(), pod.GetName(), container.GetName())
+	return []*api.ContainerUpdate{update}, nil
+}
+
+func (p *NriServer) RemovePodSandbox(_ context.Context, pod *api.PodSandbox) error {
+	podCtx := &protocol.PodContext{}
+	podCtx.FromNri(pod)
+	// todo: return error or bypass error based on PluginFailurePolicy
+	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreRemoveRunPodSandbox, podCtx)
+	if err != nil {
+		klog.Errorf("nri hooks run error: %v", err)
+		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
+			return err
+		}
+	}
+	podCtx.NriRemoveDone(p.options.Executor)
+
+	klog.V(6).Infof("handle NRI RemovePodSandbox successfully, pod %s/%s", pod.GetNamespace(), pod.GetName())
+	return nil
+}
diff --git a/pkg/koordlet/runtimehooks/nri/mock_nri.go b/pkg/koordlet/runtimehooks/nri/mock_nri.go
new file mode 100644
index 00000000..d4f6f6c3
--- /dev/null
+++ b/pkg/koordlet/runtimehooks/nri/mock_nri.go
@@ -0,0 +1,119 @@
+/*
+Copyright 2022 The Koordinator Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Code generated by MockGen. DO NOT EDIT.
+// Source: interface.go
+
+// Package nri is a generated GoMock package.
+package nri
+
+import (
+	context "context"
+	reflect "reflect"
+
+	api "github.com/containerd/nri/pkg/api"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockStubInterface is a mock of StubInterface interface.
+type MockStubInterface struct {
+	ctrl     *gomock.Controller
+	recorder *MockStubInterfaceMockRecorder
+}
+
+// MockStubInterfaceMockRecorder is the mock recorder for MockStubInterface.
+type MockStubInterfaceMockRecorder struct {
+	mock *MockStubInterface
+}
+
+// NewMockStubInterface creates a new mock instance.
+func NewMockStubInterface(ctrl *gomock.Controller) *MockStubInterface {
+	mock := &MockStubInterface{ctrl: ctrl}
+	mock.recorder = &MockStubInterfaceMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockStubInterface) EXPECT() *MockStubInterfaceMockRecorder {
+	return m.recorder
+}
+
+// Run mocks base method.
+func (m *MockStubInterface) Run(arg0 context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Run", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Run indicates an expected call of Run.
+func (mr *MockStubInterfaceMockRecorder) Run(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Run", reflect.TypeOf((*MockStubInterface)(nil).Run), arg0)
+}
+
+// Start mocks base method.
+func (m *MockStubInterface) Start(arg0 context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Start", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Start indicates an expected call of Start.
+func (mr *MockStubInterfaceMockRecorder) Start(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*MockStubInterface)(nil).Start), arg0)
+}
+
+// Stop mocks base method.
+func (m *MockStubInterface) Stop() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Stop")
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockStubInterfaceMockRecorder) Stop() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockStubInterface)(nil).Stop))
+}
+
+// UpdateContainers mocks base method.
+func (m *MockStubInterface) UpdateContainers(arg0 []*api.ContainerUpdate) ([]*api.ContainerUpdate, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateContainers", arg0)
+	ret0, _ := ret[0].([]*api.ContainerUpdate)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UpdateContainers indicates an expected call of UpdateContainers.
+func (mr *MockStubInterfaceMockRecorder) UpdateContainers(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateContainers", reflect.TypeOf((*MockStubInterface)(nil).UpdateContainers), arg0)
+}
+
+// Wait mocks base method.
+func (m *MockStubInterface) Wait() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Wait")
+}
+
+// Wait indicates an expected call of Wait.
+func (mr *MockStubInterfaceMockRecorder) Wait() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Wait", reflect.TypeOf((*MockStubInterface)(nil).Wait))
+}
diff --git a/pkg/koordlet/runtimehooks/nri/option.go b/pkg/koordlet/runtimehooks/nri/option.go
new file mode 100644
index 00000000..ce9eb9da
--- /dev/null
+++ b/pkg/koordlet/runtimehooks/nri/option.go
@@ -0,0 +1,53 @@
+/*
+Copyright 2022 The Koordinator Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package nri
+
+import (
+	"fmt"
+	"path/filepath"
+	"time"
+
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/tools/record"
+
+	"github.com/koordinator-sh/koordinator/pkg/koordlet/resourceexecutor"
+	"github.com/koordinator-sh/koordinator/pkg/koordlet/util/system"
+	rmconfig "github.com/koordinator-sh/koordinator/pkg/runtimeproxy/config"
+)
+
+type Options struct {
+	NriPluginName     string
+	NriPluginIdx      string
+	NriSocketPath     string
+	NriConnectTimeout time.Duration
+	// support stop running other hooks once someone failed
+	PluginFailurePolicy rmconfig.FailurePolicyType
+	// todo: add support for disable stages
+	DisableStages map[string]struct{}
+	Executor      resourceexecutor.ResourceUpdateExecutor
+	BackOff       wait.Backoff
+	EventRecorder record.EventRecorder
+}
+
+func (o Options) Validate() error {
+	// a fast check for the NRI support status
+	completeNriSocketPath := filepath.Join(system.Conf.VarRunRootDir, o.NriSocketPath)
+	if !system.FileExists(completeNriSocketPath) {
+		return fmt.Errorf("nri socket path %q does not exist", completeNriSocketPath)
+	}
+	return nil
+}
diff --git a/pkg/koordlet/runtimehooks/nri/server.go b/pkg/koordlet/runtimehooks/nri/server.go
index 50bd5e1b..f420266c 100644
--- a/pkg/koordlet/runtimehooks/nri/server.go
+++ b/pkg/koordlet/runtimehooks/nri/server.go
@@ -21,58 +21,45 @@ import (
 	"fmt"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/containerd/nri/pkg/api"
 	"github.com/containerd/nri/pkg/stub"
 	"go.uber.org/atomic"
 	"k8s.io/apimachinery/pkg/util/wait"
-	"k8s.io/client-go/tools/record"
 	"k8s.io/klog/v2"
-	"sigs.k8s.io/yaml"
 
-	"github.com/koordinator-sh/koordinator/pkg/koordlet/resourceexecutor"
-	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/hooks"
-	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/protocol"
 	"github.com/koordinator-sh/koordinator/pkg/koordlet/util/system"
-	rmconfig "github.com/koordinator-sh/koordinator/pkg/runtimeproxy/config"
 )
 
 type nriConfig struct {
 	Events []string `json:"events"`
 }
 
-type Options struct {
-	NriPluginName     string
-	NriPluginIdx      string
-	NriSocketPath     string
-	NriConnectTimeout time.Duration
-	// support stop running other hooks once someone failed
-	PluginFailurePolicy rmconfig.FailurePolicyType
-	// todo: add support for disable stages
-	DisableStages map[string]struct{}
-	Executor      resourceexecutor.ResourceUpdateExecutor
-	BackOff       wait.Backoff
-	EventRecorder record.EventRecorder
-}
-
-func (o Options) Validate() error {
-	// a fast check for the NRI support status
-	completeNriSocketPath := filepath.Join(system.Conf.VarRunRootDir, o.NriSocketPath)
-	if !system.FileExists(completeNriSocketPath) {
-		return fmt.Errorf("nri socket path %q does not exist", completeNriSocketPath)
-	}
-
-	return nil
+// Stub is the interface of nri stub
+//
+//go:generate mockgen -source=server.go -destination=mock_nri.go -package=nri
+type StubInterface interface {
+	stub.Stub
 }
 
 type NriServer struct {
 	cfg      nriConfig
-	stub     stub.Stub
 	mask     stub.EventMask
 	options  Options       // server options
 	stubOpts []stub.Option // nri stub options
-	stopped  *atomic.Bool  // if false, the stub will try to reconnect when stub.OnClose is invoked
+
+	mutex            sync.Mutex
+	stopped          bool // if false, the stub will try to reconnect when stub.OnClose is invoked
+	stub             StubInterface
+	connID           int64
+	connClosedSignal chan struct{}
+}
+
+// make it testable
+var mockableNewNriStub = func(p interface{}, opts []stub.Option, onClose func()) (StubInterface, error) {
+	return stub.New(p, append(opts, stub.WithOnClose(onClose))...)
 }
 
 const (
@@ -93,206 +80,168 @@ func NewNriServer(opt Options) (*NriServer, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to validate nri server, err: %w", err)
 	}
-
 	stubOpts := []stub.Option{
 		stub.WithPluginName(opt.NriPluginName),
 		stub.WithPluginIdx(opt.NriPluginIdx),
 		stub.WithSocketPath(filepath.Join(system.Conf.VarRunRootDir, opt.NriSocketPath)),
 	}
 	p := &NriServer{
-		options:  opt,
-		stubOpts: stubOpts,
-		stopped:  atomic.NewBool(false),
+		options:          opt,
+		stubOpts:         stubOpts,
+		connClosedSignal: make(chan struct{}, 10),
 	}
 	if p.mask, err = api.ParseEventMask(events); err != nil {
 		klog.Errorf("failed to parse events %v", err)
 		return p, err
 	}
 	p.cfg.Events = strings.Split(events, ",")
-
-	if p.stub, err = stub.New(p, append(p.stubOpts, stub.WithOnClose(p.onClose))...); err != nil {
-		klog.Errorf("failed to create plugin stub: %v", err)
-		return nil, err
-	}
-
 	return p, nil
 }
 
-func (p *NriServer) Start() error {
-	err := p.options.Validate()
+func (p *NriServer) Start() (err error) {
+	klog.V(4).Info("starting nri server")
+	defer func() {
+		klog.V(4).Infof("nri server started with err: %v", err)
+	}()
+
+	err = p.options.Validate()
 	if err != nil {
 		return err
 	}
-	success := time.After(p.options.NriConnectTimeout)
-	errorChan := make(chan error)
-
-	go func(chan error) {
-		if p.stub != nil {
-			err := p.stub.Run(context.Background())
-			if err != nil {
-				klog.Errorf("nri server exited with error: %v", err)
-				errorChan <- err
-			} else {
-				klog.V(4).Info("nri server started")
-			}
-		} else {
-			err := fmt.Errorf("nri stub is nil")
-			errorChan <- err
-		}
-	}(errorChan)
-
-	select {
-	case <-success:
-		return nil
-	case <-errorChan:
-		return fmt.Errorf("nri start fail, err: %w", err)
+	// connect at first when start nri server
+	if _, err := p.connect(); err != nil {
+		return fmt.Errorf("failed to connect to nri servver: %v", err)
 	}
+
+	// then keep connection alive forever until stopped
+	go p.keepAlive()
+	return nil
 }
 
 func (p *NriServer) Stop() {
-	p.stopped.Store(true)
-	p.stub.Stop()
+	if stopped := func() bool {
+		p.mutex.Lock()
+		defer p.mutex.Unlock()
+		ret := p.stopped
+		p.stopped = true
+		return ret
+	}(); stopped {
+		return
+	}
+	p.cleanConnection(-1)
 }
 
-func (p *NriServer) Configure(_ context.Context, config, runtime, version string) (stub.EventMask, error) {
-	klog.V(4).Infof("got configuration data: %q from runtime %s %s", config, runtime, version)
-	if config == "" {
-		return p.mask, nil
-	}
-
-	err := yaml.Unmarshal([]byte(config), &p.cfg)
-	if err != nil {
-		return 0, fmt.Errorf("failed to parse provided configuration: %w", err)
-	}
-
-	p.mask, err = api.ParseEventMask(p.cfg.Events...)
-	if err != nil {
-		return 0, fmt.Errorf("failed to parse events in configuration: %w", err)
+func (p *NriServer) keepAlive() {
+	// this action is important, we should not catch panic
+	// so that we can find the problem
+	defer func() {
+		p.mutex.Lock()
+		defer p.mutex.Unlock()
+		// close channel when no consumer
+		close(p.connClosedSignal)
+		p.connClosedSignal = nil
+	}()
+	for {
+		_, ok := <-p.connClosedSignal
+		if !ok {
+			return
+		}
+		if stopped := func() bool {
+			p.mutex.Lock()
+			defer p.mutex.Unlock()
+			return p.stopped
+		}(); stopped {
+			return
+		}
+		_ = wait.ExponentialBackoff(p.options.BackOff, p.connect)
+		if p.connected() {
+			continue
+		}
+		// still not connected ,retry it in next turn
+		select {
+		case p.connClosedSignal <- struct{}{}:
+		default:
+		}
 	}
-
-	klog.V(6).Infof("handle NRI Configure successfully, config %s, runtime %s, version %s",
-		config, runtime, version)
-	return p.mask, nil
 }
 
-func (p *NriServer) Synchronize(_ context.Context, pods []*api.PodSandbox, containers []*api.Container) ([]*api.ContainerUpdate, error) {
-	// todo: update existed containers configure
-	return nil, nil
+func (p *NriServer) connected() bool {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+	return p.stub != nil
 }
 
-func (p *NriServer) RunPodSandbox(_ context.Context, pod *api.PodSandbox) error {
-	podCtx := &protocol.PodContext{}
-	podCtx.FromNri(pod)
-	// todo: return error or bypass error based on PluginFailurePolicy
-	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreRunPodSandbox, podCtx)
-	if err != nil {
-		klog.Errorf("nri hooks run error: %v", err)
-		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
-			return err
-		}
+func (p *NriServer) connect() (bool, error) {
+	if p.connected() {
+		return true, nil
 	}
-	podCtx.NriDone(p.options.Executor)
-
-	klog.V(6).Infof("handle NRI RunPodSandbox successfully, pod %s/%s", pod.GetNamespace(), pod.GetName())
-	return nil
-}
-
-func (p *NriServer) CreateContainer(_ context.Context, pod *api.PodSandbox, container *api.Container) (*api.ContainerAdjustment, []*api.ContainerUpdate, error) {
-	containerCtx := &protocol.ContainerContext{}
-	containerCtx.FromNri(pod, container)
-	// todo: return error or bypass error based on PluginFailurePolicy
-	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreCreateContainer, containerCtx)
+	connID := time.Now().UnixNano()
+	klog.Infof("connect to nri server, conn_id: %d", connID)
+	stubIns, err := mockableNewNriStub(p, p.stubOpts, p.reconnect(connID))
 	if err != nil {
-		klog.Errorf("nri run hooks error: %v", err)
-		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
-			return nil, nil, err
-		}
+		return false, err
 	}
-
-	adjust, _, err := containerCtx.NriDone(p.options.Executor)
+	err = p.tryDial(connID, stubIns, p.options.NriConnectTimeout)
 	if err != nil {
-		klog.Errorf("containerCtx nri done failed: %v", err)
-		return nil, nil, nil
-	}
-
-	klog.V(6).Infof("handle NRI CreateContainer successfully, container %s/%s/%s",
-		pod.GetNamespace(), pod.GetName(), container.GetName())
-	return adjust, nil, nil
+		klog.Errorf("failed to connect to nri server, conn_id: %d, %v", connID, err)
+		return false, err
+	}
+	p.mutex.Lock()
+	p.stub = stubIns
+	p.connID = connID
+	p.mutex.Unlock()
+	klog.Infof("connect to nri server successfully, conn_id: %d", connID)
+	return true, nil
 }
 
-func (p *NriServer) UpdateContainer(_ context.Context, pod *api.PodSandbox, container *api.Container, r *api.LinuxResources) ([]*api.ContainerUpdate, error) {
-	containerCtx := &protocol.ContainerContext{}
-	containerCtx.FromNri(pod, container)
-	// todo: return error or bypass error based on PluginFailurePolicy
-	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreUpdateContainerResources, containerCtx)
-	if err != nil {
-		klog.Errorf("nri run hooks error: %v", err)
-		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
-			return nil, err
+func (p *NriServer) tryDial(connID int64, stubIns stub.Stub, atLeastLiveTime time.Duration) error {
+	success := time.After(atLeastLiveTime)
+	errChan := make(chan error)
+	returned := atomic.NewBool(false)
+	go func() {
+		err := stubIns.Run(context.Background())
+		if err == nil {
+			return
 		}
+		p.cleanConnection(connID)
+		if !returned.Load() {
+			// make sure the consumer exists, otherwise the goroutine will be blocked
+			// and never exit, which will cause memory leak
+			errChan <- err
+		}
+		close(errChan)
+	}()
+	select {
+	case <-success:
+		returned.Store(true)
+		return nil
+	case err := <-errChan:
+		return err
 	}
-
-	_, update, err := containerCtx.NriDone(p.options.Executor)
-	if err != nil {
-		klog.Errorf("containerCtx nri done failed: %v", err)
-		return nil, nil
-	}
-
-	klog.V(6).Infof("handle NRI UpdateContainer successfully, container %s/%s/%s",
-		pod.GetNamespace(), pod.GetName(), container.GetName())
-	return []*api.ContainerUpdate{update}, nil
 }
 
-func (p *NriServer) RemovePodSandbox(_ context.Context, pod *api.PodSandbox) error {
-	podCtx := &protocol.PodContext{}
-	podCtx.FromNri(pod)
-	// todo: return error or bypass error based on PluginFailurePolicy
-	err := hooks.RunHooks(p.options.PluginFailurePolicy, rmconfig.PreRemoveRunPodSandbox, podCtx)
-	if err != nil {
-		klog.Errorf("nri hooks run error: %v", err)
-		if p.options.PluginFailurePolicy == rmconfig.PolicyFail {
-			return err
-		}
+func (p *NriServer) reconnect(connID int64) func() {
+	return func() {
+		p.cleanConnection(connID)
 	}
-	podCtx.NriRemoveDone(p.options.Executor)
-
-	klog.V(6).Infof("handle NRI RemovePodSandbox successfully, pod %s/%s", pod.GetNamespace(), pod.GetName())
-	return nil
 }
 
-func (p *NriServer) onClose() {
-	//TODO: consider the pod status during restart
-	retryFunc := func() (bool, error) {
-		if p.stopped.Load() { // if set to stopped, no longer reconnect
-			return true, nil
+func (p *NriServer) cleanConnection(connID int64) {
+	klog.Warningf("nri server connection closed, conn_id: %d", connID)
+	defer klog.Infof("clean dead stub successfully, conn_id: %d", connID)
+	p.mutex.Lock()
+	if p.stub != nil && (p.connID == connID || connID == -1) {
+		p.stub.Stop()
+		// wait for the plugin to stop
+		p.stub.Wait()
+		p.stub = nil
+		p.connID = -1
+	}
+	if p.connClosedSignal != nil {
+		select {
+		case p.connClosedSignal <- struct{}{}:
+		default:
 		}
-
-		newStub, err := stub.New(p, append(p.stubOpts, stub.WithOnClose(p.onClose))...)
-		if err != nil {
-			klog.Errorf("failed to create plugin stub: %v", err)
-			return false, nil
-		}
-
-		p.stub = newStub
-		err = p.Start()
-		if err != nil {
-			completeNriSocketPath := filepath.Join(system.Conf.VarRunRootDir, p.options.NriSocketPath)
-			targetErr := fmt.Errorf("nri socket path %q does not exist", completeNriSocketPath)
-			if err.Error() == targetErr.Error() {
-				return false, err
-			}
-			//TODO: check the error type, if nri server disable nri, we should also break backoff
-			klog.Warningf("nri reconnect failed, err: %s", err)
-			return false, nil
-		} else {
-			klog.V(4).Info("nri server restart success")
-			return true, nil
-		}
-	}
-
-	// TODO: high version wait not support BackoffUntil with BackOffManger as parameters, when updated to v0.27.0 version wait, we can refine ExponentialBackoff.
-	err := wait.ExponentialBackoff(p.options.BackOff, retryFunc)
-	if err != nil {
-		klog.Errorf("nri server restart failed after several times retry")
 	}
+	p.mutex.Unlock()
 }



Test patch:
diff --git a/pkg/koordlet/runtimehooks/nri/server_test.go b/pkg/koordlet/runtimehooks/nri/server_test.go
index 4b660974..bfe421a7 100644
--- a/pkg/koordlet/runtimehooks/nri/server_test.go
+++ b/pkg/koordlet/runtimehooks/nri/server_test.go
@@ -25,6 +25,9 @@ import (
 
 	"github.com/containerd/nri/pkg/api"
 	"github.com/containerd/nri/pkg/stub"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/koordinator-sh/koordinator/pkg/koordlet/resourceexecutor"
 	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/hooks"
@@ -633,3 +636,111 @@ func TestNriServer_RemovePodSandbox(t *testing.T) {
 		})
 	}
 }
+
+type stubConstructor func(p interface{}, opts []stub.Option, onClose func()) (StubInterface, error)
+
+func mockNewStub(fn stubConstructor) func() {
+	old := mockableNewNriStub
+	mockableNewNriStub = fn
+	return func() {
+		mockableNewNriStub = old
+	}
+}
+
+func TestStartNriServerFailed(t *testing.T) {
+	a := assert.New(t)
+	ctr := gomock.NewController(t)
+	defer ctr.Finish()
+	fakeStub := NewMockStubInterface(ctr)
+	fakeStub.EXPECT().Run(gomock.Any()).Return(fmt.Errorf("error")).Times(1)
+	defer mockNewStub(func(p interface{}, opts []stub.Option, onClose func()) (StubInterface, error) {
+		return fakeStub, nil
+	})()
+	helper := system.NewFileTestUtil(t)
+	defer helper.Cleanup()
+	srv, err := NewNriServer(Options{
+		NriSocketPath:       "",
+		NriConnectTimeout:   time.Second,
+		PluginFailurePolicy: "Ignore",
+		DisableStages:       getDisableStagesMap([]string{"PreRunPodSandbox"}),
+		Executor:            nil,
+	})
+	a.Nil(err)
+	a.NotNil(srv)
+	err = srv.Start()
+	a.NotNil(err)
+}
+
+func TestReconnectToNriServer(t *testing.T) {
+	a := assert.New(t)
+	ctr := gomock.NewController(t)
+	defer ctr.Finish()
+	stubs := make(chan struct {
+		Stub    *MockStubInterface
+		RunChan chan error
+		OnClose func()
+	})
+	defer mockNewStub(func(p interface{}, opts []stub.Option, onClose func()) (StubInterface, error) {
+		fakeStub := NewMockStubInterface(ctr)
+		st := struct {
+			Stub    *MockStubInterface
+			RunChan chan error
+			OnClose func()
+		}{
+			Stub:    fakeStub,
+			RunChan: make(chan error),
+			OnClose: onClose,
+		}
+		go func() {
+			stubs <- st
+		}()
+		fakeStub.EXPECT().Run(gomock.Any()).DoAndReturn(func(ctx context.Context) error {
+			return <-st.RunChan
+		}).Times(1)
+		return fakeStub, nil
+	})()
+
+	helper := system.NewFileTestUtil(t)
+	defer helper.Cleanup()
+	srv, err := NewNriServer(Options{
+		NriSocketPath:       "",
+		NriConnectTimeout:   time.Second,
+		PluginFailurePolicy: "Ignore",
+		DisableStages:       getDisableStagesMap([]string{"PreRunPodSandbox"}),
+		Executor:            nil,
+		BackOff: wait.Backoff{
+			Duration: time.Second,
+			Factor:   2,
+			Jitter:   0.1,
+			Steps:    3,
+			Cap:      time.Second * 3,
+		},
+	})
+	a.Nil(err)
+	a.NotNil(srv)
+	err = srv.Start()
+	a.Nil(err)
+	for i := 0; i < 4; i++ {
+		st := <-stubs
+		if i == 0 {
+			// test connection closed ,should reconnect
+			st.Stub.EXPECT().Stop().Return().Times(1)
+			st.Stub.EXPECT().Wait().Return().Times(1)
+			st.OnClose()
+		}
+		if i == 1 {
+			// test run stub failed ,should reconnect
+			st.Stub.EXPECT().Stop().Return().Times(1)
+			st.Stub.EXPECT().Wait().Return().Times(1)
+			time.Sleep(time.Second * 3)
+			st.RunChan <- fmt.Errorf("error")
+		}
+		if i == 2 {
+			time.Sleep(time.Second * 3)
+			break
+		}
+	}
+	srv.mutex.Lock()
+	defer srv.mutex.Unlock()
+	a.NotNil(srv.stub)
+}