github.com/cloudwego/kitex@v0.9.0/internal/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 defines the Options of client
package client

import (
	"context"
	"time"

	"github.com/cloudwego/localsession/backup"

	"github.com/cloudwego/kitex/internal/configutil"
	"github.com/cloudwego/kitex/internal/stream"
	"github.com/cloudwego/kitex/pkg/acl"
	"github.com/cloudwego/kitex/pkg/circuitbreak"
	connpool2 "github.com/cloudwego/kitex/pkg/connpool"
	"github.com/cloudwego/kitex/pkg/diagnosis"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/event"
	"github.com/cloudwego/kitex/pkg/fallback"
	"github.com/cloudwego/kitex/pkg/http"
	"github.com/cloudwego/kitex/pkg/loadbalance"
	"github.com/cloudwego/kitex/pkg/loadbalance/lbcache"
	"github.com/cloudwego/kitex/pkg/proxy"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec/protobuf"
	"github.com/cloudwego/kitex/pkg/remote/codec/thrift"
	"github.com/cloudwego/kitex/pkg/remote/connpool"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2"
	"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/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/transmeta"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/pkg/warmup"
	"github.com/cloudwego/kitex/transport"
)

func init() {
	remote.PutPayloadCode(serviceinfo.Thrift, thrift.NewThriftCodec())
	remote.PutPayloadCode(serviceinfo.Protobuf, protobuf.NewProtobufCodec())
}

// Options is used to initialize a client.
type Options struct {
	Cli     *rpcinfo.EndpointBasicInfo
	Svr     *rpcinfo.EndpointBasicInfo
	Configs rpcinfo.RPCConfig
	Locks   *ConfigLocks
	Once    *configutil.OptionOnce

	MetaHandlers []remote.MetaHandler

	RemoteOpt        *remote.ClientOption
	Proxy            proxy.ForwardProxy
	Resolver         discovery.Resolver
	HTTPResolver     http.Resolver
	Balancer         loadbalance.Loadbalancer
	BalancerCacheOpt *lbcache.Options
	PoolCfg          *connpool2.IdleConfig
	ErrHandle        func(context.Context, error) error
	Targets          string
	CBSuite          *circuitbreak.CBSuite
	Timeouts         rpcinfo.TimeoutProvider

	ACLRules []acl.RejectFunc

	MWBs  []endpoint.MiddlewareBuilder
	IMWBs []endpoint.MiddlewareBuilder

	Bus          event.Bus
	Events       event.Queue
	ExtraTimeout time.Duration

	// DebugInfo should only contain objects that are suitable for json serialization.
	DebugInfo    utils.Slice
	DebugService diagnosis.Service

	// Observability
	TracerCtl  *rpcinfo.TraceController
	StatsLevel *stats.Level

	// retry policy
	RetryMethodPolicies map[string]retry.Policy
	RetryContainer      *retry.Container
	RetryWithResult     *retry.ShouldResultRetry

	// fallback policy
	Fallback *fallback.Policy

	CloseCallbacks []func() error
	WarmUpOption   *warmup.ClientOption

	// GRPC
	GRPCConnPoolSize uint32
	GRPCConnectOpts  *grpc.ConnectOptions

	// XDS
	XDSEnabled          bool
	XDSRouterMiddleware endpoint.Middleware

	// Context backup
	CtxBackupHandler backup.BackupHandler

	Streaming stream.StreamingConfig
}

// Apply applies all options.
func (o *Options) Apply(opts []Option) {
	for _, op := range opts {
		op.F(o, &o.DebugInfo)
	}
}

// Option is the only way to config client.
type Option struct {
	F func(o *Options, di *utils.Slice)
}

// NewOptions creates a new option.
func NewOptions(opts []Option) *Options {
	o := &Options{
		Cli:          &rpcinfo.EndpointBasicInfo{Tags: make(map[string]string)},
		Svr:          &rpcinfo.EndpointBasicInfo{Tags: make(map[string]string)},
		MetaHandlers: []remote.MetaHandler{transmeta.MetainfoClientHandler},
		RemoteOpt:    newClientRemoteOption(),
		Configs:      rpcinfo.NewRPCConfig(),
		Locks:        NewConfigLocks(),
		Once:         configutil.NewOptionOnce(),
		HTTPResolver: http.NewDefaultResolver(),
		DebugService: diagnosis.NoopService,

		Bus:    event.NewEventBus(),
		Events: event.NewQueue(event.MaxEventNum),

		TracerCtl: &rpcinfo.TraceController{},

		GRPCConnectOpts: new(grpc.ConnectOptions),
	}
	o.Apply(opts)

	o.initRemoteOpt()

	if o.RetryContainer != nil && o.DebugService != nil {
		o.DebugService.RegisterProbeFunc(diagnosis.RetryPolicyKey, o.RetryContainer.Dump)
	}

	if o.StatsLevel == nil {
		level := stats.LevelDisabled
		if o.TracerCtl.HasTracer() {
			level = stats.LevelDetailed
		}
		o.StatsLevel = &level
	}
	return o
}

func (o *Options) initRemoteOpt() {
	var zero connpool2.IdleConfig

	if o.Configs.TransportProtocol()&transport.GRPC == transport.GRPC {
		if o.PoolCfg != nil && *o.PoolCfg == zero {
			// grpc unary short connection
			o.GRPCConnectOpts.ShortConn = true
		}
		o.RemoteOpt.ConnPool = nphttp2.NewConnPool(o.Svr.ServiceName, o.GRPCConnPoolSize, *o.GRPCConnectOpts)
		o.RemoteOpt.CliHandlerFactory = nphttp2.NewCliTransHandlerFactory()
	}
	if o.RemoteOpt.ConnPool == nil {
		if o.PoolCfg != nil {
			if *o.PoolCfg == zero {
				o.RemoteOpt.ConnPool = connpool.NewShortPool(o.Svr.ServiceName)
			} else {
				o.RemoteOpt.ConnPool = connpool.NewLongPool(o.Svr.ServiceName, *o.PoolCfg)
			}
		} else {
			o.RemoteOpt.ConnPool = connpool.NewLongPool(
				o.Svr.ServiceName,
				connpool2.IdleConfig{
					MaxIdlePerAddress: 10,
					MaxIdleGlobal:     100,
					MaxIdleTimeout:    time.Minute,
				},
			)
		}
	}
}

// InitRetryContainer init retry container and add close callback
func (o *Options) InitRetryContainer() {
	if o.RetryContainer == nil {
		o.RetryContainer = retry.NewRetryContainerWithPercentageLimit()
		o.CloseCallbacks = append(o.CloseCallbacks, o.RetryContainer.Close)
	}
}