github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpollmux/mux_pool.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo 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 netpollmux
    18  
    19  import (
    20  	"context"
    21  	"net"
    22  	"sync"
    23  	"sync/atomic"
    24  
    25  	"github.com/cloudwego/netpoll"
    26  
    27  	"github.com/cloudwego/kitex/pkg/remote"
    28  
    29  	"golang.org/x/sync/singleflight"
    30  )
    31  
    32  var _ remote.LongConnPool = &MuxPool{}
    33  
    34  // NewMuxConnPool size must be adjust to 2^N
    35  func NewMuxConnPool(size int) *MuxPool {
    36  	lp := &MuxPool{
    37  		size: int32(size),
    38  	}
    39  	return lp
    40  }
    41  
    42  // MuxPool manages a pool of long connections.
    43  type MuxPool struct {
    44  	// mod     int32
    45  	size    int32
    46  	sfg     singleflight.Group
    47  	connMap sync.Map // key address, value *muxCliConn
    48  }
    49  
    50  // Get pick or generate a net.Conn and return
    51  func (mp *MuxPool) Get(ctx context.Context, network, address string, opt remote.ConnOption) (net.Conn, error) {
    52  	v, ok := mp.connMap.Load(address)
    53  	if ok {
    54  		connection := v.(*conns).get()
    55  		if connection != nil && connection.IsActive() {
    56  			return connection, nil
    57  		}
    58  	}
    59  	connection, err, _ := mp.sfg.Do(address, func() (i interface{}, e error) {
    60  		conn, err := opt.Dialer.DialTimeout(network, address, opt.ConnectTimeout)
    61  		if err != nil {
    62  			return nil, err
    63  		}
    64  		connection := newMuxCliConn(conn.(netpoll.Connection))
    65  		var cs *conns
    66  		if ok {
    67  			cs = v.(*conns)
    68  		} else {
    69  			cs = &conns{size: uint32(mp.size), conns: make([]*muxCliConn, mp.size)}
    70  		}
    71  		cs.put(connection)
    72  		mp.connMap.Store(address, cs)
    73  		return connection, nil
    74  	})
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  	return connection.(*muxCliConn), nil
    79  }
    80  
    81  // Put implements the ConnPool interface.
    82  func (mp *MuxPool) Put(conn net.Conn) error {
    83  	if _, ok := conn.(*muxCliConn); ok {
    84  		return nil
    85  	}
    86  	return conn.Close()
    87  }
    88  
    89  // Discard implements the ConnPool interface.
    90  func (mp *MuxPool) Discard(conn net.Conn) error {
    91  	if _, ok := conn.(*muxCliConn); ok {
    92  		return nil
    93  	}
    94  	return conn.Close()
    95  }
    96  
    97  // Clean implements the LongConnPool interface.
    98  func (mp *MuxPool) Clean(network, address string) {
    99  	if v, ok := mp.connMap.Load(address); ok {
   100  		mp.connMap.Delete(address)
   101  		v.(*conns).close()
   102  	}
   103  }
   104  
   105  // Close is to release resource of ConnPool, it is executed when client is closed.
   106  func (mp *MuxPool) Close() error {
   107  	mp.connMap.Range(func(addr, v interface{}) bool {
   108  		mp.connMap.Delete(addr)
   109  		v.(*conns).close()
   110  		return true
   111  	})
   112  	return nil
   113  }
   114  
   115  type conns struct {
   116  	index uint32
   117  	size  uint32
   118  	conns []*muxCliConn
   119  }
   120  
   121  func (c *conns) get() *muxCliConn {
   122  	i := atomic.AddUint32(&c.index, 1)
   123  	return c.conns[i%c.size]
   124  }
   125  
   126  // put and close together
   127  func (c *conns) put(conn *muxCliConn) {
   128  	for i := 0; i < int(c.size); i++ {
   129  		if c.conns[i] == nil {
   130  			c.conns[i] = conn
   131  			return
   132  		}
   133  		if !c.conns[i].IsActive() {
   134  			c.conns[i].close()
   135  			c.conns[i] = conn
   136  			return
   137  		}
   138  	}
   139  }
   140  
   141  func (c *conns) close() {
   142  	for i := range c.conns {
   143  		if c.conns[i] != nil {
   144  			c.conns[i].close()
   145  		}
   146  	}
   147  }