github.com/annwntech/go-micro/v2@v2.9.5/util/socket/pool.go (about)

     1  package socket
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type Pool struct {
     8  	sync.RWMutex
     9  	pool map[string]*Socket
    10  }
    11  
    12  func (p *Pool) Get(id string) (*Socket, bool) {
    13  	// attempt to get existing socket
    14  	p.RLock()
    15  	socket, ok := p.pool[id]
    16  	if ok {
    17  		p.RUnlock()
    18  		return socket, ok
    19  	}
    20  	p.RUnlock()
    21  
    22  	// save socket
    23  	p.Lock()
    24  	defer p.Unlock()
    25  	// double checked locking
    26  	socket, ok = p.pool[id]
    27  	if ok {
    28  		return socket, ok
    29  	}
    30  	// create new socket
    31  	socket = New(id)
    32  	p.pool[id] = socket
    33  
    34  	// return socket
    35  	return socket, false
    36  }
    37  
    38  func (p *Pool) Release(s *Socket) {
    39  	p.Lock()
    40  	defer p.Unlock()
    41  
    42  	// close the socket
    43  	s.Close()
    44  	delete(p.pool, s.id)
    45  }
    46  
    47  // Close the pool and delete all the sockets
    48  func (p *Pool) Close() {
    49  	p.Lock()
    50  	defer p.Unlock()
    51  	for id, sock := range p.pool {
    52  		sock.Close()
    53  		delete(p.pool, id)
    54  	}
    55  }
    56  
    57  // NewPool returns a new socket pool
    58  func NewPool() *Pool {
    59  	return &Pool{
    60  		pool: make(map[string]*Socket),
    61  	}
    62  }