github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/pkg/client/grpc/pool.go (about)

     1  /*
     2  Copyright [2014] - [2023] The Last.Backend authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package grpc
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  	"time"
    23  
    24  	"google.golang.org/grpc"
    25  	"google.golang.org/grpc/connectivity"
    26  )
    27  
    28  type PoolOptions struct {
    29  	Size *int           `env:"POOL_SIZE"  envDefault:"" comment:"Set pool size"`
    30  	TTL  *time.Duration `env:"POOL_TTL"  envDefault:"" comment:"Set pool ttl"`
    31  }
    32  
    33  type pool struct {
    34  	sync.Mutex
    35  
    36  	size       int
    37  	ttl        int64
    38  	maxStreams int
    39  	maxIdle    int
    40  	conns      map[string]*streamsPool
    41  }
    42  
    43  type streamsPool struct {
    44  	head  *poolConn
    45  	busy  *poolConn
    46  	count int
    47  	idle  int
    48  }
    49  
    50  type poolConn struct {
    51  	*grpc.ClientConn
    52  
    53  	err  error
    54  	addr string
    55  
    56  	pool    *pool
    57  	sp      *streamsPool
    58  	streams int
    59  	created int64
    60  
    61  	pre  *poolConn
    62  	next *poolConn
    63  	in   bool
    64  }
    65  
    66  func newPool() *pool {
    67  	return &pool{
    68  		size:       1,
    69  		ttl:        int64(0),
    70  		maxStreams: 1,
    71  		maxIdle:    0,
    72  		conns:      make(map[string]*streamsPool),
    73  	}
    74  }
    75  
    76  func (p *pool) Init(opts PoolOptions) {
    77  	p.Lock()
    78  	defer p.Unlock()
    79  	if opts.Size != nil {
    80  		p.size = *opts.Size
    81  	}
    82  	if opts.TTL != nil {
    83  		p.ttl = int64(opts.TTL.Seconds())
    84  	}
    85  }
    86  
    87  func (p *pool) getConn(ctx context.Context, addr string, opts ...grpc.DialOption) (*poolConn, error) {
    88  	now := time.Now().Unix()
    89  	p.Lock()
    90  	sp, ok := p.conns[addr]
    91  	if !ok {
    92  		sp = &streamsPool{head: &poolConn{}, busy: &poolConn{}, count: 0, idle: 0}
    93  		p.conns[addr] = sp
    94  	}
    95  
    96  	conn := sp.head.next
    97  
    98  	for conn != nil {
    99  		switch conn.GetState() {
   100  		case connectivity.Connecting:
   101  			conn = conn.next
   102  			continue
   103  		case connectivity.Shutdown:
   104  			next := conn.next
   105  			if conn.streams == 0 {
   106  				removeConn(conn)
   107  				sp.idle--
   108  			}
   109  			conn = next
   110  			continue
   111  		case connectivity.TransientFailure:
   112  			next := conn.next
   113  			if conn.streams == 0 {
   114  				removeConn(conn)
   115  				conn.ClientConn.Close()
   116  				sp.idle--
   117  			}
   118  			conn = next
   119  			continue
   120  		case connectivity.Ready:
   121  		case connectivity.Idle:
   122  		}
   123  
   124  		if now-conn.created > p.ttl {
   125  			next := conn.next
   126  			if conn.streams == 0 {
   127  				removeConn(conn)
   128  				conn.ClientConn.Close()
   129  				sp.idle--
   130  			}
   131  			conn = next
   132  			continue
   133  		}
   134  
   135  		if conn.streams >= p.maxStreams {
   136  			next := conn.next
   137  			removeConn(conn)
   138  			addConnAfter(conn, sp.busy)
   139  			conn = next
   140  			continue
   141  		}
   142  
   143  		if conn.streams == 0 {
   144  			sp.idle--
   145  		}
   146  
   147  		conn.streams++
   148  		p.Unlock()
   149  		return conn, nil
   150  	}
   151  
   152  	p.Unlock()
   153  
   154  	cc, err := grpc.DialContext(ctx, addr, opts...)
   155  	if err != nil {
   156  		return nil, err
   157  	}
   158  	conn = &poolConn{cc, nil, addr, p, sp, 1, time.Now().Unix(), nil, nil, false}
   159  
   160  	p.Lock()
   161  	if sp.count < p.size {
   162  		addConnAfter(conn, sp.head)
   163  	}
   164  	p.Unlock()
   165  
   166  	return conn, nil
   167  }
   168  
   169  func (p *pool) release(addr string, conn *poolConn, err error) {
   170  	p.Lock()
   171  	p, sp, created := conn.pool, conn.sp, conn.created
   172  
   173  	if !conn.in && sp.count < p.size {
   174  		addConnAfter(conn, sp.head)
   175  	}
   176  	if !conn.in {
   177  		p.Unlock()
   178  		conn.ClientConn.Close()
   179  		return
   180  	}
   181  
   182  	if conn.streams >= p.maxStreams {
   183  		removeConn(conn)
   184  		addConnAfter(conn, sp.head)
   185  	}
   186  	conn.streams--
   187  
   188  	if conn.streams == 0 {
   189  		now := time.Now().Unix()
   190  		if err != nil || sp.idle >= p.maxIdle || now-created > p.ttl {
   191  			removeConn(conn)
   192  			p.Unlock()
   193  			conn.ClientConn.Close()
   194  			return
   195  		}
   196  		sp.idle++
   197  	}
   198  	p.Unlock()
   199  }
   200  
   201  func (conn *poolConn) Close() {
   202  	conn.pool.release(conn.addr, conn, conn.err)
   203  }
   204  
   205  func removeConn(conn *poolConn) {
   206  	if conn.pre != nil {
   207  		conn.pre.next = conn.next
   208  	}
   209  	if conn.next != nil {
   210  		conn.next.pre = conn.pre
   211  	}
   212  	conn.pre = nil
   213  	conn.next = nil
   214  	conn.in = false
   215  	conn.sp.count--
   216  }
   217  
   218  func addConnAfter(conn *poolConn, after *poolConn) {
   219  	conn.next = after.next
   220  	conn.pre = after
   221  	if after.next != nil {
   222  		after.next.pre = conn
   223  	}
   224  	after.next = conn
   225  	conn.in = true
   226  	conn.sp.count++
   227  }