Repository: dragonflyoss/Dragonfly2
Pull request #934 (2021-12-15 09:32:18)

feat: add kmutex and krwmutex
Signed-off-by: Gaius <gaius.qi@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description
- Add `kmutex` and `krwmutex`
<!--- Describe your changes in detail -->

## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate):

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [x] Code compiles correctly.
- [x] My code follows the code style of this project.
- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [x] I have read the **CONTRIBUTING** document.
- [x] I have added tests to cover my changes.
- [x] All new and existing tests passed.


codecov-commenter: # [Codecov](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) Report
> Merging [#934](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (d9ff997) into [main](https://codecov.io/gh/dragonflyoss/Dragonfly2/commit/cf2c5b13eb00c24fa8af8a3daa94eb0db8c828ea?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (cf2c5b1) will **increase** coverage by `0.08%`.
> The diff coverage is `82.35%`.

[![Impacted file tree graph](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934/graphs/tree.svg?width=650&height=150&src=pr&token=PUICJYSAT7&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)

```diff
@@            Coverage Diff             @@
##             main     #934      +/-   ##
==========================================
+ Coverage   42.66%   42.74%   +0.08%     
==========================================
  Files         113      115       +2     
  Lines        8605     8639      +34     
==========================================
+ Hits         3671     3693      +22     
- Misses       4468     4477       +9     
- Partials      466      469       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| Object-compatibility-e2etests | `?` | |
| e2etests | `?` | |
| unittests | `42.74% <82.35%> (+0.08%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) | Coverage Δ | |
|---|---|---|
| [pkg/sync/krwmutex.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-cGtnL3N5bmMva3J3bXV0ZXguZ28=) | `81.81% <81.81%> (ø)` | |
| [pkg/sync/kmutex.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-cGtnL3N5bmMva211dGV4Lmdv) | `83.33% <83.33%> (ø)` | |
| [client/daemon/peer/peertask\_base.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/934/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX2Jhc2UuZ28=) | `45.38% <0.00%> (-1.16%)` | :arrow_down: |


Patch:
diff --git a/pkg/sync/kmutex.go b/pkg/sync/kmutex.go
new file mode 100644
index 00000000..833bc570
--- /dev/null
+++ b/pkg/sync/kmutex.go
@@ -0,0 +1,48 @@
+/*
+ *     Copyright 2020 The Dragonfly 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 sync
+
+import (
+	"sync"
+)
+
+type Kmutex struct {
+	m *sync.Map
+}
+
+func NewKmutex() *Kmutex {
+	m := sync.Map{}
+	return &Kmutex{&m}
+}
+
+func (k *Kmutex) Lock(key interface{}) {
+	m := sync.Mutex{}
+	rm, _ := k.m.LoadOrStore(key, &m)
+	mu := rm.(*sync.Mutex)
+	mu.Lock()
+}
+
+func (k *Kmutex) Unlock(key interface{}) {
+	rm, ok := k.m.Load(key)
+	if !ok {
+		return
+	}
+
+	mu := rm.(*sync.Mutex)
+	k.m.Delete(key)
+	mu.Unlock()
+}
diff --git a/pkg/sync/krwmutex.go b/pkg/sync/krwmutex.go
new file mode 100644
index 00000000..fe3b5743
--- /dev/null
+++ b/pkg/sync/krwmutex.go
@@ -0,0 +1,66 @@
+/*
+ *     Copyright 2020 The Dragonfly 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 sync
+
+import (
+	"sync"
+)
+
+type Krwmutex struct {
+	m *sync.Map
+}
+
+func NewKrwmutex() *Krwmutex {
+	m := sync.Map{}
+	return &Krwmutex{&m}
+}
+
+func (k *Krwmutex) Lock(key interface{}) {
+	m := sync.RWMutex{}
+	rm, _ := k.m.LoadOrStore(key, &m)
+	mu := rm.(*sync.RWMutex)
+	mu.Lock()
+}
+
+func (k *Krwmutex) Unlock(key interface{}) {
+	rm, ok := k.m.Load(key)
+	if !ok {
+		return
+	}
+
+	mu := rm.(*sync.RWMutex)
+	k.m.Delete(key)
+	mu.Unlock()
+}
+
+func (k *Krwmutex) RLock(key interface{}) {
+	m := sync.RWMutex{}
+	rm, _ := k.m.LoadOrStore(key, &m)
+	mu := rm.(*sync.RWMutex)
+	mu.RLock()
+}
+
+func (k *Krwmutex) RUnlock(key interface{}) {
+	rm, ok := k.m.Load(key)
+	if !ok {
+		return
+	}
+
+	mu := rm.(*sync.RWMutex)
+	k.m.Delete(key)
+	mu.RUnlock()
+}
diff --git a/scheduler/core/service.go b/scheduler/core/service.go
index 8f1c6666..048f27d4 100644
--- a/scheduler/core/service.go
+++ b/scheduler/core/service.go
@@ -34,7 +34,7 @@ import (
 	"d7y.io/dragonfly/v2/pkg/rpc/base"
 	"d7y.io/dragonfly/v2/pkg/rpc/base/common"
 	schedulerRPC "d7y.io/dragonfly/v2/pkg/rpc/scheduler"
-	"d7y.io/dragonfly/v2/pkg/synclock"
+	pkgsync "d7y.io/dragonfly/v2/pkg/sync"
 	"d7y.io/dragonfly/v2/scheduler/config"
 	"d7y.io/dragonfly/v2/scheduler/core/scheduler"
 	"d7y.io/dragonfly/v2/scheduler/metrics"
@@ -77,6 +77,7 @@ type SchedulerService struct {
 	monitor *monitor
 	done    chan struct{}
 	wg      sync.WaitGroup
+	kmu     *pkgsync.Krwmutex
 
 	config        *config.SchedulerConfig
 	dynconfig     config.DynconfigInterface
@@ -122,6 +123,8 @@ func NewSchedulerService(cfg *config.SchedulerConfig, pluginDir string, metricsC
 		metricsConfig: metricsConfig,
 		dynconfig:     dynConfig,
 		done:          make(chan struct{}),
+		wg:            sync.WaitGroup{},
+		kmu:           pkgsync.NewKrwmutex(),
 	}
 	if !ops.disableCDN {
 		var opts []grpc.DialOption
@@ -249,24 +252,23 @@ func (s *SchedulerService) RegisterTask(req *schedulerRPC.PeerTaskRequest, task
 func (s *SchedulerService) GetOrAddTask(ctx context.Context, task *supervisor.Task) *supervisor.Task {
 	span := trace.SpanFromContext(ctx)
 
-	synclock.Lock(task.ID, true)
+	s.kmu.RLock(task.ID)
 	task, ok := s.taskManager.GetOrAdd(task)
 	if ok {
 		span.SetAttributes(config.AttributeTaskStatus.String(task.GetStatus().String()))
 		span.SetAttributes(config.AttributeLastTriggerTime.String(task.LastTriggerAt.Load().String()))
 		if task.LastTriggerAt.Load().Add(s.config.AccessWindow).After(time.Now()) || task.IsHealth() {
 			span.SetAttributes(config.AttributeNeedSeedCDN.Bool(false))
-			synclock.UnLock(task.ID, true)
+			s.kmu.RUnlock(task.ID)
 			return task
 		}
 	} else {
 		task.Log().Infof("add new task %s", task.ID)
 	}
+	s.kmu.RUnlock(task.ID)
 
-	synclock.UnLock(task.ID, true)
-
-	synclock.Lock(task.ID, false)
-	defer synclock.UnLock(task.ID, false)
+	s.kmu.Lock(task.ID)
+	defer s.kmu.Unlock(task.ID)
 
 	// do trigger
 	span.SetAttributes(config.AttributeTaskStatus.String(task.GetStatus().String()))


Test patch:
diff --git a/pkg/sync/kmutex_test.go b/pkg/sync/kmutex_test.go
new file mode 100644
index 00000000..6154fce3
--- /dev/null
+++ b/pkg/sync/kmutex_test.go
@@ -0,0 +1,87 @@
+/*
+ *     Copyright 2020 The Dragonfly 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 sync
+
+import (
+	"testing"
+	"time"
+)
+
+func TestKmutex_Lock(t *testing.T) {
+	k := NewKmutex()
+	done := time.After(time.Second)
+
+	k.Lock("foo")
+	foo := make(chan struct{})
+	go func() {
+		k.Lock("foo")
+		close(foo)
+	}()
+
+	k.Lock("bar")
+	bar := make(chan struct{})
+	go func() {
+		k.Lock("bar")
+		close(bar)
+	}()
+
+	k.Lock("baz")
+	baz := make(chan struct{})
+	go func() {
+		k.Lock("baz")
+		close(baz)
+	}()
+
+	select {
+	case <-bar:
+	case <-baz:
+	case <-foo:
+	case <-done:
+		t.Fatal("kmutex is not unlocked")
+	default:
+	}
+
+	k.Unlock("foo")
+	select {
+	case <-bar:
+	case <-baz:
+	case <-done:
+		t.Fatal("bar and baz are not unlocked")
+	case <-foo:
+		k.Unlock("foo")
+	}
+
+	k.Unlock("bar")
+	select {
+	case <-foo:
+	case <-baz:
+	case <-done:
+		t.Fatal("foo and baz are not unlocked")
+	case <-bar:
+		k.Unlock("bar")
+	}
+
+	k.Unlock("baz")
+	select {
+	case <-foo:
+	case <-bar:
+	case <-done:
+		t.Fatal("foo and bar are not unlocked")
+	case <-baz:
+		k.Unlock("baz")
+	}
+}
diff --git a/pkg/sync/krwmutex_test.go b/pkg/sync/krwmutex_test.go
new file mode 100644
index 00000000..615f51bf
--- /dev/null
+++ b/pkg/sync/krwmutex_test.go
@@ -0,0 +1,152 @@
+/*
+ *     Copyright 2020 The Dragonfly 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 sync
+
+import (
+	"testing"
+	"time"
+)
+
+func TestKrwmutex_Lock(t *testing.T) {
+	k := NewKrwmutex()
+	done := time.After(time.Second)
+
+	k.Lock("foo")
+	foo := make(chan struct{})
+	go func() {
+		k.Lock("foo")
+		close(foo)
+	}()
+
+	k.Lock("bar")
+	bar := make(chan struct{})
+	go func() {
+		k.Lock("bar")
+		close(bar)
+	}()
+
+	k.Lock("baz")
+	baz := make(chan struct{})
+	go func() {
+		k.Lock("baz")
+		close(baz)
+	}()
+
+	select {
+	case <-bar:
+	case <-baz:
+	case <-foo:
+	case <-done:
+		t.Fatal("kmutex is not unlocked")
+	default:
+	}
+
+	k.Unlock("foo")
+	select {
+	case <-bar:
+	case <-baz:
+	case <-done:
+		t.Fatal("bar and baz are not unlocked")
+	case <-foo:
+		k.Unlock("foo")
+	}
+
+	k.Unlock("bar")
+	select {
+	case <-foo:
+	case <-baz:
+	case <-done:
+		t.Fatal("foo and baz are not unlocked")
+	case <-bar:
+		k.Unlock("bar")
+	}
+
+	k.Unlock("baz")
+	select {
+	case <-foo:
+	case <-bar:
+	case <-done:
+		t.Fatal("foo and bar are not unlocked")
+	case <-baz:
+		k.Unlock("baz")
+	}
+}
+
+func TestKrwmutex_RLock(t *testing.T) {
+	k := NewKrwmutex()
+	done := time.After(time.Second)
+
+	k.Lock("foo")
+	foo := make(chan struct{})
+	go func() {
+		k.RLock("foo")
+		close(foo)
+	}()
+
+	k.Lock("bar")
+	bar := make(chan struct{})
+	go func() {
+		k.RLock("bar")
+		close(bar)
+	}()
+
+	k.Lock("baz")
+	baz := make(chan struct{})
+	go func() {
+		k.RLock("baz")
+		close(baz)
+	}()
+
+	select {
+	case <-bar:
+	case <-baz:
+	case <-foo:
+	case <-done:
+		t.Fatal("kmutex is not unlocked")
+	default:
+	}
+
+	k.Unlock("foo")
+	select {
+	case <-bar:
+	case <-baz:
+	case <-done:
+		t.Fatal("bar and baz are not unlocked")
+	case <-foo:
+		k.RUnlock("foo")
+	}
+
+	k.Unlock("bar")
+	select {
+	case <-foo:
+	case <-baz:
+	case <-done:
+		t.Fatal("foo and baz are not unlocked")
+	case <-bar:
+		k.RUnlock("bar")
+	}
+
+	k.Unlock("baz")
+	select {
+	case <-foo:
+	case <-bar:
+	case <-done:
+		t.Fatal("foo and bar are not unlocked")
+	case <-baz:
+		k.RUnlock("baz")
+	}
+}