github.com/zly-app/zapp@v1.3.3/component/conn/conn.go (about)

     1  /*
     2  -------------------------------------------------
     3     Author :       zlyuancn
     4     date：         2021/1/1
     5     Description :
     6  -------------------------------------------------
     7  */
     8  
     9  package conn
    10  
    11  import (
    12  	"sync"
    13  
    14  	"go.uber.org/zap"
    15  
    16  	"github.com/zly-app/zapp/consts"
    17  	"github.com/zly-app/zapp/logger"
    18  	"github.com/zly-app/zapp/pkg/utils"
    19  )
    20  
    21  type CreatorFunc = func(name string) (IInstance, error)
    22  
    23  type IInstance interface {
    24  	Close()
    25  }
    26  
    27  // 连接器
    28  type Conn struct {
    29  	wgs map[string]*connWaitGroup
    30  	mx  sync.RWMutex
    31  }
    32  
    33  type connWaitGroup struct {
    34  	instance IInstance
    35  	e        error
    36  	wg       sync.WaitGroup
    37  }
    38  
    39  func NewConn() *Conn {
    40  	return &Conn{
    41  		wgs: make(map[string]*connWaitGroup),
    42  	}
    43  }
    44  
    45  // 获取实例
    46  func (c *Conn) GetInstance(creator CreatorFunc, name ...string) IInstance {
    47  	if len(name) == 0 {
    48  		return c.getInstance(creator, consts.DefaultComponentName)
    49  	}
    50  	return c.getInstance(creator, name[0])
    51  }
    52  
    53  func (c *Conn) getInstance(creator CreatorFunc, name string) IInstance {
    54  	c.mx.RLock()
    55  	wg, ok := c.wgs[name]
    56  	c.mx.RUnlock()
    57  
    58  	if ok {
    59  		wg.wg.Wait()
    60  		if wg.e != nil {
    61  			logger.Log.Panic(wg.e.Error(), zap.String("name", name))
    62  		}
    63  		return wg.instance
    64  	}
    65  
    66  	c.mx.Lock()
    67  
    68  	// 再获取一次, 它可能在获取锁的过程中完成了
    69  	if wg, ok = c.wgs[name]; ok {
    70  		c.mx.Unlock()
    71  		wg.wg.Wait()
    72  		if wg.e != nil {
    73  			logger.Log.Panic(wg.e.Error(), zap.String("name", name))
    74  		}
    75  		return wg.instance
    76  	}
    77  
    78  	// 占位置
    79  	wg = new(connWaitGroup)
    80  	wg.wg.Add(1)
    81  	c.wgs[name] = wg
    82  	c.mx.Unlock()
    83  
    84  	var err error
    85  	err = utils.Recover.WrapCall(func() error {
    86  		wg.instance, err = creator(name)
    87  		return err
    88  	})
    89  
    90  	// 如果出现错误, 删除占位
    91  	if err != nil {
    92  		wg.e = err
    93  		wg.wg.Done()
    94  		c.mx.Lock()
    95  		delete(c.wgs, name)
    96  		c.mx.Unlock()
    97  		logger.Log.Panic(utils.Recover.GetRecoverErrorDetail(wg.e), zap.String("name", name))
    98  	}
    99  
   100  	wg.wg.Done()
   101  	return wg.instance
   102  }
   103  
   104  // 移除实例, 移除时会关闭示例
   105  func (c *Conn) Remove(name string) {
   106  	c.mx.Lock()
   107  	defer c.mx.Unlock()
   108  	wg, ok := c.wgs[name]
   109  	if ok {
   110  		wg.instance.Close()
   111  		delete(c.wgs, name)
   112  	}
   113  }
   114  
   115  // 关闭所有实例的链接
   116  func (c *Conn) CloseAll() {
   117  	c.mx.Lock()
   118  	defer c.mx.Unlock()
   119  
   120  	for _, wg := range c.wgs {
   121  		if wg.instance != nil {
   122  			wg.instance.Close()
   123  		}
   124  	}
   125  	c.wgs = make(map[string]*connWaitGroup)
   126  }