Get number of connections of a Transport in Go

You could try and check out hlts2/least-connections, a least-connections balancing algorithm written in golang.

Example:

lc, err := New([]*url.URL{
    {Host: "192.168.33.10"},
    {Host: "192.168.33.11"},
    {Host: "192.168.33.12"},
})

src1, done1 := lc.Next() // {Host: "192.168.33.10"}

src2, done2 := lc.Next() // {Host: "192.168.33.11"}

done1() // Reduce connection of src1

src3, done3 := lc.Next() // {Host: "192.168.33.10"}

It does use a sync.Mutex

type leastConnections struct {
    conns []conn
    mu    *sync.Mutex
}

Another more complex example: panjf2000/gnet presented in "Releasing a high-performance and lightweight event-loop networking library for Go", by Andy Pan.

gnet is an event-driven networking framework that is fast and lightweight.

Supporting multiple load-balancing algorithms: Round-Robin, Source Addr Hash and Least-Connections

Example

    // start a server
    // connect 10 clients
    // each client will pipe random data for 1-3 seconds.
    // the writes to the server will be random sizes. 0KB - 1MB.
    // the server will echo back the data.
    // waits for graceful connection closing.
    t.Run("poll", func(t *testing.T) {
        t.Run("tcp", func(t *testing.T) {
            t.Run("1-loop", func(t *testing.T) {
                testServe("tcp", ":9991", false, false, false, 10, RoundRobin)
            })
            t.Run("N-loop", func(t *testing.T) {
                testServe("tcp", ":9992", false, true, false, 10, LeastConnections)
            })
        })

Again, the structure is straightforward:

    // leastConnectionsEventLoopSet with Least-Connections algorithm.
    leastConnectionsEventLoopSet struct {
        sync.RWMutex
        minHeap                 minEventLoopHeap
        cachedRoot              *eventloop
        threshold               int32
        calibrateConnsThreshold int32
    }

You can wrap the field of DialContext in the *http.Transport to get the number of open connections from the Transport .

Here is my example code: https://gist.github.com/HattieWebb/2000c514b37a6a0bb5e8b55fadbe3433

// do not support if DialTLS is set.
func NewTransportWithConnectNum(old *http.Transport) *Transport{
    connectCounter:=&connectCounter_t{}
    tran:=&Transport{
        Transport:        old,
        connectCounter: connectCounter,
    }
    oldDialer:=tran.getOldDialer()
    tran.Transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error){
        conn1,err:=oldDialer(ctx,network,addr)
        if err!=nil{
            return nil,err
        }
        connectCounter.add(1)
        conn2:=&connWithCounter{
            Conn: conn1,
            connectCounter: connectCounter,
        }
        return conn2,nil
    }
    return tran
}

type Transport struct{
    *http.Transport
    connectCounter *connectCounter_t
}

func (tran *Transport) GetConnectionNum() int{
    return tran.connectCounter.get()
}

var zeroDialer net.Dialer

func (tran *Transport) getOldDialer() func(ctx context.Context, network, addr string) (net.Conn, error){
    if tran.DialContext!=nil{
        return tran.DialContext
    }
    if tran.Dial!=nil{
        return func(ctx context.Context, network, addr string) (net.Conn, error){
            return tran.Dial(network,addr)
        }
    }
    return zeroDialer.DialContext
}

type connectCounter_t struct{
    connectNum int
    connectNumLocker sync.Mutex
}

func (c *connectCounter_t) add(num int){
    c.connectNumLocker.Lock()
    c.connectNum+=num
    c.connectNumLocker.Unlock()
}
func (c *connectCounter_t) get()(num int){
    c.connectNumLocker.Lock()
    num = c.connectNum
    c.connectNumLocker.Unlock()
    return num
}

type connWithCounter struct{
    net.Conn
    closeCounterSyncOnce sync.Once
    connectCounter *connectCounter_t
}

func (conn *connWithCounter) Close() (err error){
    err = conn.Conn.Close()
    conn.closeCounterSyncOnce.Do(func(){
        conn.connectCounter.add(-1)
    })
    return err
}