mux_pool.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package netpollmux

import (
	"context"
	"net"
	"sync"
	"sync/atomic"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/remote"

	"golang.org/x/sync/singleflight"
)

var _ remote.LongConnPool = &MuxPool{}

// NewMuxConnPool size must be adjust to 2^N
func NewMuxConnPool(size int) *MuxPool {
	lp := &MuxPool{
		size: int32(size),
	}
	return lp
}

// MuxPool manages a pool of long connections.
type MuxPool struct {
	// mod     int32
	size    int32
	sfg     singleflight.Group
	connMap sync.Map // key address, value *muxCliConn
}

// Get pick or generate a net.Conn and return
func (mp *MuxPool) Get(ctx context.Context, network, address string, opt remote.ConnOption) (net.Conn, error) {
	v, ok := mp.connMap.Load(address)
	if ok {
		connection := v.(*conns).get()
		if connection != nil && connection.IsActive() {
			return connection, nil
		}
	}
	connection, err, _ := mp.sfg.Do(address, func() (i interface{}, e error) {
		conn, err := opt.Dialer.DialTimeout(network, address, opt.ConnectTimeout)
		if err != nil {
			return nil, err
		}
		connection := newMuxCliConn(conn.(netpoll.Connection))
		var cs *conns
		if ok {
			cs = v.(*conns)
		} else {
			cs = &conns{size: mp.size, conns: make([]*muxCliConn, mp.size)}
		}
		cs.put(connection)
		mp.connMap.Store(address, cs)
		return connection, nil
	})
	if err != nil {
		return nil, err
	}
	return connection.(*muxCliConn), nil
}

// Put implements the ConnPool interface.
func (mp *MuxPool) Put(conn net.Conn) error {
	if _, ok := conn.(*muxCliConn); ok {
		return nil
	}
	return conn.Close()
}

// Discard implements the ConnPool interface.
func (mp *MuxPool) Discard(conn net.Conn) error {
	if _, ok := conn.(*muxCliConn); ok {
		return nil
	}
	return conn.Close()
}

// Clean implements the LongConnPool interface.
func (mp *MuxPool) Clean(network, address string) {
	if v, ok := mp.connMap.Load(address); ok {
		mp.connMap.Delete(address)
		v.(*conns).close()
	}
}

// Close is to release resource of ConnPool, it is executed when client is closed.
func (mp *MuxPool) Close() error {
	mp.connMap.Range(func(addr, v interface{}) bool {
		mp.connMap.Delete(addr)
		v.(*conns).close()
		return true
	})
	return nil
}

type conns struct {
	index int32
	size  int32
	conns []*muxCliConn
}

func (c *conns) get() *muxCliConn {
	i := atomic.AddInt32(&c.index, 1)
	return c.conns[i%c.size]
}

// put and close together
func (c *conns) put(conn *muxCliConn) {
	for i := 0; i < int(c.size); i++ {
		if c.conns[i] == nil {
			c.conns[i] = conn
			return
		}
		if !c.conns[i].IsActive() {
			c.conns[i].close()
			c.conns[i] = conn
			return
		}
	}
}

func (c *conns) close() {
	for i := range c.conns {
		if c.conns[i] != nil {
			c.conns[i].close()
		}
	}
}