github.com/cloudwego/kitex@v0.9.0/pkg/generic/binary_test/generic_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 test
    18  
    19  import (
    20  	"context"
    21  	"encoding/binary"
    22  	"net"
    23  	"runtime"
    24  	"runtime/debug"
    25  	"strings"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/apache/thrift/lib/go/thrift"
    30  
    31  	"github.com/cloudwego/kitex/client"
    32  	"github.com/cloudwego/kitex/client/callopt"
    33  	"github.com/cloudwego/kitex/client/genericclient"
    34  	kt "github.com/cloudwego/kitex/internal/mocks/thrift"
    35  	"github.com/cloudwego/kitex/internal/test"
    36  	"github.com/cloudwego/kitex/pkg/generic"
    37  	"github.com/cloudwego/kitex/pkg/utils"
    38  	"github.com/cloudwego/kitex/server"
    39  )
    40  
    41  var addr = test.GetLocalAddress()
    42  
    43  func TestRun(t *testing.T) {
    44  	t.Run("RawThriftBinary", rawThriftBinary)
    45  	t.Run("RawThriftBinaryError", rawThriftBinaryError)
    46  	t.Run("RawThriftBinaryMockReq", rawThriftBinaryMockReq)
    47  	t.Run("RawThriftBinary2NormalServer", rawThriftBinary2NormalServer)
    48  }
    49  
    50  func rawThriftBinary(t *testing.T) {
    51  	svr := initRawThriftBinaryServer(new(GenericServiceImpl))
    52  	defer svr.Stop()
    53  
    54  	cli := initRawThriftBinaryClient()
    55  
    56  	method := "myMethod"
    57  	buf := genBinaryReqBuf(method)
    58  
    59  	resp, err := cli.GenericCall(context.Background(), method, buf, callopt.WithRPCTimeout(1*time.Second))
    60  	test.Assert(t, err == nil, err)
    61  	respBuf, ok := resp.([]byte)
    62  	test.Assert(t, ok)
    63  	test.Assert(t, string(respBuf[12+len(method):]) == respMsg)
    64  }
    65  
    66  func rawThriftBinaryError(t *testing.T) {
    67  	svr := initRawThriftBinaryServer(new(GenericServiceErrorImpl))
    68  	defer svr.Stop()
    69  
    70  	cli := initRawThriftBinaryClient()
    71  
    72  	method := "myMethod"
    73  	buf := genBinaryReqBuf(method)
    74  
    75  	_, err := cli.GenericCall(context.Background(), method, buf, callopt.WithRPCTimeout(100*time.Second))
    76  	test.Assert(t, err != nil)
    77  	test.Assert(t, strings.Contains(err.Error(), errResp), err.Error())
    78  }
    79  
    80  func rawThriftBinaryMockReq(t *testing.T) {
    81  	svr := initRawThriftBinaryServer(new(GenericServiceMockImpl))
    82  	defer svr.Stop()
    83  
    84  	cli := initRawThriftBinaryClient()
    85  
    86  	req := kt.NewMockReq()
    87  	req.Msg = reqMsg
    88  	strMap := make(map[string]string)
    89  	strMap["aa"] = "aa"
    90  	strMap["bb"] = "bb"
    91  	req.StrMap = strMap
    92  	args := kt.NewMockTestArgs()
    93  	args.Req = req
    94  
    95  	// encode
    96  	rc := utils.NewThriftMessageCodec()
    97  	buf, err := rc.Encode("Test", thrift.CALL, 100, args)
    98  	test.Assert(t, err == nil, err)
    99  
   100  	resp, err := cli.GenericCall(context.Background(), "Test", buf)
   101  	test.Assert(t, err == nil, err)
   102  
   103  	// decode
   104  	buf = resp.([]byte)
   105  	var result kt.MockTestResult
   106  	method, seqID, err := rc.Decode(buf, &result)
   107  	test.Assert(t, err == nil, err)
   108  	test.Assert(t, method == "Test", method)
   109  	test.Assert(t, seqID != 100, seqID)
   110  	test.Assert(t, *result.Success == respMsg)
   111  
   112  	seqID2, err2 := generic.GetSeqID(buf)
   113  	test.Assert(t, err2 == nil, err2)
   114  	test.Assert(t, seqID2 == seqID, seqID2)
   115  }
   116  
   117  func rawThriftBinary2NormalServer(t *testing.T) {
   118  	svr := initMockServer(new(MockImpl))
   119  	defer svr.Stop()
   120  
   121  	cli := initRawThriftBinaryClient()
   122  
   123  	req := kt.NewMockReq()
   124  	req.Msg = reqMsg
   125  	strMap := make(map[string]string)
   126  	strMap["aa"] = "aa"
   127  	strMap["bb"] = "bb"
   128  	req.StrMap = strMap
   129  	args := kt.NewMockTestArgs()
   130  	args.Req = req
   131  
   132  	// encode
   133  	rc := utils.NewThriftMessageCodec()
   134  	buf, err := rc.Encode("Test", thrift.CALL, 100, args)
   135  	test.Assert(t, err == nil, err)
   136  
   137  	resp, err := cli.GenericCall(context.Background(), "Test", buf, callopt.WithRPCTimeout(100*time.Second))
   138  	test.Assert(t, err == nil, err)
   139  
   140  	// decode
   141  	buf = resp.([]byte)
   142  	var result kt.MockTestResult
   143  	method, seqID, err := rc.Decode(buf, &result)
   144  	test.Assert(t, err == nil, err)
   145  	test.Assert(t, method == "Test", method)
   146  	// seqID会在kitex中覆盖，避免TTHeader和Payload codec 不一致问题
   147  	test.Assert(t, seqID != 100, seqID)
   148  	test.Assert(t, *result.Success == respMsg)
   149  }
   150  
   151  func initRawThriftBinaryClient() genericclient.Client {
   152  	g := generic.BinaryThriftGeneric()
   153  	cli := newGenericClient("destServiceName", g, addr)
   154  	return cli
   155  }
   156  
   157  func initRawThriftBinaryServer(handler generic.Service) server.Server {
   158  	addr, _ := net.ResolveTCPAddr("tcp", addr)
   159  	g := generic.BinaryThriftGeneric()
   160  	svr := newGenericServer(g, addr, handler)
   161  	return svr
   162  }
   163  
   164  func initMockServer(handler kt.Mock) server.Server {
   165  	tcpAddr, _ := net.ResolveTCPAddr("tcp", addr)
   166  	svr := NewMockServer(handler, tcpAddr)
   167  	return svr
   168  }
   169  
   170  func genBinaryReqBuf(method string) []byte {
   171  	idx := 0
   172  	buf := make([]byte, 12+len(method)+len(reqMsg))
   173  	binary.BigEndian.PutUint32(buf, thrift.VERSION_1)
   174  	idx += 4
   175  	binary.BigEndian.PutUint32(buf[idx:idx+4], uint32(len(method)))
   176  	idx += 4
   177  	copy(buf[idx:idx+len(method)], method)
   178  	idx += len(method)
   179  	binary.BigEndian.PutUint32(buf[idx:idx+4], 100)
   180  	idx += 4
   181  	copy(buf[idx:idx+len(reqMsg)], reqMsg)
   182  	return buf
   183  }
   184  
   185  func TestBinaryThriftGenericClientClose(t *testing.T) {
   186  	debug.SetGCPercent(-1)
   187  	defer debug.SetGCPercent(100)
   188  
   189  	var ms runtime.MemStats
   190  	runtime.ReadMemStats(&ms)
   191  
   192  	t.Logf("Before new clients, allocation: %f Mb, Number of allocation: %d\n", mb(ms.HeapAlloc), ms.HeapObjects)
   193  
   194  	cliCnt := 10000
   195  	clis := make([]genericclient.Client, cliCnt)
   196  	for i := 0; i < cliCnt; i++ {
   197  		g := generic.BinaryThriftGeneric()
   198  		clis[i] = newGenericClient("destServiceName", g, addr, client.WithShortConnection())
   199  	}
   200  
   201  	runtime.ReadMemStats(&ms)
   202  	preHeapAlloc, preHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   203  	t.Logf("After new clients, allocation: %f Mb, Number of allocation: %d\n", preHeapAlloc, preHeapObjects)
   204  
   205  	for _, cli := range clis {
   206  		_ = cli.Close()
   207  	}
   208  	runtime.GC()
   209  	runtime.ReadMemStats(&ms)
   210  	afterGCHeapAlloc, afterGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   211  	t.Logf("After close clients and GC be executed, allocation: %f Mb, Number of allocation: %d\n", afterGCHeapAlloc, afterGCHeapObjects)
   212  	test.Assert(t, afterGCHeapAlloc < preHeapAlloc && afterGCHeapObjects < preHeapObjects)
   213  
   214  	// Trigger the finalizer of kclient be executed
   215  	time.Sleep(200 * time.Millisecond) // ensure the finalizer be executed
   216  	runtime.GC()
   217  	runtime.ReadMemStats(&ms)
   218  	secondGCHeapAlloc, secondGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   219  	t.Logf("After second GC, allocation: %f Mb, Number of allocation: %d\n", secondGCHeapAlloc, secondGCHeapObjects)
   220  	test.Assert(t, secondGCHeapAlloc/2 < afterGCHeapAlloc && secondGCHeapObjects/2 < afterGCHeapObjects)
   221  }
   222  
   223  func TestBinaryThriftGenericClientFinalizer(t *testing.T) {
   224  	debug.SetGCPercent(-1)
   225  	defer debug.SetGCPercent(100)
   226  
   227  	var ms runtime.MemStats
   228  	runtime.ReadMemStats(&ms)
   229  	t.Logf("Before new clients, allocation: %f Mb, Number of allocation: %d\n", mb(ms.HeapAlloc), ms.HeapObjects)
   230  
   231  	cliCnt := 10000
   232  	clis := make([]genericclient.Client, cliCnt)
   233  	for i := 0; i < cliCnt; i++ {
   234  		g := generic.BinaryThriftGeneric()
   235  		clis[i] = newGenericClient("destServiceName", g, addr, client.WithShortConnection())
   236  	}
   237  
   238  	runtime.ReadMemStats(&ms)
   239  	t.Logf("After new clients, allocation: %f Mb, Number of allocation: %d\n", mb(ms.HeapAlloc), ms.HeapObjects)
   240  
   241  	runtime.GC()
   242  	runtime.ReadMemStats(&ms)
   243  	firstGCHeapAlloc, firstGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   244  	t.Logf("After first GC, allocation: %f Mb, Number of allocation: %d\n", firstGCHeapAlloc, firstGCHeapObjects)
   245  
   246  	// Trigger the finalizer of generic client be executed
   247  	time.Sleep(200 * time.Millisecond) // ensure the finalizer be executed
   248  	runtime.GC()
   249  	runtime.ReadMemStats(&ms)
   250  	secondGCHeapAlloc, secondGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   251  	t.Logf("After second GC, allocation: %f Mb, Number of allocation: %d\n", secondGCHeapAlloc, secondGCHeapObjects)
   252  	test.Assert(t, secondGCHeapAlloc < firstGCHeapAlloc && secondGCHeapObjects < firstGCHeapObjects)
   253  
   254  	// Trigger the finalizer of kClient be executed
   255  	time.Sleep(200 * time.Millisecond) // ensure the finalizer be executed
   256  	runtime.GC()
   257  	runtime.ReadMemStats(&ms)
   258  	thirdGCHeapAlloc, thirdGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
   259  	t.Logf("After third GC, allocation: %f Mb, Number of allocation: %d\n", thirdGCHeapAlloc, thirdGCHeapObjects)
   260  	test.Assert(t, thirdGCHeapAlloc < secondGCHeapAlloc/2 && thirdGCHeapObjects < secondGCHeapObjects/2)
   261  }
   262  
   263  func mb(byteSize uint64) float32 {
   264  	return float32(byteSize) / float32(1024*1024)
   265  }