client/option.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"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/cloudwego/kitex/internal/client"
	internal_stats "github.com/cloudwego/kitex/internal/stats"
	"github.com/cloudwego/kitex/pkg/circuitbreak"
	"github.com/cloudwego/kitex/pkg/connpool"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/http"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/loadbalance"
	"github.com/cloudwego/kitex/pkg/loadbalance/lbcache"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/netpollmux"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/retry"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/pkg/warmup"
	"github.com/cloudwego/kitex/transport"
)

// Option is the only way to config client.
type Option = client.Option

// Options is used to initialize a client.
type Options = client.Options

// A Suite is a collection of Options. It is useful to assemble multiple associated
// Options as a single one to keep the order or presence in a desired manner.
type Suite interface {
	Options() []Option
}

// WithTransportProtocol sets the transport protocol for client.
func WithTransportProtocol(tp transport.Protocol) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		tpName := tp.String()
		if tpName == transport.Unknown {
			panic(fmt.Errorf("WithTransportProtocol: invalid '%v'", tp))
		}
		di.Push(fmt.Sprintf("WithTransportProtocol(%s)", tpName))
		rpcinfo.AsMutableRPCConfig(o.Configs).SetTransportProtocol(tp)
	}}
}

// WithSuite adds a option suite for client.
func WithSuite(suite Suite) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		var nested struct {
			Suite   string
			Options utils.Slice
		}
		nested.Suite = fmt.Sprintf("%T(%+v)", suite, suite)

		for _, op := range suite.Options() {
			op.F(o, &nested.Options)
		}
		di.Push(nested)
	}}
}

// WithMiddleware adds middleware for client to handle request.
func WithMiddleware(mw endpoint.Middleware) Option {
	mwb := func(ctx context.Context) endpoint.Middleware {
		return mw
	}
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMiddleware(%+v)", utils.GetFuncName(mw)))
		o.MWBs = append(o.MWBs, mwb)
	}}
}

// WithMiddlewareBuilder adds middleware that depend on context for client to handle request
func WithMiddlewareBuilder(mwb endpoint.MiddlewareBuilder) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMiddlewareBuilder(%+v)", utils.GetFuncName(mwb)))
		o.MWBs = append(o.MWBs, mwb)
	}}
}

// WithInstanceMW adds middleware for client to handle request after service discovery and loadbalance process.
func WithInstanceMW(mw endpoint.Middleware) Option {
	imwb := func(ctx context.Context) endpoint.Middleware {
		return mw
	}
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithInstanceMW(%+v)", utils.GetFuncName(mw)))
		o.IMWBs = append(o.IMWBs, imwb)
	}}
}

// WithDestService specifies the name of target service.
func WithDestService(svr string) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithDestService(%s)", svr))
		o.Svr.ServiceName = svr
	}}
}

// WithHostPorts specifies the target instance addresses when doing service discovery.
// It overwrites the results from the Resolver.
func WithHostPorts(hostports ...string) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithHostPorts(%v)", hostports))

		var ins []discovery.Instance
		for _, hp := range hostports {
			if _, err := net.ResolveTCPAddr("tcp", hp); err == nil {
				ins = append(ins, discovery.NewInstance("tcp", hp, discovery.DefaultWeight, nil))
				continue
			}
			if _, err := net.ResolveUnixAddr("unix", hp); err == nil {
				ins = append(ins, discovery.NewInstance("unix", hp, discovery.DefaultWeight, nil))
				continue
			}
			panic(fmt.Errorf("WithHostPorts: invalid '%s'", hp))
		}

		if len(ins) == 0 {
			panic("WithHostPorts() requires at least one argument")
		}

		o.Targets = strings.Join(hostports, ",")
		o.Resolver = &discovery.SynthesizedResolver{
			ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
				return discovery.Result{
					Cacheable: true,
					CacheKey:  "fixed",
					Instances: ins,
				}, nil
			},
			NameFunc: func() string { return o.Targets },
			TargetFunc: func(ctx context.Context, target rpcinfo.EndpointInfo) string {
				return o.Targets
			},
		}
	}}
}

// WithResolver provides the Resolver for kitex client.
func WithResolver(r discovery.Resolver) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithResolver(%T)", r))

		o.Resolver = r
	}}
}

// WithHTTPResolver specifies resolver for url (which specified by WithURL).
func WithHTTPResolver(r http.Resolver) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithHTTPResolver(%T)", r))

		o.HTTPResolver = r
	}}
}

// WithShortConnection forces kitex to close connection after each call is finished.
func WithShortConnection() Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push("WithShortConnection")

		o.PoolCfg = new(connpool.IdleConfig)
	}}
}

// WithLongConnection enables long connection with kitex's built-in pooling implementation.
func WithLongConnection(cfg connpool.IdleConfig) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithLongConnection(%+v)", cfg))

		o.PoolCfg = connpool.CheckPoolConfig(cfg)
	}}
}

// WithMuxConnection specifies the transport type to be mux.
func WithMuxConnection(connNum int) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push("WithMuxConnection")

		o.RemoteOpt.ConnPool = netpollmux.NewMuxConnPool(connNum)
		WithTransHandlerFactory(netpollmux.NewCliTransHandlerFactory()).F(o, di)
		rpcinfo.AsMutableRPCConfig(o.Configs).SetTransportProtocol(transport.TTHeader)
	}}
}

// WithLogger sets the Logger for kitex client.
// Deprecated: client uses the global klog.DefaultLogger.
func WithLogger(logger klog.FormatLogger) Option {
	panic("client.WithLogger is deprecated")
}

// WithLoadBalancer sets the loadbalancer for client.
func WithLoadBalancer(lb loadbalance.Loadbalancer, opts ...*lbcache.Options) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithLoadBalancer(%+v, %+v)", lb, opts))
		o.Balancer = lb
		if len(opts) > 0 {
			o.BalancerCacheOpt = opts[0]
		}
	}}
}

// WithRPCTimeout specifies the RPC timeout.
func WithRPCTimeout(d time.Duration) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithRPCTimeout(%dms)", d.Milliseconds()))

		rpcinfo.AsMutableRPCConfig(o.Configs).SetRPCTimeout(d)
		o.Locks.Bits |= rpcinfo.BitRPCTimeout
	}}
}

// WithConnectTimeout specifies the connection timeout.
func WithConnectTimeout(d time.Duration) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithConnectTimeout(%dms)", d.Milliseconds()))

		rpcinfo.AsMutableRPCConfig(o.Configs).SetConnectTimeout(d)
		o.Locks.Bits |= rpcinfo.BitConnectTimeout
	}}
}

// WithTimeoutProvider adds a TimeoutProvider to the client.
// Note that the timeout settings provided by the TimeoutProvider
// will be applied before the other timeout options in this package
// and those in the callopt package. Thus it can not modify the
// timeouts set by WithRPCTimeout or WithConnectTimeout.
func WithTimeoutProvider(p rpcinfo.TimeoutProvider) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithTimeoutProvider(%T(%+v))", p, p))
		o.Timeouts = p
	}}
}

// WithTag sets the customize tag for service discovery, eg: idc, cluster.
func WithTag(key, val string) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithTag(%s=%s)", key, val))

		o.Svr.Tags[key] = val
		o.Locks.Tags[key] = struct{}{}
	}}
}

// WithTracer adds a tracer to client.
func WithTracer(c stats.Tracer) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithTracer(%T{%+v})", c, c))

		if o.TracerCtl == nil {
			o.TracerCtl = &internal_stats.Controller{}
		}
		o.TracerCtl.Append(c)
	}}
}

// WithStatsLevel sets the stats level for client.
func WithStatsLevel(level stats.Level) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithStatsLevel(%+v)", level))
		l := level
		o.StatsLevel = &l
	}}
}

// WithCodec to set a codec that handle other protocols which not support by kitex
func WithCodec(c remote.Codec) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithCodec(%+v)", c))

		o.RemoteOpt.Codec = c
	}}
}

// WithPayloadCodec to set a payloadCodec that handle other payload which not support by kitex
func WithPayloadCodec(c remote.PayloadCodec) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithPayloadCodec(%+v)", c))

		o.RemoteOpt.PayloadCodec = c
	}}
}

// WithConnReporterEnabled to enable reporting connection pool stats.
func WithConnReporterEnabled() Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push("WithConnReporterEnabled()")

		o.RemoteOpt.EnableConnPoolReporter = true
	}}
}

// WithFailureRetry sets the failure retry policy for client.
func WithFailureRetry(p *retry.FailurePolicy) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		if p == nil {
			return
		}
		di.Push(fmt.Sprintf("WithFailureRetry(%+v)", *p))
		if o.RetryPolicy == nil {
			o.RetryPolicy = &retry.Policy{}
		}
		if o.RetryPolicy.BackupPolicy != nil {
			panic("BackupPolicy has been setup, cannot support Failure Retry at same time")
		}
		o.RetryPolicy.FailurePolicy = p
		o.RetryPolicy.Enable = true
		o.RetryPolicy.Type = retry.FailureType
	}}
}

// WithBackupRequest sets the backup request policy for client.
func WithBackupRequest(p *retry.BackupPolicy) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		if p == nil {
			return
		}
		di.Push(fmt.Sprintf("WithBackupRequest(%+v)", *p))
		if o.RetryPolicy == nil {
			o.RetryPolicy = &retry.Policy{}
		}
		if o.RetryPolicy.FailurePolicy != nil {
			panic("Failure Retry has been setup, cannot support Backup Request at same time")
		}
		o.RetryPolicy.BackupPolicy = p
		o.RetryPolicy.Enable = true
		o.RetryPolicy.Type = retry.BackupType
	}}
}

// WithCircuitBreaker adds a circuitbreaker suite for the client.
func WithCircuitBreaker(s *circuitbreak.CBSuite) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push("WithCircuitBreaker()")
		o.CBSuite = s
	}}
}

// WithGRPCConnPoolSize sets the value for the client connection pool size.
// In general, you should not adjust the size of the connection pool, otherwise it may cause performance degradation.
// You should adjust the size according to the actual situation.
func WithGRPCConnPoolSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCConnPoolSize(%d)", s))
		o.GRPCConnPoolSize = s
	}}
}

// WithGRPCWriteBufferSize determines how much data can be batched before doing a
// write on the wire. The corresponding memory allocation for this buffer will
// be twice the size to keep syscalls low. The default value for this buffer is
// 32KB.
//
// Zero will disable the write buffer such that each write will be on underlying
// connection. Note: A Send call may not directly translate to a write.
// It corresponds to the WithWriteBufferSize DialOption of gRPC.
func WithGRPCWriteBufferSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCWriteBufferSize(%d)", s))
		o.GRPCConnectOpts.WriteBufferSize = s
	}}
}

// WithGRPCReadBufferSize lets you set the size of read buffer, this determines how
// much data can be read at most for each read syscall.
//
// The default value for this buffer is 32KB. Zero will disable read buffer for
// a connection so data framer can access the underlying conn directly.
// It corresponds to the WithReadBufferSize DialOption of gRPC.
func WithGRPCReadBufferSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCReadBufferSize(%d)", s))
		o.GRPCConnectOpts.ReadBufferSize = s
	}}
}

// WithGRPCInitialWindowSize sets the value for initial window size on a grpc stream.
// The lower bound for window size is 64K and any value smaller than that will be ignored.
// It corresponds to the WithInitialWindowSize DialOption of gRPC.
func WithGRPCInitialWindowSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCInitialWindowSize(%d)", s))
		o.GRPCConnectOpts.InitialWindowSize = s
	}}
}

// WithGRPCInitialConnWindowSize sets the value for initial window size on a connection of grpc.
// The lower bound for window size is 64K and any value smaller than that will be ignored.
// It corresponds to the WithInitialConnWindowSize DialOption of gRPC.
func WithGRPCInitialConnWindowSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCInitialConnWindowSize(%d)", s))
		o.GRPCConnectOpts.InitialConnWindowSize = s
	}}
}

// WithGRPCMaxHeaderListSize returns a DialOption that specifies the maximum
// (uncompressed) size of header list that the client is prepared to accept.
// It corresponds to the WithMaxHeaderListSize DialOption of gRPC.
func WithGRPCMaxHeaderListSize(s uint32) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCMaxHeaderListSize(%d)", s))
		o.GRPCConnectOpts.MaxHeaderListSize = &s
	}}
}

// WithGRPCKeepaliveParams returns a DialOption that specifies keepalive parameters for the client transport.
// It corresponds to the WithKeepaliveParams DialOption of gRPC.
func WithGRPCKeepaliveParams(kp grpc.ClientKeepalive) Option {
	if kp.Time < grpc.KeepaliveMinPingTime {
		kp.Time = grpc.KeepaliveMinPingTime
	}
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCKeepaliveParams(%+v)", kp))
		o.GRPCConnectOpts.KeepaliveParams = kp
	}}
}

// WithWarmingUp forces the client to do some warm-ups at the end of the initialization.
func WithWarmingUp(wuo *warmup.ClientOption) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithWarmingUp(%+v)", wuo))
		o.WarmUpOption = wuo
	}}
}