resource_set.go

package utils

import (
	"sync"
)

type resourceSet struct {
	sync.Map
	mapLock  sync.Mutex
	locks    map[interface{}]*sync.Mutex
	initFunc func(key interface{}) (interface{}, error)
}

func NewResourceSet(f func(key interface{}) (interface{}, error)) *resourceSet {
	return &resourceSet{
		locks:    make(map[interface{}]*sync.Mutex),
		initFunc: f,
	}
}

func (rs *resourceSet) CreateOrGet(key interface{}) (value interface{}, err error) {
	// 大部分情况会直接取到
	if v, ok := rs.Load(key); ok {
		return v, nil
	}

	// 保证相同的 key 只创建一个锁，不同的 key 使用不同的锁
	// 可以并行地为不同 key 执行初始化函数，互不影响
	rs.mapLock.Lock()
	lock, ok := rs.locks[key]
	if !ok {
		lock = &sync.Mutex{}
		rs.locks[key] = lock
	}
	rs.mapLock.Unlock()

	// 保证相同的 key 只创建一个资源
	lock.Lock()
	defer lock.Unlock()
	// 拿到锁后再判断是否已经创建好了资源
	if v, ok := rs.Load(key); ok {
		return v, nil
	}

	resource, err := rs.initFunc(key)
	if err != nil {
		return nil, err
	}
	rs.Store(key, resource)

	return resource, nil
}