github.com/cloudwego/kitex@v0.9.0/pkg/rpcinfo/remoteinfo/remoteInfo_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 remoteinfo_test

import (
	"sync"
	"testing"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/internal"
	mocksdiscovery "github.com/cloudwego/kitex/internal/mocks/discovery"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

func TestRemoteInfo(t *testing.T) {
	test.PanicAt(t, func() {
		_ = remoteinfo.NewRemoteInfo(nil, "method")
	}, func(err interface{}) bool {
		e, ok := err.(error)
		return ok && e.Error() == "nil basicInfo"
	})

	bi := &rpcinfo.EndpointBasicInfo{
		ServiceName: "service",
		Method:      "method0",
		Tags:        map[string]string{"a": "b"},
	}
	ri := remoteinfo.NewRemoteInfo(bi, "method1")
	test.Assert(t, ri != nil)
	test.Assert(t, ri.ServiceName() == "service")
	test.Assert(t, ri.Method() == "method1")
	test.Assert(t, ri.DefaultTag("a", "-") == "b")
	test.Assert(t, ri.DefaultTag("x", "-") == "-")
}

func TestAsRemoteInfo(t *testing.T) {
	test.Assert(t, remoteinfo.AsRemoteInfo(nil) == nil)

	ei := rpcinfo.EmptyEndpointInfo()
	test.Assert(t, remoteinfo.AsRemoteInfo(ei) == nil)

	bi := &rpcinfo.EndpointBasicInfo{
		ServiceName: "service",
		Method:      "method0",
		Tags:        map[string]string{"a": "b", "1": "2"},
	}
	ri := remoteinfo.NewRemoteInfo(bi, "method1")
	ri.SetTagLock("1")
	test.Assert(t, ri.GetInstance() == nil)

	ei = ri.ImmutableView()
	ri2 := remoteinfo.AsRemoteInfo(ei)
	test.Assert(t, ri2 != nil)

	ri2.SetServiceName("service2")
	test.Assert(t, ri2.SetTag("a", "aa") == nil)
	test.Assert(t, ri2.SetTag("1", "11") != nil)

	test.Assert(t, ri.ServiceName() == "service2")
	test.Assert(t, ri.DefaultTag("a", "-") == "aa")
	test.Assert(t, ri.DefaultTag("1", "-") == "2")
	test.Assert(t, ri.GetInstance() == nil)
	test.Assert(t, ri.Address() == nil)

	ins := discovery.NewInstance("n", "a", 1, nil)
	ri2.SetInstance(ins)
	test.Assert(t, ri.GetInstance() == ins)
	na := ri.Address()
	test.Assert(t, na.Network() == "n" && na.String() == "a")

	_, ok := ins.(remoteinfo.RefreshableInstance)
	na = utils.NewNetAddr("nnn", "aaa")
	test.Assert(t, ri2.SetRemoteAddr(na) == ok)
	if ok {
		na2 := ri.GetInstance().Address()
		test.Assert(t, na2 != nil)
		test.Assert(t, na2.Network() == "nnn", na2.Network(), na2)
		test.Assert(t, na2.String() == "aaa")
	}
}

func TestSetTag(t *testing.T) {
	lockKey := "lock"
	unlockKey := "unlock"
	bi := &rpcinfo.EndpointBasicInfo{
		ServiceName: "service",
		Tags:        map[string]string{lockKey: "a", unlockKey: "b"},
	}
	ri := remoteinfo.NewRemoteInfo(bi, "method1")

	ri.SetTagLock(lockKey)
	val, exist := ri.Tag(lockKey)
	test.Assert(t, val == "a")
	test.Assert(t, exist)
	val, exist = ri.Tag(unlockKey)
	test.Assert(t, val == "b")
	test.Assert(t, exist)

	// lock key cannot be reset
	ri.SetTag(lockKey, "aa")
	val, _ = ri.Tag(lockKey)
	test.Assert(t, val == "a")

	// lock key still can be reset with ForceSetTag
	ri.ForceSetTag(lockKey, "aa")
	val, _ = ri.Tag(lockKey)
	test.Assert(t, val == "aa")

	// unlock key can be reset
	ri.SetTag(unlockKey, "bb")
	val, _ = ri.Tag(unlockKey)
	test.Assert(t, val == "bb")

	// unlock key can be reset
	ri.ForceSetTag(unlockKey, "bb")
	val, _ = ri.Tag(unlockKey)
	test.Assert(t, val == "bb")
}

func TestGetTag(t *testing.T) {
	clusterKey, idcKey, myCluster, myIDC, mock := "cluster", "idc", "myCluster", "myIDC", "mock"

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	inst := mocksdiscovery.NewMockInstance(ctrl)
	inst.EXPECT().Tag(gomock.Any()).DoAndReturn(
		func(key string) (string, bool) {
			switch key {
			case clusterKey:
				return myCluster, true
			case idcKey:
				return myIDC, true
			}
			return "", false
		}).AnyTimes()

	bi := &rpcinfo.EndpointBasicInfo{
		ServiceName: "service",
	}
	ri := remoteinfo.NewRemoteInfo(bi, "method1")

	// case1: no cluster and idc
	valCluster, clusterOk := ri.Tag(clusterKey)
	valIDC, idcOk := ri.Tag(idcKey)
	test.Assert(t, !clusterOk)
	test.Assert(t, valCluster == "", valCluster)
	test.Assert(t, !idcOk)
	test.Assert(t, valIDC == "", valIDC)

	// case2: have cluster and idc which value are from tag set
	ri.SetTag(clusterKey, mock)
	ri.SetTag(idcKey, mock)
	valCluster, clusterOk = ri.Tag(clusterKey)
	valIDC, idcOk = ri.Tag(idcKey)
	test.Assert(t, clusterOk)
	test.Assert(t, valCluster == mock, valCluster)
	test.Assert(t, idcOk)
	test.Assert(t, valIDC == mock, valIDC)

	// case3: have cluster and idc which value are from instance, the priority of tags from instance is higher than tag set
	ri.SetInstance(inst)
	valCluster, clusterOk = ri.Tag(clusterKey)
	valIDC, idcOk = ri.Tag(idcKey)
	test.Assert(t, clusterOk)
	test.Assert(t, valCluster == myCluster, valCluster)
	test.Assert(t, idcOk)
	test.Assert(t, valIDC == myIDC, valIDC)
}

func TestRecycleRace(t *testing.T) {
	ri := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{ServiceName: "service", Tags: map[string]string{"key1": "val1"}}, "method1")

	// test the data race problem caused by tag modification
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		ri.ForceSetTag("key11", "val11")
		wg.Done()
	}()
	go func() {
		ri.(internal.Reusable).Recycle()
		wg.Done()
	}()
	wg.Wait()
}