From 9c45e8a974fd84f5489441d891706e54bd6c1fc9 Mon Sep 17 00:00:00 2001
From: Wei Huang <wei.huang1@ibm.com>
Date: Wed, 5 May 2021 12:00:52 -0700
Subject: [PATCH] sched: make CycleState's Read()/Write()/Delete() thread-safe

- add internal locking to CycleState's Read()/Write()/Delete() functions
- remove Lock() and Unlock() functions
---
 pkg/scheduler/framework/cycle_state.go        | 39 ++++++-------------
 .../plugins/examples/multipoint/multipoint.go |  6 ---
 2 files changed, 11 insertions(+), 34 deletions(-)

diff --git a/pkg/scheduler/framework/cycle_state.go b/pkg/scheduler/framework/cycle_state.go
index fd7f5543d8c..652a257d9a1 100644
--- a/pkg/scheduler/framework/cycle_state.go
+++ b/pkg/scheduler/framework/cycle_state.go
@@ -86,9 +86,10 @@ func (c *CycleState) Clone() *CycleState {
 
 // Read retrieves data with the given "key" from CycleState. If the key is not
 // present an error is returned.
-// This function is not thread safe. In multi-threaded code, lock should be
-// acquired first.
+// This function is thread safe by acquiring an internal lock first.
 func (c *CycleState) Read(key StateKey) (StateData, error) {
+	c.mx.RLock()
+	defer c.mx.RUnlock()
 	if v, ok := c.storage[key]; ok {
 		return v, nil
 	}
@@ -96,35 +97,17 @@ func (c *CycleState) Read(key StateKey) (StateData, error) {
 }
 
 // Write stores the given "val" in CycleState with the given "key".
-// This function is not thread safe. In multi-threaded code, lock should be
-// acquired first.
+// This function is thread safe by acquiring an internal lock first.
 func (c *CycleState) Write(key StateKey, val StateData) {
-	c.storage[key] = val
-}
-
-// Delete deletes data with the given key from CycleState.
-// This function is not thread safe. In multi-threaded code, lock should be
-// acquired first.
-func (c *CycleState) Delete(key StateKey) {
-	delete(c.storage, key)
-}
-
-// Lock acquires CycleState lock.
-func (c *CycleState) Lock() {
 	c.mx.Lock()
-}
-
-// Unlock releases CycleState lock.
-func (c *CycleState) Unlock() {
+	c.storage[key] = val
 	c.mx.Unlock()
 }
 
-// RLock acquires CycleState read lock.
-func (c *CycleState) RLock() {
-	c.mx.RLock()
-}
-
-// RUnlock releases CycleState read lock.
-func (c *CycleState) RUnlock() {
-	c.mx.RUnlock()
+// Delete deletes data with the given key from CycleState.
+// This function is thread safe by acquiring an internal lock first.
+func (c *CycleState) Delete(key StateKey) {
+	c.mx.Lock()
+	delete(c.storage, key)
+	c.mx.Unlock()
 }
diff --git a/pkg/scheduler/framework/plugins/examples/multipoint/multipoint.go b/pkg/scheduler/framework/plugins/examples/multipoint/multipoint.go
index ba483a3afc2..f7e9bf9f4ec 100644
--- a/pkg/scheduler/framework/plugins/examples/multipoint/multipoint.go
+++ b/pkg/scheduler/framework/plugins/examples/multipoint/multipoint.go
@@ -56,9 +56,7 @@ func (mc CommunicatingPlugin) Reserve(ctx context.Context, state *framework.Cycl
 		return framework.NewStatus(framework.Error, "pod cannot be nil")
 	}
 	if pod.Name == "my-test-pod" {
-		state.Lock()
 		state.Write(framework.StateKey(pod.Name), &stateData{data: "never bind"})
-		state.Unlock()
 	}
 	return nil
 }
@@ -67,12 +65,10 @@ func (mc CommunicatingPlugin) Reserve(ctx context.Context, state *framework.Cycl
 // during "reserve" extension point or later.
 func (mc CommunicatingPlugin) Unreserve(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) {
 	if pod.Name == "my-test-pod" {
-		state.Lock()
 		// The pod is at the end of its lifecycle -- let's clean up the allocated
 		// resources. In this case, our clean up is simply deleting the key written
 		// in the Reserve operation.
 		state.Delete(framework.StateKey(pod.Name))
-		state.Unlock()
 	}
 }
 
@@ -81,8 +77,6 @@ func (mc CommunicatingPlugin) PreBind(ctx context.Context, state *framework.Cycl
 	if pod == nil {
 		return framework.NewStatus(framework.Error, "pod cannot be nil")
 	}
-	state.RLock()
-	defer state.RUnlock()
 	if v, e := state.Read(framework.StateKey(pod.Name)); e == nil {
 		if value, ok := v.(*stateData); ok && value.data == "never bind" {
 			return framework.NewStatus(framework.Unschedulable, "pod is not permitted")