github.com/cloudwego/kitex@v0.9.0/client/callopt/options.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 callopt contains options that control the behavior of client on request level.
package callopt

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/cloudwego/kitex/internal/client"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/fallback"
	"github.com/cloudwego/kitex/pkg/http"
	"github.com/cloudwego/kitex/pkg/retry"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
)

var callOptionsPool = sync.Pool{
	New: newOptions,
}

type CallOptions struct {
	configs      rpcinfo.MutableRPCConfig
	svr          remoteinfo.RemoteInfo
	locks        *client.ConfigLocks
	httpResolver http.Resolver

	// export field for using in client
	RetryPolicy    retry.Policy
	Fallback       *fallback.Policy
	CompressorName string
}

func newOptions() interface{} {
	return &CallOptions{
		locks: &client.ConfigLocks{
			Tags: make(map[string]struct{}),
		},
	}
}

// Recycle zeros the call option and put it to the pool.
func (co *CallOptions) Recycle() {
	if co == nil {
		return
	}
	co.configs = nil
	co.svr = nil
	co.RetryPolicy = retry.Policy{}
	co.Fallback = nil
	co.locks.Zero()
	callOptionsPool.Put(co)
}

// Option is a series of options used at the beginning of a RPC call.
type Option struct {
	f func(o *CallOptions, di *strings.Builder)
}

// F returns the function of the option.
// It's useful for creating streamcall.Option from existing callopt.Option
// Note: not all callopt.Option(s) are available for stream clients.
func (o Option) F() func(o *CallOptions, di *strings.Builder) {
	return o.f
}

// NewOption returns a new Option with the given function.
// It's useful for converting streamcall.Option back to a callopt.Option
func NewOption(f func(o *CallOptions, di *strings.Builder)) Option {
	return Option{f: f}
}

// WithHostPort specifies the target address for a RPC call.
// The given address will overwrite the result from Resolver.
func WithHostPort(hostport string) Option {
	return Option{func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithHostPort(")
		di.WriteString(hostport)
		di.WriteString("),")

		if err := setInstance(o.svr, hostport); err != nil {
			panic(fmt.Errorf("WithHostPort: %w", err))
		}
	}}
}

func setInstance(svr remoteinfo.RemoteInfo, hostport string) error {
	if _, err := net.ResolveTCPAddr("tcp", hostport); err == nil {
		svr.SetInstance(discovery.NewInstance("tcp", hostport, discovery.DefaultWeight, nil))
	} else if _, err := net.ResolveUnixAddr("unix", hostport); err == nil {
		svr.SetInstance(discovery.NewInstance("unix", hostport, discovery.DefaultWeight, nil))
	} else {
		return fmt.Errorf("invalid '%s'", hostport)
	}
	return nil
}

// WithURL specifies the target for a RPC call with url.
// The given url will be resolved to hostport and overwrites the result from Resolver.
func WithURL(url string) Option {
	return Option{func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithURL(")
		di.WriteString(url)
		di.WriteString("),")

		o.svr.SetTag(rpcinfo.HTTPURL, url)
		hostport, err := o.httpResolver.Resolve(url)
		if err != nil {
			panic(fmt.Errorf("http resolve failed: %w", err))
		}

		if err := setInstance(o.svr, hostport); err != nil {
			panic(fmt.Errorf("WithURL: %w", err))
		}
	}}
}

// WithHTTPHost specifies host in http header(work when RPC over http).
func WithHTTPHost(host string) Option {
	return Option{func(o *CallOptions, di *strings.Builder) {
		o.svr.SetTag(rpcinfo.HTTPHost, host)
	}}
}

// WithRPCTimeout specifies the RPC timeout for a RPC call.
// FIXME: callopt.WithRPCTimeout works only when client.WithRPCTimeout or
// client.WithTimeoutProvider is specified.
func WithRPCTimeout(d time.Duration) Option {
	return Option{func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithRPCTimeout()")

		o.configs.SetRPCTimeout(d)
		o.locks.Bits |= rpcinfo.BitRPCTimeout

		// TODO SetReadWriteTimeout 是否考虑删除
		o.configs.SetReadWriteTimeout(d)
		o.locks.Bits |= rpcinfo.BitReadWriteTimeout
	}}
}

// WithConnectTimeout specifies the connection timeout for a RPC call.
func WithConnectTimeout(d time.Duration) Option {
	return Option{func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithConnectTimeout()")

		o.configs.SetConnectTimeout(d)
		o.locks.Bits |= rpcinfo.BitConnectTimeout
	}}
}

// WithTag sets the tags for service discovery for a RPC call.
func WithTag(key, val string) Option {
	return Option{f: func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithTag")
		di.WriteByte('(')
		di.WriteString(key)
		di.WriteByte('=')
		di.WriteString(val)
		di.WriteByte(')')

		o.svr.SetTag(key, val)
		o.locks.Tags[key] = struct{}{}
	}}
}

// WithRetryPolicy sets the retry policy for a RPC call.
// Build retry.Policy with retry.BuildFailurePolicy or retry.BuildBackupRequest instead of building retry.Policy directly.
// Demos are provided below:
//
//	  demo1. call with failure retry policy, default retry error is Timeout
//	  	`resp, err := cli.Mock(ctx, req, callopt.WithRetryPolicy(retry.BuildFailurePolicy(retry.NewFailurePolicy())))`
//	  demo2. call with backup request policy
//	  	`bp := retry.NewBackupPolicy(10)
//		 	bp.WithMaxRetryTimes(1)
//	  	resp, err := cli.Mock(ctx, req, callopt.WithRetryPolicy(retry.BuildBackupRequest(bp)))`
func WithRetryPolicy(p retry.Policy) Option {
	return Option{f: func(o *CallOptions, di *strings.Builder) {
		if !p.Enable {
			return
		}
		if p.Type == retry.BackupType {
			di.WriteString("WithBackupRequest")
		} else {
			di.WriteString("WithFailureRetry")
		}
		o.RetryPolicy = p
	}}
}

// WithFallback is used to set the fallback policy for a RPC call.
// Demos are provided below:
//
//	demo1. call with fallback for error
//		`resp, err := cli.Mock(ctx, req, callopt.WithFallback(fallback.ErrorFallback(yourFBFunc))`
//	demo2. call with fallback for error and enable reportAsFallback, which sets reportAsFallback to be true and will do report(metric) as Fallback result
//		`resp, err := cli.Mock(ctx, req, callopt.WithFallback(fallback.ErrorFallback(yourFBFunc).EnableReportAsFallback())`
func WithFallback(fb *fallback.Policy) Option {
	return Option{f: func(o *CallOptions, di *strings.Builder) {
		if !fallback.IsPolicyValid(fb) {
			return
		}
		di.WriteString("WithFallback")
		o.Fallback = fb
	}}
}

// WithGRPCCompressor specifies the compressor for the GRPC frame payload.
// Supported compressor names: identity, gzip
// Custom compressors can be registered via `encoding.RegisterCompressor`
func WithGRPCCompressor(compressorName string) Option {
	return Option{f: func(o *CallOptions, di *strings.Builder) {
		di.WriteString("WithGRPCCompressor")
		di.WriteByte('(')
		di.WriteString(compressorName)
		di.WriteByte(')')
		o.CompressorName = compressorName
	}}
}

// Apply applies call options to the rpcinfo.RPCConfig and internal.RemoteInfo of kitex client.
// The return value records the name and arguments of each option.
// This function is for internal purpose only.
func Apply(cos []Option, cfg rpcinfo.MutableRPCConfig, svr remoteinfo.RemoteInfo, locks *client.ConfigLocks, httpResolver http.Resolver) (string, *CallOptions) {
	var buf strings.Builder
	buf.Grow(64)

	co := callOptionsPool.Get().(*CallOptions)
	co.configs = cfg
	co.svr = svr
	co.locks.Merge(locks)
	co.httpResolver = httpResolver

	buf.WriteByte('[')
	for i := range cos {
		if i > 0 {
			buf.WriteByte(',')
		}
		cos[i].f(co, &buf)
	}
	buf.WriteByte(']')

	co.locks.ApplyLocks(cfg, svr)
	return buf.String(), co
}