v2ray-core/transport/internet/internal/pool.go

package internal

import (
	"net"
	"sync"
	"time"

	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
)

// ConnectionRecyler is the interface for recycling connections.
type ConnectionRecyler interface {
	// Put returns a connection back to a connection pool.
	Put(ConnectionID, net.Conn)
}

// ConnectionID is the ID of a connection.
type ConnectionID struct {
	Local      v2net.Address
	Remote     v2net.Address
	RemotePort v2net.Port
}

// NewConnectionID creates a new ConnectionId.
func NewConnectionID(source v2net.Address, dest v2net.Destination) ConnectionID {
	return ConnectionID{
		Local:      source,
		Remote:     dest.Address,
		RemotePort: dest.Port,
	}
}

// ExpiringConnection is a connection that will expire in certain time.
type ExpiringConnection struct {
	conn   net.Conn
	expire time.Time
}

// Expired returns true if the connection has expired.
func (ec *ExpiringConnection) Expired() bool {
	return ec.expire.Before(time.Now())
}

// Pool is a connection pool.
type Pool struct {
	sync.Mutex
	connsByDest map[ConnectionID][]*ExpiringConnection
	cleanupOnce signal.Once
}

// NewConnectionPool creates a new Pool.
func NewConnectionPool() *Pool {
	return &Pool{
		connsByDest: make(map[ConnectionID][]*ExpiringConnection),
	}
}

// Get returns a connection with matching connection ID. Nil if not found.
func (p *Pool) Get(id ConnectionID) net.Conn {
	p.Lock()
	defer p.Unlock()

	list, found := p.connsByDest[id]
	if !found {
		return nil
	}
	connIdx := -1
	for idx, conn := range list {
		if !conn.Expired() {
			connIdx = idx
			break
		}
	}
	if connIdx == -1 {
		return nil
	}
	listLen := len(list)
	conn := list[connIdx]
	if connIdx != listLen-1 {
		list[connIdx] = list[listLen-1]
	}
	list = list[:listLen-1]
	p.connsByDest[id] = list
	return conn.conn
}

func (p *Pool) cleanup() {
	defer p.cleanupOnce.Reset()

	for len(p.connsByDest) > 0 {
		time.Sleep(time.Second * 5)
		expiredConns := make([]net.Conn, 0, 16)
		p.Lock()
		for dest, list := range p.connsByDest {
			validConns := make([]*ExpiringConnection, 0, len(list))
			for _, conn := range list {
				if conn.Expired() {
					expiredConns = append(expiredConns, conn.conn)
				} else {
					validConns = append(validConns, conn)
				}
			}
			if len(validConns) != len(list) {
				p.connsByDest[dest] = validConns
			}
		}
		p.Unlock()
		for _, conn := range expiredConns {
			conn.Close()
		}
	}
}

// Put implements ConnectionRecyler.Put().
func (p *Pool) Put(id ConnectionID, conn net.Conn) {
	expiringConn := &ExpiringConnection{
		conn:   conn,
		expire: time.Now().Add(time.Second * 4),
	}

	p.Lock()
	defer p.Unlock()

	list, found := p.connsByDest[id]
	if !found {
		list = []*ExpiringConnection{expiringConn}
	} else {
		list = append(list, expiringConn)
	}
	p.connsByDest[id] = list

	p.cleanupOnce.Do(func() {
		go p.cleanup()
	})
}