github.com/cloudwego/kitex@v0.9.0/client/rpctimeout_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 client

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpctimeout"
)

var panicMsg = "mock panic"

func block(ctx context.Context, request, response interface{}) (err error) {
	time.Sleep(1 * time.Second)
	return nil
}

func pass(ctx context.Context, request, response interface{}) (err error) {
	time.Sleep(200 * time.Millisecond)
	return nil
}

func panicEp(ctx context.Context, request, response interface{}) (err error) {
	panic(panicMsg)
}

func TestNewRPCTimeoutMW(t *testing.T) {
	t.Parallel()

	s := rpcinfo.NewEndpointInfo("mockService", "mockMethod", nil, nil)
	c := rpcinfo.NewRPCConfig()
	r := rpcinfo.NewRPCInfo(nil, s, nil, c, rpcinfo.NewRPCStats())
	m := rpcinfo.AsMutableRPCConfig(c)
	m.SetRPCTimeout(time.Millisecond * 500)

	var moreTimeout time.Duration
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), r)
	mwCtx := context.Background()
	mwCtx = context.WithValue(mwCtx, rpctimeout.TimeoutAdjustKey, &moreTimeout)

	var err error
	var mw1, mw2 endpoint.Middleware

	// 1. normal
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	err = mw1(mw2(pass))(ctx, nil, nil)
	test.Assert(t, err == nil)

	// 2. block to mock timeout
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	err = mw1(mw2(block))(ctx, nil, nil)
	test.Assert(t, err != nil, err)
	test.Assert(t, err.(*kerrors.DetailedError).ErrorType() == kerrors.ErrRPCTimeout)

	// 3. block, pass more timeout, timeout won't happen
	mw1 = rpctimeout.MiddlewareBuilder(510 * time.Millisecond)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	err = mw1(mw2(block))(ctx, nil, nil)
	test.Assert(t, err == nil)

	// 4. cancel
	cancelCtx, cancelFunc := context.WithCancel(ctx)
	time.AfterFunc(100*time.Millisecond, func() {
		cancelFunc()
	})
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	err = mw1(mw2(block))(cancelCtx, nil, nil)
	test.Assert(t, errors.Is(err, context.Canceled), err)

	// 5. panic with timeout, the panic is recovered
	// < v1.1.* panic happen, >=v1.1* wrap panic to error
	m.SetRPCTimeout(time.Millisecond * 500)
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	err = mw1(mw2(panicEp))(ctx, nil, nil)
	test.Assert(t, strings.Contains(err.Error(), panicMsg))

	// 6. panic without timeout, panic won't be recovered in timeout mw, but it will be recoverd in client.Call
	m.SetRPCTimeout(0)
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	test.Panic(t, func() { mw1(mw2(panicEp))(ctx, nil, nil) })

	// 7. panic with streaming, panic won't be recovered in timeout mw, but it will be recoverd in client.Call
	m = rpcinfo.AsMutableRPCConfig(c)
	m.SetInteractionMode(rpcinfo.Streaming)
	mw1 = rpctimeout.MiddlewareBuilder(0)(mwCtx)
	mw2 = rpcTimeoutMW(mwCtx)
	test.Panic(t, func() { mw1(mw2(panicEp))(ctx, nil, nil) })
}

func TestIsBusinessTimeout(t *testing.T) {
	type args struct {
		start        time.Time
		kitexTimeout time.Duration
		actualDDL    time.Time
		threshold    time.Duration
	}
	start := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
	tests := []struct {
		name string
		args args
		want bool
	}{
		{
			name: "business ddl ahead of kitex ddl by more than threshold",
			args: args{
				start:        start,
				kitexTimeout: time.Second * 2,
				actualDDL:    start.Add(time.Second),
				threshold:    time.Millisecond * 500,
			},
			want: true,
		},
		{
			name: "business ddl ahead of kitex ddl by less than threshold",
			args: args{
				start:        start,
				kitexTimeout: time.Second,
				actualDDL:    start.Add(time.Millisecond * 800),
				threshold:    time.Millisecond * 500,
			},
			want: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := isBusinessTimeout(tt.args.start, tt.args.kitexTimeout, tt.args.actualDDL, tt.args.threshold); got != tt.want {
				t.Errorf("isBusinessTimeout() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestRpcTimeoutMWTimeoutByBusiness(t *testing.T) {
	runTimeoutMW := func(timeout time.Duration) error {
		mw := rpcTimeoutMW(context.Background())
		processor := mw(func(ctx context.Context, req, rsp interface{}) error {
			time.Sleep(time.Millisecond * 100)
			return nil
		})

		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*30)
		defer cancel()

		ctx = rpcinfo.NewCtxWithRPCInfo(ctx, mockRPCInfo(timeout))

		return processor(ctx, nil, nil)
	}

	t.Run("Timeout by business with no need fine grained err code", func(t *testing.T) {
		rpctimeout.DisableGlobalNeedFineGrainedErrCode()
		if err := runTimeoutMW(time.Second); !kerrors.IsTimeoutError(err) {
			t.Errorf("rpcTimeoutMW(1s) = %v, want %v", err, kerrors.ErrRPCTimeout)
		}
	})

	t.Run("Timeout by business with no rpc timeout and no need fine grained err code", func(t *testing.T) {
		rpctimeout.DisableGlobalNeedFineGrainedErrCode()
		if err := runTimeoutMW(0); !kerrors.IsTimeoutError(err) {
			t.Errorf("rpcTimeoutMW(0) = %v, want %v", err, kerrors.ErrRPCTimeout)
		}
	})

	t.Run("Timeout by business with need fine grained err code", func(t *testing.T) {
		rpctimeout.EnableGlobalNeedFineGrainedErrCode()
		if err := runTimeoutMW(time.Second); err.(*kerrors.DetailedError).ErrorType() != kerrors.ErrTimeoutByBusiness {
			t.Errorf("rpcTimeoutMW(1s) = %v, want %v", err, kerrors.ErrTimeoutByBusiness)
		}
		rpctimeout.DisableGlobalNeedFineGrainedErrCode()
	})
}

func TestRpcTimeoutMWCancelByBusiness(t *testing.T) {
	runTimeoutMW := func() error {
		mw := rpcTimeoutMW(context.Background())
		processor := mw(func(ctx context.Context, req, rsp interface{}) error {
			time.Sleep(time.Millisecond * 100)
			return nil
		})

		ctx, cancel := context.WithCancel(context.Background())
		go func() {
			time.Sleep(10 * time.Millisecond)
			cancel()
		}()

		ctx = rpcinfo.NewCtxWithRPCInfo(ctx, mockRPCInfo(time.Second))

		return processor(ctx, nil, nil)
	}

	t.Run("Cancel by business with no need fine grained err code", func(t *testing.T) {
		rpctimeout.DisableGlobalNeedFineGrainedErrCode()
		if err := runTimeoutMW(); !kerrors.IsTimeoutError(err) {
			t.Errorf("rpcTimeoutMW() = %v, want %v", err, kerrors.ErrRPCTimeout)
		}
	})

	t.Run("Cancel by business with need fine grained err code", func(t *testing.T) {
		rpctimeout.EnableGlobalNeedFineGrainedErrCode()
		if err := runTimeoutMW(); err.(*kerrors.DetailedError).ErrorType() != kerrors.ErrCanceledByBusiness {
			t.Errorf("rpcTimeoutMW() = %v, want %v", err, kerrors.ErrCanceledByBusiness)
		}
		rpctimeout.DisableGlobalNeedFineGrainedErrCode()
	})
}

func mockRPCInfo(timeout time.Duration) rpcinfo.RPCInfo {
	s := rpcinfo.NewEndpointInfo("mockService", "mockMethod", nil, nil)
	c := rpcinfo.NewRPCConfig()
	mc := rpcinfo.AsMutableRPCConfig(c)
	_ = mc.SetRPCTimeout(timeout)
	return rpcinfo.NewRPCInfo(nil, s, nil, c, rpcinfo.NewRPCStats())
}

func Test_isBusinessTimeout(t *testing.T) {
	type args struct {
		start        time.Time
		kitexTimeout time.Duration
		actualDDL    time.Time
		threshold    time.Duration
	}
	now, _ := time.Parse(time.RFC3339, "2023-01-01T00:00:00Z08:00")
	threshold := rpctimeout.LoadBusinessTimeoutThreshold()
	tests := []struct {
		name string
		args args
		want bool
	}{
		{
			name: "zero-timeout",
			args: args{
				start:        now,
				kitexTimeout: 0,
				actualDDL:    now.Add(time.Second),
				threshold:    threshold,
			},
			want: true,
		},
		{
			name: "after-timeout",
			args: args{
				start:        now,
				kitexTimeout: time.Second,
				actualDDL:    now.Add(time.Second).Add(time.Millisecond),
				threshold:    threshold,
			},
			want: false,
		},
		{
			name: "after-(timeout-threshold)",
			args: args{
				start:        now,
				kitexTimeout: time.Second,
				actualDDL:    now.Add(time.Second).Add(-threshold).Add(time.Millisecond),
				threshold:    threshold,
			},
			want: false,
		},
		{
			name: "before-(timeout-threshold)",
			args: args{
				start:        now,
				kitexTimeout: time.Second,
				actualDDL:    now.Add(time.Second).Add(-threshold).Add(-time.Millisecond),
				threshold:    threshold,
			},
			want: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := isBusinessTimeout(tt.args.start, tt.args.kitexTimeout, tt.args.actualDDL, tt.args.threshold); got != tt.want {
				t.Errorf("isBusinessTimeout() = %v, want %v", got, tt.want)
			}
		})
	}
}