syncx: Map 支持 LoadOrStoreFunc 方法 (#194)

.CHANGELOG.md

@@ -6,6 +6,7 @@
 - [mapx: 修复红黑树删除节点问题](https://github.com/ecodeclub/ekit/pull/183)
 - [sqlx: 构建Scanner抽象替代现有ScanRows及ScanAll](https://github.com/ecodeclub/ekit/pull/182)
 - [pool: 重构TaskPool](https://github.com/ecodeclub/ekit/pull/184)
+- [syncx:Map 支持 LoadOrStoreFunc 方法](https://github.com/ecodeclub/ekit/pull/194)
 - [mapx: MutipleTreeMap](https://github.com/ecodeclub/ekit/pull/187)
 - [mapx: 为 MultipleMap 添加 PutVals 方法](https://github.com/ecodeclub/ekit/pull/189)
 - [mapx: LinkedMap 特性](https://github.com/ecodeclub/ekit/pull/191)

syncx/map.go

@@ -41,6 +41,7 @@ func (m *Map[K, V]) Store(key K, value V) {
 }
 
 // LoadOrStore 加载或者存储一个键值对
+// true 代表是加载的，false 代表执行了 store
 func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 	var anyVal any
 	anyVal, loaded = m.m.LoadOrStore(key, value)
@@ -50,6 +51,27 @@ func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 	return
 }
 
+// LoadOrStoreFunc 是一个优化，也就是使用该方法能够避免无意义的创建实例。
+// 如果你的初始化过程非常消耗资源，那么使用这个方法是有价值的。
+// 它的代价就是 Key 不存在的时候会多一次 Load 调用。
+// 当 fn 返回 error 的时候，LoadOrStoreFunc 也会返回 error。
+func (m *Map[K, V]) LoadOrStoreFunc(key K, fn func() (V, error)) (actual V, loaded bool, err error) {
+	var anyVal any
+	val, ok := m.Load(key)
+	if ok {
+		return val, true, nil
+	}
+	val, err = fn()
+	if err != nil {
+		return
+	}
+	anyVal, loaded = m.m.LoadOrStore(key, val)
+	if anyVal != nil {
+		actual = anyVal.(V)
+	}
+	return
+}
+
 // LoadAndDelete 加载并且删除一个键值对
 func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
 	var anyVal any

syncx/map_test.go

@@ -15,6 +15,7 @@
 package syncx
 
 import (
+	"errors"
 	"fmt"
 	"sync"
 	"testing"
@@ -87,6 +88,32 @@ func TestMap_LoadOrStore(t *testing.T) {
 	assert.Nil(t, val)
 }
 
+func TestMap_LoadOrStoreFunc(t *testing.T) {
+	var m = Map[string, *User]{}
+	val, loaded, err := m.LoadOrStoreFunc("Tom", func() (*User, error) {
+		return &User{Name: "Tom"}, nil
+	})
+	assert.NoError(t, err)
+	assert.False(t, loaded)
+	assert.Equal(t, &User{Name: "Tom"}, val)
+
+	// 测试 Tom 存在的情况
+	val, loaded, err = m.LoadOrStoreFunc("Tom", func() (*User, error) {
+		return &User{Name: "Tom"}, nil
+	})
+	assert.NoError(t, err)
+	assert.True(t, loaded)
+	assert.Equal(t, &User{Name: "Tom"}, val)
+
+	// 测试初始化失败
+	val, loaded, err = m.LoadOrStoreFunc("Jerry", func() (*User, error) {
+		return nil, errors.New("初始话失败")
+	})
+	assert.Equal(t, err, errors.New("初始话失败"))
+	assert.False(t, loaded)
+	assert.Equal(t, (*User)(nil), val)
+}
+
 func TestMap_LoadAndDelete(t *testing.T) {
 	var m = Map[string, *User]{}
 	m.Store("Tom", nil)
@@ -204,6 +231,54 @@ func ExampleMap_LoadOrStore() {
 	// 加载旧值 <nil>
 }
 
+func ExampleMap_LoadOrStoreFunc() {
+	var m = Map[string, *User]{}
+	_, loaded, _ := m.LoadOrStoreFunc("Tom", func() (*User, error) {
+		return &User{Name: "Tom"}, nil
+	})
+	// 执行存储
+	if !loaded {
+		fmt.Println("设置了新值 Tom")
+	}
+
+	_, loaded, _ = m.LoadOrStoreFunc("Tom", func() (*User, error) {
+		return &User{Name: "Tom-copy"}, nil
+	})
+	// Tom 这个 key 已经存在，执行加载
+	if loaded {
+		fmt.Println("加载旧值 Tom")
+	}
+
+	_, loaded, _ = m.LoadOrStoreFunc("Jerry", func() (*User, error) {
+		return nil, nil
+	})
+	// 执行存储，注意值是 nil
+	if !loaded {
+		fmt.Println("设置了新值 nil")
+	}
+	val, loaded, _ := m.LoadOrStoreFunc("Jerry", func() (*User, error) {
+		return &User{Name: "Jerry"}, nil
+	})
+	// Jerry 这个 key 已经存在，执行加载，于是把原本的 nil 加载出来
+	if loaded {
+		fmt.Printf("加载旧值 %v\n", val)
+	}
+
+	_, _, err := m.LoadOrStoreFunc("Kitty", func() (*User, error) {
+		return nil, errors.New("初始化失败")
+	})
+	if err != nil {
+		fmt.Println(err.Error())
+	}
+
+	// Output:
+	// 设置了新值 Tom
+	// 加载旧值 Tom
+	// 设置了新值 nil
+	// 加载旧值 <nil>
+	// 初始化失败
+}
+
 func ExampleMap_Range() {
 	var m Map[string, int]
 	m.Store("Tom", 18)