From 5b6a9b324603945b4fc49e9e38492285bea72881 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Wed, 1 May 2024 06:09:36 +0000
Subject: [PATCH] 1

---
 sync/Map.go      | 44 ++++++++++++++++++++++++++++++++++++++------
 sync/Map_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/sync/Map.go b/sync/Map.go
index 734adcf..18ef018 100644
--- a/sync/Map.go
+++ b/sync/Map.go
@@ -81,19 +81,20 @@ type MapExceeded[K, V any] struct {
 type mapExceededItem[V any] struct {
 	data     *V
 	exceeded time.Time
+	wait     sync.RWMutex
 }
 
 func (t *MapExceeded[K, V]) Store(k K, v *V, dur time.Duration) {
-	t.m.Store(k, mapExceededItem[V]{
+	t.m.Store(k, &mapExceededItem[V]{
 		data:     v,
 		exceeded: time.Now().Add(dur),
 	})
 }
 
 func (t *MapExceeded[K, V]) Load(k K) (*V, bool) {
-	if v, ok := t.m.Load(k); ok {
-		if v.(mapExceededItem[V]).exceeded.After(time.Now()) {
-			return v.(mapExceededItem[V]).data, true
+	if v, ok := t.m.LoadV(k).(*mapExceededItem[V]); ok {
+		if v.exceeded.After(time.Now()) {
+			return v.data, true
 		}
 		t.Delete(k)
 	}
@@ -102,7 +103,7 @@ func (t *MapExceeded[K, V]) Load(k K) (*V, bool) {
 
 func (t *MapExceeded[K, V]) Range(f func(key K, value *V) bool) {
 	t.m.Range(func(key, value any) bool {
-		if value.(mapExceededItem[V]).exceeded.After(time.Now()) {
+		if value.(*mapExceededItem[V]).exceeded.After(time.Now()) {
 			return f(key.(K), value.(*V))
 		}
 		t.Delete(key.(K))
@@ -116,7 +117,7 @@ func (t *MapExceeded[K, V]) Len() int {
 
 func (t *MapExceeded[K, V]) GC() {
 	t.m.Range(func(key, value any) bool {
-		if value.(mapExceededItem[V]).exceeded.Before(time.Now()) {
+		if value.(*mapExceededItem[V]).exceeded.Before(time.Now()) {
 			t.Delete(key.(K))
 		}
 		return true
@@ -126,3 +127,34 @@ func (t *MapExceeded[K, V]) GC() {
 func (t *MapExceeded[K, V]) Delete(k K) {
 	t.m.Delete(k)
 }
+
+func (t *MapExceeded[K, V]) LoadOrStore(k K) (vr *V, loaded bool, store func(v1 *V, dur time.Duration)) {
+	store = func(v1 *V, dur time.Duration) {}
+	var actual any
+	actual, loaded = t.m.LoadOrStore(k, &mapExceededItem[V]{})
+	v := actual.(*mapExceededItem[V])
+	v.wait.RLock()
+	exp := v.exceeded
+	vr = v.data
+	v.wait.RUnlock()
+	if loaded && time.Now().Before(exp) {
+		return
+	}
+	if !loaded || (loaded && !exp.IsZero()) {
+		store = func(v1 *V, dur time.Duration) {
+			v.wait.Lock()
+			v.data = v1
+			v.exceeded = time.Now().Add(dur)
+			v.wait.Unlock()
+		}
+		return
+	}
+	for loaded && exp.IsZero() {
+		time.Sleep(time.Millisecond * 20)
+		v.wait.RLock()
+		exp = v.exceeded
+		vr = v.data
+		v.wait.RUnlock()
+	}
+	return
+}
diff --git a/sync/Map_test.go b/sync/Map_test.go
index fa71d0f..ef821ec 100644
--- a/sync/Map_test.go
+++ b/sync/Map_test.go
@@ -346,3 +346,30 @@ func TestMapExceeded1(t *testing.T) {
 		t.Fatal()
 	}
 }
+
+func TestMapExceeded2(t *testing.T) {
+	var m MapExceeded[string, []byte]
+	var data = []byte("1")
+	if v, loaded, f := m.LoadOrStore("1"); v != nil || loaded {
+		t.Fatal()
+	} else {
+		f(&data, time.Second)
+		if v, ok := m.Load("1"); !ok || v == nil || !bytes.Equal(data, *v) {
+			t.Fatal()
+		}
+	}
+
+	var w sync.WaitGroup
+	w.Add(10)
+	for i := 0; i < 10; i++ {
+		go func() {
+			v, loaded, f := m.LoadOrStore("2")
+			if (!loaded && v != nil) || (loaded && !bytes.Equal(data, *v)) {
+				panic("")
+			}
+			f(&data, time.Second)
+			w.Done()
+		}()
+	}
+	w.Wait()
+}
-- 
2.39.2