Repository: kubernetes-sigs/kwok
Pull request #319 (2023-02-22 11:41:41)

clean: Wrapping sync.Map with Generics
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide/first-contribution.md#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/sig-release/release.md#issuepr-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

#### What type of PR is this?

<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

/kind cleanup

#### What this PR does / why we need it:

#### Which issue(s) this PR fixes:
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

#### Special notes for your reviewer:

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, a release note is required:
Enter your extended release note in the block below. If the PR requires additional action from users switching to the new release, include the string "action required".

For more information on release notes see: https://git.k8s.io/community/contributors/guide/release-notes.md
-->
```release-note

```

#### Additional documentation e.g., KEPs (Kubernetes Enhancement Proposals), usage docs, etc.:

<!--
This section can be blank if this pull request does not require a release note.

When adding links which point to resources within git repositories, like
KEPs or supporting documentation, please reference a specific commit and avoid
linking directly to the master branch. This ensures that links reference a
specific point in time, rather than a document that may change over time.

See here for guidance on getting permanent links to files: https://help.github.com/en/articles/getting-permanent-links-to-files

Please use the following format for linking documentation:
- [KEP]: <link>
- [Usage]: <link>
- [Other doc]: <link>
-->
```docs

```


k8s-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/kwok/pull/319#" title="Author self-approved">wzshiming</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fkwok).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubernetes-sigs/kwok/blob/main/OWNERS)~~ [wzshiming]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
netlify[bot]: ### <span aria-hidden="true">👷</span> Deploy Preview for *k8s-kwok* processing.


|  Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | b904f6c4985254c1a62d844a0f685de4de96f366 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/k8s-kwok/deploys/63f5ff77e77e96000866794f |

Patch:
diff --git a/pkg/apis/v1alpha1/kwokctl_configuration_types.go b/pkg/apis/v1alpha1/kwokctl_configuration_types.go
index 07e2133d..b209f494 100644
--- a/pkg/apis/v1alpha1/kwokctl_configuration_types.go
+++ b/pkg/apis/v1alpha1/kwokctl_configuration_types.go
@@ -143,7 +143,7 @@ type KwokctlConfigurationOptions struct {
 	KubeSchedulerImage string `json:"kubeSchedulerImage,omitempty"`
 
 	// KwokControllerImage is the image of Kwok.
-	// is the default value for flag --controller-image and env KWOK_CONTROLLER_IMAGE
+	// is the default value for flag --contoller-image and env KWOK_CONTROLLER_IMAGE
 	KwokControllerImage string `json:"kwokControllerImage,omitempty"`
 
 	// PrometheusImage is the image of Prometheus.
diff --git a/pkg/kwok/controllers/controller.go b/pkg/kwok/controllers/controller.go
index caabe62b..530ceb5c 100644
--- a/pkg/kwok/controllers/controller.go
+++ b/pkg/kwok/controllers/controller.go
@@ -145,7 +145,6 @@ func NewController(conf Config) (*Controller, error) {
 		DisregardStatusWithLabelSelector:      conf.DisregardStatusWithLabelSelector,
 		Stages:                                conf.PodStages,
 		LockPodParallelism:                    16,
-		NodeHasFunc:                           nodes.Has, // just handle pods that are on nodes we have
 		NodeGetFunc:                           nodes.Get,
 		FuncMap:                               defaultFuncMap,
 		Recorder:                              recorder,
diff --git a/pkg/kwok/controllers/node_controller.go b/pkg/kwok/controllers/node_controller.go
index e613f423..75fabad4 100644
--- a/pkg/kwok/controllers/node_controller.go
+++ b/pkg/kwok/controllers/node_controller.go
@@ -22,7 +22,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
-	"sync"
 	"text/template"
 	"time"
 
@@ -43,6 +42,7 @@ import (
 	"sigs.k8s.io/kwok/pkg/apis/internalversion"
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/expression"
+	"sigs.k8s.io/kwok/pkg/utils/maps"
 )
 
 // NodeController is a fake nodes implementation that can be used to test
@@ -57,13 +57,13 @@ type NodeController struct {
 	manageNodesWithLabelSelector          string
 	nodeSelectorFunc                      func(node *corev1.Node) bool
 	lockPodsOnNodeFunc                    func(ctx context.Context, nodeName string) error
-	nodesSets                             *stringSets
+	nodesSets                             maps.SyncMap[string, *NodeInfo]
 	renderer                              *renderer
 	nodeChan                              chan *corev1.Node
 	parallelTasks                         *parallelTasks
 	lifecycle                             Lifecycle
 	cronjob                               *cron.Cron
-	delayJobsCancels                      sync.Map
+	delayJobsCancels                      maps.SyncMap[string, cron.DoFunc]
 	recorder                              record.EventRecorder
 }
 
@@ -118,7 +118,6 @@ func NewNodeController(conf NodeControllerConfig) (*NodeController, error) {
 		nodeIP:                                conf.NodeIP,
 		nodeName:                              conf.NodeName,
 		nodePort:                              conf.NodePort,
-		nodesSets:                             newStringSets(),
 		cronjob:                               cron.NewCron(),
 		lifecycle:                             lifecycles,
 		parallelTasks:                         newParallelTasks(conf.LockNodeParallelism),
@@ -240,7 +239,7 @@ func (c *NodeController) WatchNodes(ctx context.Context, ch chan<- *corev1.Node,
 					}
 				case watch.Deleted:
 					node := event.Object.(*corev1.Node)
-					if c.nodesSets.Has(node.Name) {
+					if _, has := c.nodesSets.Load(node.Name); has {
 						c.nodesSets.Delete(node.Name)
 					}
 				}
@@ -419,17 +418,17 @@ func (c *NodeController) LockNode(ctx context.Context, node *corev1.Node) error
 	)
 	cancelFunc, ok := c.cronjob.AddWithCancel(cron.Order(time.Now().Add(delay)), func() {
 		c.parallelTasks.Add(func() {
-			old, ok := c.delayJobsCancels.LoadAndDelete(key)
+			cancelOld, ok := c.delayJobsCancels.LoadAndDelete(key)
 			if ok {
-				old.(cron.DoFunc)()
+				cancelOld()
 			}
 			do()
 		})
 	})
 	if ok {
-		old, ok := c.delayJobsCancels.LoadOrStore(key, cancelFunc)
+		cancelOld, ok := c.delayJobsCancels.LoadOrStore(key, cancelFunc)
 		if ok {
-			old.(cron.DoFunc)()
+			cancelOld()
 		}
 	}
 	return nil
@@ -526,15 +525,16 @@ func (c *NodeController) putNodeInfo(node *corev1.Node) {
 			}
 		}
 	}
-	nodeInfo := NodeInfo{
+	nodeInfo := &NodeInfo{
 		HostIPs: nodeIPs,
 	}
-	c.nodesSets.Put(node.Name, nodeInfo)
+	c.nodesSets.Store(node.Name, nodeInfo)
 }
 
 // Has returns true if the node is existed
 func (c *NodeController) Has(nodeName string) bool {
-	return c.nodesSets.Has(nodeName)
+	_, has := c.nodesSets.Load(nodeName)
+	return has
 }
 
 // Size returns the number of nodes
@@ -543,10 +543,10 @@ func (c *NodeController) Size() int {
 }
 
 // Get returns Has bool and corev1.Node if the node is existed
-func (c *NodeController) Get(nodeName string) (NodeInfo, bool) {
-	value, has := c.nodesSets.Get(nodeName)
-	if nodeInfo, converted := value.(NodeInfo); has && converted {
+func (c *NodeController) Get(nodeName string) (*NodeInfo, bool) {
+	nodeInfo, has := c.nodesSets.Load(nodeName)
+	if has {
 		return nodeInfo, has
 	}
-	return NodeInfo{}, has
+	return nil, has
 }
diff --git a/pkg/kwok/controllers/pod_controller.go b/pkg/kwok/controllers/pod_controller.go
index c24e2de3..87a62421 100644
--- a/pkg/kwok/controllers/pod_controller.go
+++ b/pkg/kwok/controllers/pod_controller.go
@@ -22,7 +22,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
-	"sync"
 	"text/template"
 	"time"
 
@@ -44,6 +43,7 @@ import (
 	"sigs.k8s.io/kwok/pkg/kwok/cni"
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/expression"
+	"sigs.k8s.io/kwok/pkg/utils/maps"
 )
 
 var (
@@ -59,15 +59,14 @@ type PodController struct {
 	disregardStatusWithLabelSelector      labels.Selector
 	nodeIP                                string
 	cidrIPNet                             *net.IPNet
-	nodeHasFunc                           func(nodeName string) bool
-	nodeGetFunc                           func(nodeName string) (NodeInfo, bool)
+	nodeGetFunc                           func(nodeName string) (*NodeInfo, bool)
 	ipPool                                *ipPool
 	renderer                              *renderer
 	lockPodChan                           chan *corev1.Pod
 	parallelTasks                         *parallelTasks
 	lifecycle                             Lifecycle
 	cronjob                               *cron.Cron
-	delayJobsCancels                      sync.Map
+	delayJobsCancels                      maps.SyncMap[string, cron.DoFunc]
 	recorder                              record.EventRecorder
 }
 
@@ -79,8 +78,7 @@ type PodControllerConfig struct {
 	DisregardStatusWithLabelSelector      string
 	NodeIP                                string
 	CIDR                                  string
-	NodeHasFunc                           func(nodeName string) bool
-	NodeGetFunc                           func(nodeName string) (NodeInfo, bool)
+	NodeGetFunc                           func(nodeName string) (*NodeInfo, bool)
 	Stages                                []*internalversion.Stage
 	LockPodParallelism                    int
 	FuncMap                               template.FuncMap
@@ -117,7 +115,6 @@ func NewPodController(conf PodControllerConfig) (*PodController, error) {
 		nodeIP:                                conf.NodeIP,
 		cidrIPNet:                             cidrIPNet,
 		ipPool:                                newIPPool(cidrIPNet),
-		nodeHasFunc:                           conf.NodeHasFunc,
 		nodeGetFunc:                           conf.NodeGetFunc,
 		cronjob:                               cron.NewCron(),
 		lifecycle:                             lifecycles,
@@ -305,17 +302,17 @@ func (c *PodController) LockPod(ctx context.Context, pod *corev1.Pod) error {
 	)
 	cancelFunc, ok := c.cronjob.AddWithCancel(cron.Order(time.Now().Add(delay)), func() {
 		c.parallelTasks.Add(func() {
-			old, ok := c.delayJobsCancels.LoadAndDelete(key)
+			cancelOld, ok := c.delayJobsCancels.LoadAndDelete(key)
 			if ok {
-				old.(cron.DoFunc)()
+				cancelOld()
 			}
 			do()
 		})
 	})
 	if ok {
-		old, ok := c.delayJobsCancels.LoadOrStore(key, cancelFunc)
+		cancelOld, ok := c.delayJobsCancels.LoadOrStore(key, cancelFunc)
 		if ok {
-			old.(cron.DoFunc)()
+			cancelOld()
 		}
 	}
 
@@ -361,7 +358,7 @@ func (c *PodController) LockPods(ctx context.Context, pods <-chan *corev1.Pod) {
 }
 
 func (c *PodController) needLockPod(pod *corev1.Pod) bool {
-	if !c.nodeHasFunc(pod.Spec.NodeName) {
+	if _, has := c.nodeGetFunc(pod.Spec.NodeName); !has {
 		return false
 	}
 
@@ -426,16 +423,18 @@ func (c *PodController) WatchPods(ctx context.Context, lockChan chan<- *corev1.P
 
 				case watch.Deleted:
 					pod := event.Object.(*corev1.Pod)
-					if c.nodeHasFunc(pod.Spec.NodeName) && !pod.Spec.HostNetwork {
-						if !c.enableCNI {
-							// Recycling PodIP
-							if pod.Status.PodIP != "" && c.cidrIPNet.Contains(net.ParseIP(pod.Status.PodIP)) {
-								c.ipPool.Put(pod.Status.PodIP)
-							}
-						} else {
-							err := cni.Remove(context.Background(), string(pod.UID), pod.Name, pod.Namespace)
-							if err != nil {
-								logger.Error("cni remove", err)
+					// Recycling PodIP
+					if !pod.Spec.HostNetwork {
+						if _, has := c.nodeGetFunc(pod.Spec.NodeName); has {
+							if !c.enableCNI {
+								if pod.Status.PodIP != "" && c.cidrIPNet.Contains(net.ParseIP(pod.Status.PodIP)) {
+									c.ipPool.Put(pod.Status.PodIP)
+								}
+							} else {
+								err := cni.Remove(context.Background(), string(pod.UID), pod.Name, pod.Namespace)
+								if err != nil {
+									logger.Error("cni remove", err)
+								}
 							}
 						}
 					}
diff --git a/pkg/kwok/controllers/renderer.go b/pkg/kwok/controllers/renderer.go
index 8f6faf1a..74086521 100644
--- a/pkg/kwok/controllers/renderer.go
+++ b/pkg/kwok/controllers/renderer.go
@@ -25,10 +25,12 @@ import (
 	"text/template"
 
 	"sigs.k8s.io/yaml"
+
+	"sigs.k8s.io/kwok/pkg/utils/maps"
 )
 
 type renderer struct {
-	cache      sync.Map
+	cache      maps.SyncMap[string, *template.Template]
 	bufferPool sync.Pool
 	funcMap    template.FuncMap
 }
@@ -36,7 +38,6 @@ type renderer struct {
 func newRenderer(funcMap template.FuncMap) *renderer {
 	return &renderer{
 		funcMap: funcMap,
-		cache:   sync.Map{},
 		bufferPool: sync.Pool{
 			New: func() interface{} {
 				return bytes.NewBuffer(make([]byte, 4*1024))
@@ -48,16 +49,15 @@ func newRenderer(funcMap template.FuncMap) *renderer {
 // renderToJSON renders the template with the given text and original object.
 func (r *renderer) renderToJSON(text string, original interface{}) ([]byte, error) {
 	text = strings.TrimSpace(text)
-	v, ok := r.cache.Load(text)
+	temp, ok := r.cache.Load(text)
 	if !ok {
-		temp, err := template.New("_").Funcs(r.funcMap).Parse(text)
+		var err error
+		temp, err = template.New("_").Funcs(r.funcMap).Parse(text)
 		if err != nil {
 			return nil, err
 		}
 		r.cache.Store(text, temp)
-		v = temp
 	}
-	temp := v.(*template.Template)
 	buf := r.bufferPool.Get().(*bytes.Buffer)
 	defer r.bufferPool.Put(buf)
 
diff --git a/pkg/kwok/controllers/utils.go b/pkg/kwok/controllers/utils.go
index 3b6ff82e..d696e010 100644
--- a/pkg/kwok/controllers/utils.go
+++ b/pkg/kwok/controllers/utils.go
@@ -160,57 +160,6 @@ func (p *parallelTasks) Wait() {
 	p.wg.Wait()
 }
 
-type stringSets struct {
-	mut  sync.RWMutex
-	sets map[string]any
-}
-
-func newStringSets() *stringSets {
-	return &stringSets{
-		sets: make(map[string]any),
-	}
-}
-
-func (s *stringSets) Size() int {
-	s.mut.RLock()
-	defer s.mut.RUnlock()
-	return len(s.sets)
-}
-
-func (s *stringSets) Put(key string, value any) {
-	s.mut.Lock()
-	defer s.mut.Unlock()
-	s.sets[key] = value
-}
-
-func (s *stringSets) Delete(key string) {
-	s.mut.Lock()
-	defer s.mut.Unlock()
-	delete(s.sets, key)
-}
-
-func (s *stringSets) Has(key string) bool {
-	s.mut.RLock()
-	defer s.mut.RUnlock()
-	_, ok := s.sets[key]
-	return ok
-}
-
-func (s *stringSets) Foreach(f func(string)) {
-	s.mut.RLock()
-	defer s.mut.RUnlock()
-	for k := range s.sets {
-		f(k)
-	}
-}
-
-func (s *stringSets) Get(key string) (any, bool) {
-	s.mut.RLock()
-	defer s.mut.RUnlock()
-	value, ok := s.sets[key]
-	return value, ok
-}
-
 func labelsParse(selector string) (labels.Selector, error) {
 	if selector == "" {
 		return nil, nil
diff --git a/pkg/kwokctl/pki/pkiutil.go b/pkg/kwokctl/pki/pkiutil.go
index 52cf0dca..3e454a78 100644
--- a/pkg/kwokctl/pki/pkiutil.go
+++ b/pkg/kwokctl/pki/pkiutil.go
@@ -238,7 +238,7 @@ func writeKey(pkiPath, name string, key crypto.Signer) error {
 	return nil
 }
 
-// encodeCertPEM returns PEM-encoded certificate data
+// encodeCertPEM returns PEM-endcoded certificate data
 func encodeCertPEM(cert *x509.Certificate) []byte {
 	block := pem.Block{
 		Type:  CertificateBlockType,
diff --git a/pkg/utils/maps/doc.go b/pkg/utils/maps/doc.go
new file mode 100644
index 00000000..9024b43f
--- /dev/null
+++ b/pkg/utils/maps/doc.go
@@ -0,0 +1,18 @@
+/*
+Copyright 2023 The Kubernetes 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 maps provides utilities to manipulate maps.
+package maps
diff --git a/pkg/utils/maps/sync.go b/pkg/utils/maps/sync.go
new file mode 100644
index 00000000..56be819f
--- /dev/null
+++ b/pkg/utils/maps/sync.go
@@ -0,0 +1,81 @@
+/*
+Copyright 2023 The Kubernetes 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 maps
+
+import (
+	"sync"
+)
+
+// SyncMap is a wrapper around sync.Map that provides a few additional methods.
+type SyncMap[K comparable, V any] struct {
+	m sync.Map
+}
+
+// Load returns the value stored in the map for a key,
+// or nil if no  value is present.
+func (m *SyncMap[K, V]) Load(key K) (value V, ok bool) {
+	v, ok := m.m.Load(key)
+	if !ok {
+		return value, false
+	}
+	return v.(V), true
+}
+
+// Store sets the value for a key.
+func (m *SyncMap[K, V]) Store(key K, value V) {
+	m.m.Store(key, value)
+}
+
+// Delete deletes the value for a key.
+func (m *SyncMap[K, V]) Delete(key K) {
+	m.m.Delete(key)
+}
+
+// Range calls f sequentially for each key and value present in the map.
+func (m *SyncMap[K, V]) Range(f func(key K, value V) bool) {
+	m.m.Range(func(key, value interface{}) bool {
+		return f(key.(K), value.(V))
+	})
+}
+
+// LoadAndDelete deletes the value for a key, returning the previous value if any.
+func (m *SyncMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
+	v, loaded := m.m.LoadAndDelete(key)
+	if !loaded {
+		return value, loaded
+	}
+	return v.(V), loaded
+}
+
+// LoadOrStore returns the existing value for the key if present.
+func (m *SyncMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
+	v, loaded := m.m.LoadOrStore(key, value)
+	if loaded {
+		return v.(V), loaded
+	}
+	return value, loaded
+}
+
+// Size returns the number of items in the map.
+func (m *SyncMap[K, V]) Size() int {
+	size := 0
+	m.m.Range(func(key, value interface{}) bool {
+		size++
+		return true
+	})
+	return size
+}


Test patch:
diff --git a/pkg/kwok/controllers/pod_controller_test.go b/pkg/kwok/controllers/pod_controller_test.go
index 3d14bea0..0a99ff5c 100644
--- a/pkg/kwok/controllers/pod_controller_test.go
+++ b/pkg/kwok/controllers/pod_controller_test.go
@@ -94,8 +94,8 @@ func TestPodController(t *testing.T) {
 	nodeHasFunc := func(nodeName string) bool {
 		return strings.HasPrefix(nodeName, "node")
 	}
-	nodeGetFunc := func(nodeName string) (NodeInfo, bool) {
-		node := NodeInfo{HostIPs: []string{defaultNodeIP}}
+	nodeGetFunc := func(nodeName string) (*NodeInfo, bool) {
+		node := &NodeInfo{HostIPs: []string{defaultNodeIP}}
 		return node, nodeHasFunc(nodeName)
 	}
 	podStageStatus, _ := NewStagesFromYaml([]byte(stages.DefaultPodStages))
@@ -106,7 +106,6 @@ func TestPodController(t *testing.T) {
 		CIDR:                                  "10.0.0.1/24",
 		DisregardStatusWithAnnotationSelector: annotationSelector.String(),
 		Stages:                                podStageStatus,
-		NodeHasFunc:                           nodeHasFunc,
 		NodeGetFunc:                           nodeGetFunc,
 		FuncMap:                               defaultFuncMap,
 		LockPodParallelism:                    2,