github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/client/loadbalance/balancer_test.go (about)

     1  package loadbalance
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/nyan233/littlerpc/core/common/logger"
     6  	"github.com/nyan233/littlerpc/core/common/transport"
     7  	"github.com/stretchr/testify/assert"
     8  	"net"
     9  	"runtime"
    10  	"strconv"
    11  	"testing"
    12  	"time"
    13  )
    14  
    15  func TestBalancer(t *testing.T) {
    16  	b := New(Config{
    17  		Scheme:                 "hash",
    18  		ResolverUpdateInterval: time.Second * 120,
    19  		Resolver: func() ([]RpcNode, error) {
    20  			return []RpcNode{
    21  				{Address: "192.168.10.1:8090"},
    22  				{Address: "192.168.10.1:8091"},
    23  				{Address: "192.168.10.1:8092"},
    24  				{Address: "192.168.10.1:8093"},
    25  				{Address: "192.168.10.1:8094"},
    26  				{Address: "192.168.10.1:8095"},
    27  				{Address: "192.168.10.1:8096"},
    28  			}, nil
    29  		},
    30  		ConnectionFactory: func(node RpcNode) (transport.ConnAdapter, error) {
    31  			return new(testConn), nil
    32  		},
    33  		CloseFunc: func(conn transport.ConnAdapter) {
    34  			_, ok := conn.(*testConn)
    35  			if !ok {
    36  				panic("conn type is not *net.TCPConn")
    37  			}
    38  		},
    39  		Logger:      logger.DefaultLogger,
    40  		MuxConnSize: 8,
    41  	})
    42  	for i := 1; i < 1<<10; i++ {
    43  		runtime.Gosched()
    44  		assert.NotNil(t, b.Target("/lrpc/internal/v1/Hello"+strconv.Itoa(i)))
    45  	}
    46  }
    47  
    48  func BenchmarkBalancer(b *testing.B) {
    49  	const NodeMax = 5000
    50  	for gSize := 1; gSize <= 16384; gSize <<= 1 {
    51  		b.Run(fmt.Sprintf("Concurrent-Node(%d)-%d", NodeMax, gSize), func(b *testing.B) {
    52  			b.ReportAllocs()
    53  			b.StopTimer()
    54  			nodes := make([]RpcNode, NodeMax)
    55  			for j := 0; j < len(nodes); j++ {
    56  				nodes[j] = RpcNode{
    57  					Address: fmt.Sprintf("10.18.2.3:%d", 1024+j),
    58  					Weight:  0,
    59  				}
    60  			}
    61  			balancer := New(Config{
    62  				Scheme:                 "hash",
    63  				ResolverUpdateInterval: time.Second * 120,
    64  				Resolver: func() ([]RpcNode, error) {
    65  					return nodes, nil
    66  				},
    67  				ConnectionFactory: func(node RpcNode) (transport.ConnAdapter, error) {
    68  					return new(testConn), nil
    69  				},
    70  				CloseFunc: func(conn transport.ConnAdapter) {
    71  					_, ok := conn.(*testConn)
    72  					if !ok {
    73  						panic("conn type is not *net.TCPConn")
    74  					}
    75  				},
    76  				Logger:      logger.DefaultLogger,
    77  				MuxConnSize: 16,
    78  			})
    79  			balancerFodder := make([]string, NodeMax)
    80  			for j := 0; j < len(balancerFodder); j++ {
    81  				balancerFodder[j] = fmt.Sprintf("lrpc/internal/v1/Hello.Test%d", j)
    82  			}
    83  			b.StartTimer()
    84  			b.SetParallelism(gSize)
    85  			b.RunParallel(func(pb *testing.PB) {
    86  				var count int
    87  				for pb.Next() {
    88  					count++
    89  					balancer.Target(balancerFodder[count%len(balancerFodder)])
    90  				}
    91  			})
    92  		})
    93  	}
    94  }
    95  
    96  type testConn struct {
    97  	net.TCPConn
    98  }
    99  
   100  func (c *testConn) SetSource(s interface{}) {
   101  	return
   102  }
   103  
   104  func (c *testConn) Source() interface{} {
   105  	return nil
   106  }