github.com/cloudwego/kitex@v0.9.0/pkg/circuitbreak/cbsuite_test.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 circuitbreak

import (
	"context"
	"errors"
	"net"
	"reflect"
	"testing"
	"time"

	"github.com/bytedance/gopkg/cloud/circuitbreaker"

	"github.com/cloudwego/kitex/internal/test"
	kd "github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/event"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

var (
	errMock  = errors.New("mock error")
	addrMock = utils.NewNetAddr("mock", "127.0.0.1:8888")
)

func TestNewCBSuite(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	test.Assert(t, cb.servicePanel == nil)
	test.Assert(t, cb.instancePanel == nil)

	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	test.Assert(t, cb.instancePanel != nil)
	test.Assert(t, cb.servicePanel != nil)

	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return nil
	})
	ctx := prepareCtx()
	err := eps(ctx, nil, nil)
	test.Assert(t, err == nil)

	eps = endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return errMock
	})
	err = eps(ctx, nil, nil)
	test.Assert(t, err.Error() == errMock.Error(), err)

	cb.Close()
	test.Assert(t, cb.instancePanel == nil)
	test.Assert(t, cb.servicePanel == nil)
}

func TestGetServiceCB(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	cb.ServicePanel()
	cb.ServiceControl()
	test.Assert(t, cb.servicePanel != nil)
	test.Assert(t, cb.serviceControl != nil)
}

func TestServiceCB(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	cb.initServiceCB()
	opts := circuitbreaker.Options{
		ShouldTripWithKey: cb.svcTripFunc,
		CoolingTimeout:    100 * time.Millisecond,
		DetectTimeout:     100 * time.Millisecond,
	}
	cb.servicePanel, _ = circuitbreaker.NewPanel(cb.onServiceStateChange, opts)

	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return kerrors.ErrRPCTimeout.WithCause(errMock)
	})
	succEps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return nil
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}

	cfgKey := RPCInfo2Key(rpcinfo.GetRPCInfo(ctx))
	cb.UpdateServiceCBConfig(cfgKey, CBConfig{Enable: false})
	// disable circuit breaker
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout), err, i)
		}
	}

	// enable circuit breaker
	cb.UpdateServiceCBConfig(cfgKey, CBConfig{
		Enable:    true,
		ErrRate:   0.1,
		MinSample: 100,
	})
	// recover
	time.Sleep(200 * time.Millisecond)
	for i := 0; i < 300; i++ {
		err := succEps(ctx, nil, nil)
		if i == 2 {
			time.Sleep(200 * time.Millisecond)
		}
		if i == 1 || i == 2 {
			// half open: within detect timeout, still return cb err
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		} else {
			// i == 0 : after cooling timeout, half open
			test.Assert(t, err == nil, err)
		}
	}
	cb.Close()
}

func TestUpdateConfigErrRate(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	opts := circuitbreaker.Options{
		ShouldTripWithKey: cb.svcTripFunc,
		CoolingTimeout:    100 * time.Millisecond,
		DetectTimeout:     100 * time.Millisecond,
	}
	cb.servicePanel, _ = circuitbreaker.NewPanel(cb.onServiceStateChange, opts)

	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	count := 0
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		if count%10 == 0 {
			return kerrors.ErrRPCTimeout.WithCause(errMock)
		}
		return
	})
	cfgKey := RPCInfo2Key(rpcinfo.GetRPCInfo(ctx))
	cb.UpdateServiceCBConfig(cfgKey, CBConfig{
		Enable:    true,
		ErrRate:   0.1,
		MinSample: 100,
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 100 {
			test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cb.Close()
}

func TestUpdateConfigSamples(t *testing.T) {
	cb := NewCBSuite(nil)
	opts := circuitbreaker.Options{
		ShouldTripWithKey: cb.svcTripFunc,
		CoolingTimeout:    100 * time.Millisecond,
		DetectTimeout:     100 * time.Millisecond,
	}
	cb.servicePanel, _ = circuitbreaker.NewPanel(cb.onServiceStateChange, opts)

	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return kerrors.ErrRPCTimeout.WithCause(errMock)
	})
	cfgKey := RPCInfo2Key(rpcinfo.GetRPCInfo(ctx))
	cb.UpdateServiceCBConfig(cfgKey, CBConfig{
		Enable:    true,
		ErrRate:   0.5,
		MinSample: 100,
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 100 {
			test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cb.Close()
}

func TestDisableCB(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		return kerrors.ErrRPCTimeout.WithCause(errMock)
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cb.Close()
}

func TestInstanceCB(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	var mws []endpoint.Middleware
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	count := 0
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		if count%10 == 0 {
			return kerrors.ErrGetConnection.WithCause(errMock)
		}
		return
	})
	cfgKey := RPCInfo2Key(rpcinfo.GetRPCInfo(ctx))
	cb.UpdateServiceCBConfig(cfgKey, CBConfig{
		Enable:    true,
		ErrRate:   0.1,
		MinSample: 100,
	})
	cb.UpdateInstanceCBConfig(CBConfig{
		Enable:    true,
		ErrRate:   0.1,
		MinSample: 100,
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 100 {
			test.Assert(t, errors.Is(err, kerrors.ErrGetConnection), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cb.Close()
}

func TestCBSuite_AddEvent4CB(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	cb.SetEventBusAndQueue(event.NewEventBus(), event.NewQueue(event.MaxEventNum))
	var mws []endpoint.Middleware
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	count := 0
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		if count%10 == 0 {
			return kerrors.ErrGetConnection.WithCause(errMock)
		}
		return
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrGetConnection), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cb.Close()
}

func TestRemoveInstBreaker(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	bus := event.NewEventBus()
	queue := event.NewQueue(event.MaxEventNum)
	cb.SetEventBusAndQueue(bus, queue)

	var mws []endpoint.Middleware
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	count := 0
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		if count%10 == 0 {
			return kerrors.ErrGetConnection.WithCause(errMock)
		}
		return
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrGetConnection), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	test.Assert(t, len(cb.instancePanel.DumpBreakers()) == 1)
	instCBKey := addrMock.String()
	_, ok := cb.instancePanel.DumpBreakers()[instCBKey]
	test.Assert(t, ok)

	now := time.Now()
	bus.Dispatch(&event.Event{
		Name: kd.ChangeEventName,
		Time: now,
		Extra: &kd.Change{
			Removed: []kd.Instance{inst},
		},
	})
	time.Sleep(10 * time.Millisecond)
	_, ok = cb.instancePanel.DumpBreakers()[instCBKey]
	test.Assert(t, !ok)
	test.Assert(t, len(cb.instancePanel.DumpBreakers()) == 0)
	cb.Close()
}

func TestCBSuite_Dump(t *testing.T) {
	cb := NewCBSuite(RPCInfo2Key)
	bus := event.NewEventBus()
	queue := event.NewQueue(event.MaxEventNum)
	cb.SetEventBusAndQueue(bus, queue)

	var mws []endpoint.Middleware
	mws = append(mws, cb.ServiceCBMW())
	mws = append(mws, cb.InstanceCBMW())
	ctx := prepareCtx()
	count := 0
	eps := endpoint.Chain(mws...)(func(ctx context.Context, req, resp interface{}) (err error) {
		if count%10 == 0 {
			return kerrors.ErrGetConnection.WithCause(errMock)
		}
		return
	})
	for i := 0; i < 300; i++ {
		err := eps(ctx, nil, nil)
		if i < 200 {
			test.Assert(t, errors.Is(err, kerrors.ErrGetConnection), err, i)
		} else {
			test.Assert(t, errors.Is(err, kerrors.ErrCircuitBreak), err, i)
		}
	}
	cfgKey := RPCInfo2Key(rpcinfo.GetRPCInfo(ctx))
	cbDump := cb.Dump().(map[string]interface{})
	test.Assert(t, len(cbDump) == 3)
	cbCfg := cbDump[cbConfig].(map[string]interface{})
	test.Assert(t, len(cbCfg) == 2, len(cbCfg), cbCfg)

	instCfg, ok := cbCfg[instanceCBKey].(CBConfig)
	test.Assert(t, ok)
	test.Assert(t, instCfg.Enable)
	test.Assert(t, instCfg.MinSample == GetDefaultCBConfig().MinSample)
	test.Assert(t, instCfg.ErrRate == GetDefaultCBConfig().ErrRate)

	svcCfg, ok := cbCfg[serviceCBKey].(map[string]interface{})
	test.Assert(t, ok)
	cfg, ok := svcCfg[cfgKey].(CBConfig)
	test.Assert(t, ok)
	test.Assert(t, cfg.Enable)
	test.Assert(t, cfg.MinSample == GetDefaultCBConfig().MinSample)
	test.Assert(t, cfg.ErrRate == GetDefaultCBConfig().ErrRate)

	// update config, then dump
	newCfg := CBConfig{
		Enable:    true,
		ErrRate:   0.1,
		MinSample: 100,
	}
	cb.UpdateServiceCBConfig(cfgKey, newCfg)
	cb.UpdateInstanceCBConfig(newCfg)
	cbDump = cb.Dump().(map[string]interface{})
	test.Assert(t, len(cbDump) == 3)

	cbCfg = cbDump[cbConfig].(map[string]interface{})
	test.Assert(t, len(cbCfg) == 2, len(cbCfg), cbCfg)

	instCfg, ok = cbCfg[instanceCBKey].(CBConfig)
	test.Assert(t, ok)
	test.Assert(t, instCfg.Enable)
	test.Assert(t, instCfg.MinSample == newCfg.MinSample)
	test.Assert(t, instCfg.ErrRate == newCfg.ErrRate)

	svcCfg, ok = cbCfg[serviceCBKey].(map[string]interface{})
	test.Assert(t, ok)
	cfg, ok = svcCfg[cfgKey].(CBConfig)
	test.Assert(t, ok)
	test.Assert(t, cfg.Enable)
	test.Assert(t, cfg.MinSample == newCfg.MinSample)
	test.Assert(t, cfg.ErrRate == newCfg.ErrRate)

	cb.Close()
}

func prepareCtx() context.Context {
	from := rpcinfo.NewEndpointInfo("caller", "", nil, nil)
	to := rpcinfo.NewEndpointInfo("callee", "method", addrMock, nil)
	ri := rpcinfo.NewRPCInfo(from, to, nil, nil, nil)
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	return ctx
}

var inst kd.Instance = &mockInst{}

type mockInst struct{}

func (m mockInst) Address() net.Addr {
	return addrMock
}

func (m mockInst) Weight() int {
	return 10
}

func (m mockInst) Tag(key string) (value string, exist bool) {
	return
}

func TestCBConfig_DeepCopy(t *testing.T) {
	type fields struct {
		c *CBConfig
	}
	tests := []struct {
		name   string
		fields fields
		want   *CBConfig
	}{
		{
			name: "test_nil_copy",
			fields: fields{
				c: nil,
			},
			want: nil,
		},
		{
			name: "test_all_copy",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			want: &CBConfig{
				Enable:    true,
				ErrRate:   0.1,
				MinSample: 10,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := tt.fields.c.DeepCopy(); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("DeepCopy() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestCBConfig_Equals(t *testing.T) {
	type fields struct {
		c *CBConfig
	}
	type args struct {
		other *CBConfig
	}
	tests := []struct {
		name   string
		fields fields
		args   args
		want   bool
	}{
		{
			name: "test_nil_equal",
			fields: fields{
				c: nil,
			},
			args: args{
				other: nil,
			},
			want: true,
		},
		{
			name: "test_nil_not_equal",
			fields: fields{
				c: nil,
			},
			args: args{
				other: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			want: false,
		},
		{
			name: "test_other_nil_not_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: nil,
			},
			want: false,
		},
		{
			name: "test_all_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			want: true,
		},
		{
			name: "test_all_not_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: &CBConfig{
					Enable:    false,
					ErrRate:   0.2,
					MinSample: 20,
				},
			},
			want: false,
		},
		{
			name: "test_enable_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: &CBConfig{
					Enable:    true,
					ErrRate:   0.2,
					MinSample: 20,
				},
			},
			want: false,
		},
		{
			name: "test_err_rate_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: &CBConfig{
					Enable:    false,
					ErrRate:   0.1,
					MinSample: 20,
				},
			},
			want: false,
		},
		{
			name: "test_min_sample_equal",
			fields: fields{
				c: &CBConfig{
					Enable:    true,
					ErrRate:   0.1,
					MinSample: 10,
				},
			},
			args: args{
				other: &CBConfig{
					Enable:    false,
					ErrRate:   0.2,
					MinSample: 10,
				},
			},
			want: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := tt.fields.c.Equals(tt.args.other); got != tt.want {
				t.Errorf("Equals() = %v, want %v", got, tt.want)
			}
		})
	}
}