github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/lbcache/cache_test.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 lbcache
    18  
    19  import (
    20  	"context"
    21  	"encoding/json"
    22  	"fmt"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/golang/mock/gomock"
    27  
    28  	mocksloadbalance "github.com/cloudwego/kitex/internal/mocks/loadbalance"
    29  	"github.com/cloudwego/kitex/internal/test"
    30  	"github.com/cloudwego/kitex/pkg/discovery"
    31  	"github.com/cloudwego/kitex/pkg/loadbalance"
    32  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    33  )
    34  
    35  var defaultOptions = Options{
    36  	RefreshInterval: defaultRefreshInterval,
    37  	ExpireInterval:  defaultExpireInterval,
    38  }
    39  
    40  func TestBuilder(t *testing.T) {
    41  	ctrl := gomock.NewController(t)
    42  	defer ctrl.Finish()
    43  
    44  	ins := discovery.NewInstance("tcp", "127.0.0.1:8888", 10, nil)
    45  	r := &discovery.SynthesizedResolver{
    46  		ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
    47  			return discovery.Result{Cacheable: true, CacheKey: key, Instances: []discovery.Instance{ins}}, nil
    48  		},
    49  		TargetFunc: func(ctx context.Context, target rpcinfo.EndpointInfo) string {
    50  			return "mockRoute"
    51  		},
    52  		NameFunc: func() string { return t.Name() },
    53  	}
    54  	lb := mocksloadbalance.NewMockLoadbalancer(ctrl)
    55  	lb.EXPECT().GetPicker(gomock.Any()).DoAndReturn(func(res discovery.Result) loadbalance.Picker {
    56  		test.Assert(t, res.Cacheable)
    57  		test.Assert(t, res.CacheKey == t.Name()+":mockRoute", res.CacheKey)
    58  		test.Assert(t, len(res.Instances) == 1)
    59  		test.Assert(t, res.Instances[0].Address().String() == "127.0.0.1:8888")
    60  		picker := mocksloadbalance.NewMockPicker(ctrl)
    61  		return picker
    62  	}).AnyTimes()
    63  	lb.EXPECT().Name().Return("Synthesized").AnyTimes()
    64  	NewBalancerFactory(r, lb, Options{Cacheable: true})
    65  	b, ok := balancerFactories.Load(cacheKey(t.Name(), "Synthesized", defaultOptions))
    66  	test.Assert(t, ok)
    67  	test.Assert(t, b != nil)
    68  	bl, err := b.(*BalancerFactory).Get(context.Background(), nil)
    69  	test.Assert(t, err == nil)
    70  	test.Assert(t, bl.GetPicker() != nil)
    71  	dump := Dump()
    72  	dumpJson, err := json.Marshal(dump)
    73  	test.Assert(t, err == nil)
    74  	test.Assert(t, string(dumpJson) == `{"TestBuilder|Synthesized|{5s 15s}":{"mockRoute":[{"Address":"tcp://127.0.0.1:8888","Weight":10}]}}`)
    75  }
    76  
    77  func TestCacheKey(t *testing.T) {
    78  	uniqueKey := cacheKey("hello", "world", Options{RefreshInterval: 15 * time.Second, ExpireInterval: 5 * time.Minute})
    79  	test.Assert(t, uniqueKey == "hello|world|{15s 5m0s}")
    80  }
    81  
    82  func TestBalancerCache(t *testing.T) {
    83  	count := 10
    84  	inss := []discovery.Instance{}
    85  	for i := 0; i < count; i++ {
    86  		inss = append(inss, discovery.NewInstance("tcp", fmt.Sprint(i), 10, nil))
    87  	}
    88  	r := &discovery.SynthesizedResolver{
    89  		TargetFunc: func(ctx context.Context, target rpcinfo.EndpointInfo) string {
    90  			return target.ServiceName()
    91  		},
    92  		ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
    93  			return discovery.Result{Cacheable: true, CacheKey: "svc", Instances: inss}, nil
    94  		},
    95  		NameFunc: func() string { return t.Name() },
    96  	}
    97  	lb := loadbalance.NewWeightedBalancer()
    98  	for i := 0; i < count; i++ {
    99  		blf := NewBalancerFactory(r, lb, Options{})
   100  		info := rpcinfo.NewEndpointInfo("svc", "", nil, nil)
   101  		b, err := blf.Get(context.Background(), info)
   102  		test.Assert(t, err == nil)
   103  		p := b.GetPicker()
   104  		for a := 0; a < count; a++ {
   105  			addr := p.Next(context.Background(), nil).Address().String()
   106  			t.Logf("count: %d addr: %s\n", i, addr)
   107  		}
   108  	}
   109  }
   110  
   111  type mockRebalancer struct {
   112  	rebalanceFunc func(ch discovery.Change)
   113  	deleteFunc    func(ch discovery.Change)
   114  }
   115  
   116  // Rebalance implements the Rebalancer interface.
   117  func (m *mockRebalancer) Rebalance(ch discovery.Change) {
   118  	if m.rebalanceFunc != nil {
   119  		m.rebalanceFunc(ch)
   120  	}
   121  }
   122  
   123  // Delete implements the Rebalancer interface.
   124  func (m *mockRebalancer) Delete(ch discovery.Change) {
   125  	if m.deleteFunc != nil {
   126  		m.deleteFunc(ch)
   127  	}
   128  }