github.com/lingyao2333/mo-zero@v1.4.1/core/syncx/resourcemanager.go (about)

     1  package syncx
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/errorx"
     8  )
     9  
    10  // A ResourceManager is a manager that used to manage resources.
    11  type ResourceManager struct {
    12  	resources    map[string]io.Closer
    13  	singleFlight SingleFlight
    14  	lock         sync.RWMutex
    15  }
    16  
    17  // NewResourceManager returns a ResourceManager.
    18  func NewResourceManager() *ResourceManager {
    19  	return &ResourceManager{
    20  		resources:    make(map[string]io.Closer),
    21  		singleFlight: NewSingleFlight(),
    22  	}
    23  }
    24  
    25  // Close closes the manager.
    26  // Don't use the ResourceManager after Close() called.
    27  func (manager *ResourceManager) Close() error {
    28  	manager.lock.Lock()
    29  	defer manager.lock.Unlock()
    30  
    31  	var be errorx.BatchError
    32  	for _, resource := range manager.resources {
    33  		if err := resource.Close(); err != nil {
    34  			be.Add(err)
    35  		}
    36  	}
    37  
    38  	// release resources to avoid using it later
    39  	manager.resources = nil
    40  
    41  	return be.Err()
    42  }
    43  
    44  // GetResource returns the resource associated with given key.
    45  func (manager *ResourceManager) GetResource(key string, create func() (io.Closer, error)) (io.Closer, error) {
    46  	val, err := manager.singleFlight.Do(key, func() (interface{}, error) {
    47  		manager.lock.RLock()
    48  		resource, ok := manager.resources[key]
    49  		manager.lock.RUnlock()
    50  		if ok {
    51  			return resource, nil
    52  		}
    53  
    54  		resource, err := create()
    55  		if err != nil {
    56  			return nil, err
    57  		}
    58  
    59  		manager.lock.Lock()
    60  		defer manager.lock.Unlock()
    61  		manager.resources[key] = resource
    62  
    63  		return resource, nil
    64  	})
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	return val.(io.Closer), nil
    70  }
    71  
    72  // Inject injects the resource associated with given key.
    73  func (manager *ResourceManager) Inject(key string, resource io.Closer) {
    74  	manager.lock.Lock()
    75  	manager.resources[key] = resource
    76  	manager.lock.Unlock()
    77  }