github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/transport/pool.go (about)

     1  package transport
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  var (
     8  	poolUint32Data = sync.Pool{
     9  		New: func() interface{} {
    10  			data := make([]byte, 4)
    11  			return &data
    12  		},
    13  	}
    14  
    15  	msgPool = sync.Pool{
    16  		New: func() interface{} {
    17  			return newMessage()
    18  		},
    19  	}
    20  )
    21  
    22  type (
    23  	Pool struct {
    24  		sync.RWMutex
    25  		pool map[string]*Socket
    26  	}
    27  )
    28  
    29  func GetMessageFromPool() *Message {
    30  	return msgPool.Get().(*Message)
    31  }
    32  
    33  func PutMessageToPool(msg *Message) {
    34  	if msg != nil {
    35  		msg.Reset()
    36  		msgPool.Put(msg)
    37  	}
    38  }
    39  
    40  // NewPool returns a new socket pool
    41  func NewPool() *Pool {
    42  	return &Pool{
    43  		pool: make(map[string]*Socket),
    44  	}
    45  }
    46  
    47  func (p *Pool) Get(id string) (*Socket, bool) {
    48  	// attempt to get existing socket
    49  	p.RLock()
    50  	socket, ok := p.pool[id]
    51  	if ok {
    52  		p.RUnlock()
    53  		return socket, ok
    54  	}
    55  	p.RUnlock()
    56  
    57  	// save socket
    58  	p.Lock()
    59  	defer p.Unlock()
    60  	// double checked locking
    61  	socket, ok = p.pool[id]
    62  	if ok {
    63  		return socket, ok
    64  	}
    65  	// create new socket
    66  	socket = New(id)
    67  	p.pool[id] = socket
    68  
    69  	// return socket
    70  	return socket, false
    71  }
    72  
    73  func (p *Pool) Release(s *Socket) {
    74  	p.Lock()
    75  	defer p.Unlock()
    76  
    77  	// close the socket
    78  	s.Close()
    79  	delete(p.pool, s.id)
    80  }
    81  
    82  // Close the pool and delete all the sockets
    83  func (p *Pool) Close() {
    84  	p.Lock()
    85  	defer p.Unlock()
    86  	for id, sock := range p.pool {
    87  		sock.Close()
    88  		delete(p.pool, id)
    89  	}
    90  }