github.com/cloudwego/kitex@v0.9.0/pkg/generic/reflect_test/reflect_test.go

/*
 * Copyright 2023 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 test

import (
	"context"
	"errors"
	"math/rand"
	"net"
	"os"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/cloudwego/kitex/client"
	"github.com/cloudwego/kitex/client/callopt"
	"github.com/cloudwego/kitex/client/genericclient"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/server"
	"github.com/cloudwego/kitex/server/genericserver"

	"github.com/apache/thrift/lib/go/thrift"
	dt "github.com/cloudwego/dynamicgo/thrift"
	dg "github.com/cloudwego/dynamicgo/thrift/generic"
)

func TestMain(m *testing.M) {
	if runtime.GOOS == "windows" {
		klog.Infof("skip generic reflect_test on windows")
		return
	}
	initExampleDescriptor()
	addr := test.GetLocalAddress()
	svr := initServer(addr)
	cli = initClient(addr)

	mAddr := test.GetLocalAddress()
	msvr := initThriftMapServer(mAddr, "./idl/example.thrift", new(exampeServerImpl))
	mcli = initThriftMapClient(mAddr, "./idl/example.thrift")

	ret := m.Run()

	cli.Close()
	mcli.Close()
	svr.Stop()
	msvr.Stop()

	os.Exit(ret)
}

var (
	SampleListSize = 100
	SampleMapSize  = 100
)

func TestThriftReflectExample(t *testing.T) {
	testThriftReflectExample_Node(t)
	testThriftReflectExample_DOM(t)
}

func BenchmarkThriftReflectExample_Node(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testThriftReflectExample_Node(b)
	}
}

func BenchmarkThriftReflectExample_DOM(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testThriftReflectExample_DOM(b)
	}
}

func testThriftReflectExample_Node(t testing.TB) {
	log_id := strconv.Itoa(rand.Int())

	// make a request body
	req, err := makeExampleReqBinary(true, reqMsg, log_id)
	if err != nil {
		t.Fatal(err)
	}

	// wrap request as thrift CALL message
	buf, err := dt.WrapBinaryBody(req, method, dt.CALL, 1, 0)
	if err != nil {
		t.Fatal(err)
	}

	// generic call
	out, err := cli.GenericCall(context.Background(), method, buf, callopt.WithRPCTimeout(1*time.Second))
	if err != nil {
		t.Fatal(err)
	}

	// println("data size:", len(out.([]byte)) + len(req), "B")

	// unwrap REPLY message and get resp body
	_, _, _, _, body, err := dt.UnwrapBinaryMessage(out.([]byte))
	if err != nil {
		t.Fatal(err)
	}

	// biz logic...
	err = exampleClientHandler_Node(body, log_id)
	if err != nil {
		t.Fatal(err)
	}
}

func testThriftReflectExample_DOM(t testing.TB) {
	log_id := strconv.Itoa(rand.Int())

	// make a request body
	req, err := makeExampleReqBinary(true, reqMsg, log_id)
	if err != nil {
		t.Fatal(err)
	}

	// wrap request as thrift CALL message
	buf, err := dt.WrapBinaryBody(req, method, dt.CALL, 1, 0)
	if err != nil {
		t.Fatal(err)
	}

	// generic call
	out, err := cli.GenericCall(context.Background(), method, buf, callopt.WithRPCTimeout(1*time.Second))
	if err != nil {
		t.Fatal(err)
	}

	// unwrap REPLY message and get resp body
	_, _, _, _, body, err := dt.UnwrapBinaryMessage(out.([]byte))
	if err != nil {
		t.Fatal(err)
	}

	// biz logic...
	err = exampleClientHandler_DOM(body, log_id)
	if err != nil {
		t.Fatal(err)
	}
}

var (
	exampleReqDesc   *dt.TypeDescriptor
	exampleRespDesc  *dt.TypeDescriptor
	strDesc          *dt.TypeDescriptor
	baseLogidPath    []dg.Path
	dynamicgoOptions = &dg.Options{}
)

func initExampleDescriptor() {
	sdesc, err := dt.NewDescritorFromPath(context.Background(), "idl/example.thrift")
	if err != nil {
		panic(err)
	}
	exampleReqDesc = sdesc.Functions()["ExampleMethod"].Request().Struct().FieldById(1).Type()
	exampleRespDesc = sdesc.Functions()["ExampleMethod"].Response().Struct().FieldById(0).Type()
	strDesc = exampleReqDesc.Struct().FieldById(1).Type()
	baseLogidPath = []dg.Path{dg.NewPathFieldName("Base"), dg.NewPathFieldName("LogID")}
}

func initServer(addr string) server.Server {
	// init special server
	ip, _ := net.ResolveTCPAddr("tcp", addr)
	g := generic.BinaryThriftGeneric()
	svr := genericserver.NewServer(new(ExampleValueServiceImpl), g,
		server.WithServiceAddr(ip), server.WithExitWaitTime(time.Millisecond*10))
	go func() {
		err := svr.Run()
		if err != nil {
			panic(err)
		}
	}()
	test.WaitServerStart(addr)
	return svr
}

var cli genericclient.Client

func initClient(addr string) genericclient.Client {
	g := generic.BinaryThriftGeneric()
	genericCli, _ := genericclient.NewClient("destServiceName", g, client.WithHostPorts(addr))
	return genericCli
}

// makeExampleRespBinary make a Thrift-Binary-Encoding response using ExampleResp DOM
// Except msg, require_field and logid, which are reset everytime
func makeExampleRespBinary(msg, require_field, logid string) ([]byte, error) {
	dom := &dg.PathNode{
		Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
		Next: []dg.PathNode{
			{
				Path: dg.NewPathFieldId(1),
				Node: dg.NewNodeString(msg),
			},
			{
				Path: dg.NewPathFieldId(2),
				Node: dg.NewNodeString(require_field),
			},
			{
				Path: dg.NewPathFieldId(255),
				Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
				Next: []dg.PathNode{
					{
						Path: dg.NewPathFieldId(1),
						Node: dg.NewNodeString(logid),
					},
				},
			},
		},
	}
	return dom.Marshal(dynamicgoOptions)
}

// makeExampleReqBinary make a Thrift-Binary-Encoding request using ExampleReq DOM
// Except B, A and logid, which are reset everytime
func makeExampleReqBinary(B bool, A, logid string) ([]byte, error) {
	list := make([]dg.PathNode, SampleListSize+1)
	list[0] = dg.PathNode{
		Path: dg.NewPathIndex(0),
		Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
		Next: []dg.PathNode{
			{
				Path: dg.NewPathFieldId(1),
				Node: dg.NewNodeString(A),
			},
		},
	}
	for i := 1; i < len(list); i++ {
		list[i] = dg.PathNode{
			Path: dg.NewPathIndex(i),
			Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
			Next: []dg.PathNode{
				{
					Path: dg.NewPathFieldId(1),
					Node: dg.NewNodeString(A),
				},
			},
		}
	}
	m := make([]dg.PathNode, SampleListSize+1)
	m[0] = dg.PathNode{
		Path: dg.NewPathStrKey("a"),
		Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
		Next: []dg.PathNode{
			{
				Path: dg.NewPathFieldId(1),
				Node: dg.NewNodeString(A),
			},
		},
	}
	for i := 1; i < len(list); i++ {
		list[i] = dg.PathNode{
			Path: dg.NewPathStrKey(strconv.Itoa(i)),
			Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
			Next: []dg.PathNode{
				{
					Path: dg.NewPathFieldId(1),
					Node: dg.NewNodeString(A),
				},
			},
		}
	}

	dom := dg.PathNode{
		Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
		Next: []dg.PathNode{
			{
				Path: dg.NewPathFieldId(1),
				Node: dg.NewNodeString("Hello"),
			},
			{
				Path: dg.NewPathFieldId(2),
				Node: dg.NewNodeInt32(1),
			},
			{
				Path: dg.NewPathFieldId(3),
				Node: dg.NewTypedNode(thrift.LIST, thrift.STRUCT, 0),
				Next: list,
			},
			{
				Path: dg.NewPathFieldId(4),
				Node: dg.NewTypedNode(thrift.MAP, thrift.STRUCT, thrift.STRING),
				Next: m,
			},
			{
				Path: dg.NewPathFieldId(6),
				Node: dg.NewTypedNode(thrift.LIST, thrift.I64, 0),
				Next: []dg.PathNode{
					{
						Path: dg.NewPathIndex(0),
						Node: dg.NewNodeInt64(1),
					},
					{
						Path: dg.NewPathIndex(1),
						Node: dg.NewNodeInt64(2),
					},
					{
						Path: dg.NewPathIndex(2),
						Node: dg.NewNodeInt64(3),
					},
				},
			},
			{
				Path: dg.NewPathFieldId(7),
				Node: dg.NewNodeBool(B),
			},
			{
				Path: dg.NewPathFieldId(255),
				Node: dg.NewTypedNode(thrift.STRUCT, 0, 0),
				Next: []dg.PathNode{
					{
						Path: dg.NewPathFieldId(1),
						Node: dg.NewNodeString(logid),
					},
					{
						Path: dg.NewPathFieldId(2),
						Node: dg.NewNodeString("a.b.c"),
					},
					{
						Path: dg.NewPathFieldId(3),
						Node: dg.NewNodeString("127.0.0.1"),
					},
					{
						Path: dg.NewPathFieldId(4),
						Node: dg.NewNodeString("dynamicgo"),
					},
				},
			},
		},
	}
	return dom.Marshal(dynamicgoOptions)
}

const (
	method  = "ExampleMethod2"
	reqMsg  = "pending"
	respMsg = "ok"
)

// ExampleValueServiceImpl ...
type ExampleValueServiceImpl struct{}

// GenericCall ...
func (g *ExampleValueServiceImpl) GenericCall(ctx context.Context, method string, request interface{}) (interface{}, error) {
	// get and unwrap body with message
	in := request.([]byte)

	// unwarp thrift message and get request body
	methodName, _, seqID, _, body, err := dt.UnwrapBinaryMessage(in)
	if err != nil {
		return nil, err
	}

	// biz logic
	resp, err := exampleServerHandler(body)
	if err != nil {
		return nil, err
	}

	// wrap response as thrift REPLY message
	return dt.WrapBinaryBody(resp, methodName, dt.REPLY, 0, seqID)
}

// biz logic
func exampleServerHandler(request []byte) (resp []byte, err error) {
	// wrap body as Value
	req := dg.NewValue(exampleReqDesc, request)
	if err != nil {
		return nil, err
	}

	required_field := ""
	logid := ""
	// if B == true then get logid and required_field
	if b, err := req.FieldByName("B").Bool(); err == nil && b {
		if e := req.GetByPath(baseLogidPath...); e.Error() != "" {
			return nil, e
		} else {
			logid, _ = e.String()
		}
		if a := req.FieldByName("TestMap").GetByStr("a"); a.Error() != "" {
			return nil, a
		} else {
			required_field, _ = a.FieldByName("Bar").String()
		}
	}

	// make response with checked values
	return makeExampleRespBinary(respMsg, required_field, logid)
}

var clientRespPool = sync.Pool{
	New: func() interface{} {
		return &dg.PathNode{}
	},
}

// biz logic...
func exampleClientHandler_Node(response []byte, log_id string) error {
	// make dynamicgo/generic.Node with body
	resp := dg.NewNode(dt.STRUCT, response)

	// check node values by Node APIs
	msg, err := resp.Field(1).String()
	if err != nil {
		return err
	}
	if msg != respMsg {
		return errors.New("msg does not match")
	}
	require_field, err := resp.Field(2).String()
	if err != nil {
		return err
	}
	if require_field != reqMsg {
		return errors.New("require_field does not match")
	}

	return nil
}

func exampleClientHandler_DOM(response []byte, log_id string) error {
	// get dom from memory pool
	root := clientRespPool.Get().(*dg.PathNode)
	root.Node = dg.NewNode(dt.STRUCT, response)

	// load **first layer** children
	err := root.Load(false, dynamicgoOptions)
	if err != nil {
		return err
	}
	// spew.Dump(root) // -- only root.Next is set
	// check node values by PathNode APIs
	require_field2, err := root.Field(2, dynamicgoOptions).Node.String()
	if err != nil {
		return err
	}
	if require_field2 != reqMsg {
		return errors.New("require_field2 does not match")
	}

	// load **all layers** children
	err = root.Load(true, dynamicgoOptions)
	if err != nil {
		return err
	}

	// spew.Dump(root) // -- every PathNode.Next will be set if it is a nesting-typed (LIST/SET/MAP/STRUCT)
	// check node values by PathNode APIs
	logid, err := root.Field(255, dynamicgoOptions).Field(1, dynamicgoOptions).Node.String()
	if err != nil {
		return err
	}
	if logid != log_id {
		return errors.New("logid not match")
	}

	// recycle DOM
	root.ResetValue()
	clientRespPool.Put(root)
	return nil
}