github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/util/pool/default.go (about)

     1  // Copyright 2020 Asim Aslam
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // Original source: github.com/micro/go-micro/v3/util/pool/default.go
    16  
    17  package pool
    18  
    19  import (
    20  	"sync"
    21  	"time"
    22  
    23  	"github.com/google/uuid"
    24  	"github.com/tickoalcantara12/micro/v3/service/network/transport"
    25  )
    26  
    27  type pool struct {
    28  	size int
    29  	ttl  time.Duration
    30  	tr   transport.Transport
    31  
    32  	sync.Mutex
    33  	conns map[string][]*poolConn
    34  }
    35  
    36  type poolConn struct {
    37  	transport.Client
    38  	id      string
    39  	created time.Time
    40  }
    41  
    42  func newPool(options Options) *pool {
    43  	return &pool{
    44  		size:  options.Size,
    45  		tr:    options.Transport,
    46  		ttl:   options.TTL,
    47  		conns: make(map[string][]*poolConn),
    48  	}
    49  }
    50  
    51  func (p *pool) Close() error {
    52  	p.Lock()
    53  	for k, c := range p.conns {
    54  		for _, conn := range c {
    55  			conn.Client.Close()
    56  		}
    57  		delete(p.conns, k)
    58  	}
    59  	p.Unlock()
    60  	return nil
    61  }
    62  
    63  // NoOp the Close since we manage it
    64  func (p *poolConn) Close() error {
    65  	return nil
    66  }
    67  
    68  func (p *poolConn) Id() string {
    69  	return p.id
    70  }
    71  
    72  func (p *poolConn) Created() time.Time {
    73  	return p.created
    74  }
    75  
    76  func (p *pool) Get(addr string, opts ...transport.DialOption) (Conn, error) {
    77  	p.Lock()
    78  	conns := p.conns[addr]
    79  
    80  	// while we have conns check age and then return one
    81  	// otherwise we'll create a new conn
    82  	for len(conns) > 0 {
    83  		conn := conns[len(conns)-1]
    84  		conns = conns[:len(conns)-1]
    85  		p.conns[addr] = conns
    86  
    87  		// if conn is old kill it and move on
    88  		if d := time.Since(conn.Created()); d > p.ttl {
    89  			conn.Client.Close()
    90  			continue
    91  		}
    92  
    93  		// we got a good conn, lets unlock and return it
    94  		p.Unlock()
    95  
    96  		return conn, nil
    97  	}
    98  
    99  	p.Unlock()
   100  
   101  	// create new conn
   102  	c, err := p.tr.Dial(addr, opts...)
   103  	if err != nil {
   104  		return nil, err
   105  	}
   106  	return &poolConn{
   107  		Client:  c,
   108  		id:      uuid.New().String(),
   109  		created: time.Now(),
   110  	}, nil
   111  }
   112  
   113  func (p *pool) Release(conn Conn, err error) error {
   114  	// don't store the conn if it has errored
   115  	if err != nil {
   116  		return conn.(*poolConn).Client.Close()
   117  	}
   118  
   119  	// otherwise put it back for reuse
   120  	p.Lock()
   121  	conns := p.conns[conn.Remote()]
   122  	if len(conns) >= p.size {
   123  		p.Unlock()
   124  		return conn.(*poolConn).Client.Close()
   125  	}
   126  	p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
   127  	p.Unlock()
   128  
   129  	return nil
   130  }