client/option_advanced.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

// Notice!! This file defines the advanced Options of client, normal user should not use it.
// It is used for customized extension.

import (
	"fmt"

	"github.com/cloudwego/kitex/internal/client"
	"github.com/cloudwego/kitex/pkg/acl"
	"github.com/cloudwego/kitex/pkg/diagnosis"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/proxy"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
	"github.com/cloudwego/kitex/pkg/retry"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/transport"
)

// WithHTTPConnection specifies client use RPC over http.
func WithHTTPConnection() Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push("WithHTTPConnection")

		WithTransHandlerFactory(netpoll.NewHTTPCliTransHandlerFactory()).F(o, di)
	}}
}

// WithClientBasicInfo provides initial information for client endpoint in RPCInfo.
func WithClientBasicInfo(ebi *rpcinfo.EndpointBasicInfo) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithClientBasicInfo(%+v)", ebi))
		if ebi != nil {
			o.Cli = ebi
		}
	}}
}

// WithDiagnosisService sets the diagnosis service for gathering debug information.
func WithDiagnosisService(ds diagnosis.Service) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithDiagnosisService(%+v)", ds))
		o.DebugService = ds
	}}
}

// WithACLRules adds ACL rules.
// Note that the ACL checking process happens before service discovery.
func WithACLRules(rules ...acl.RejectFunc) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		var names []string
		for _, r := range rules {
			names = append(names, utils.GetFuncName(r))
		}
		di.Push(fmt.Sprintf("WithACLRules(%+v)", names))

		o.ACLRules = append(o.ACLRules, rules...)
	}}
}

// WithFirstMetaHandler adds a MetaHandler.
func WithFirstMetaHandler(h remote.MetaHandler) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithFirstMetaHandler(%T)", h))
		handlers := []remote.MetaHandler{h}
		o.MetaHandlers = append(handlers, o.MetaHandlers...)
	}}
}

// WithMetaHandler adds a MetaHandler.
func WithMetaHandler(h remote.MetaHandler) client.Option {
	return client.Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMetaHandler(%T)", h))

		o.MetaHandlers = append(o.MetaHandlers, h)
	}}
}

// WithProxy sets the forward Proxy for client.
func WithProxy(p proxy.ForwardProxy) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithProxy(%T)", p))

		if o.Proxy != nil {
			panic(fmt.Errorf("reassignment of Proxy is not allowed: %T -> %T", o.Proxy, p))
		}
		o.Proxy = p
	}}
}

// WithTransHandlerFactory sets the TransHandlerFactory for client.
func WithTransHandlerFactory(f remote.ClientTransHandlerFactory) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithTransHandlerFactory(%T)", f))

		o.RemoteOpt.CliHandlerFactory = f
	}}
}

// WithDialer sets the Dialer for creating connections.
func WithDialer(d remote.Dialer) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithDialer(%T)", d))

		if d == nil {
			panic("invalid Dialer: nil")
		}
		o.RemoteOpt.Dialer = d
	}}
}

// WithConnPool sets the connection pool.
// Note that this option can only be specified once. If more
// than one pool is specified by this option, only the first
// one will be used.
func WithConnPool(pool remote.ConnPool) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithConnPool((%T)", pool))

		if o.RemoteOpt.ConnPool == nil {
			o.RemoteOpt.ConnPool = pool
		} else {
			klog.Warnf("The connection pool has been initialized. The call to WithConnPool will not take effect.")
		}
	}}
}

// WithRetryContainer sets Container
func WithRetryContainer(rc *retry.Container) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push("WithRetryContainer()")

		if rc == nil {
			panic("invalid retry container: nil")
		}
		o.RetryContainer = rc
	}}
}

// WithGeneric set Generic type for generic call
func WithGeneric(g generic.Generic) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithGeneric(%T)", g))

		if g == nil {
			panic("invalid Generic: nil")
		}
		if g.Framed() {
			rpcinfo.AsMutableRPCConfig(o.Configs).SetTransportProtocol(transport.Framed)
		}
		o.RemoteOpt.PayloadCodec = g.PayloadCodec()
	}}
}

// WithCloseCallbacks adds callback to Close
func WithCloseCallbacks(callback func() error) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithCloseCallbacks(%+v)", utils.GetFuncName(callback)))

		if callback == nil {
			panic("invalid Close Callback: nil")
		}
		o.CloseCallbacks = append(o.CloseCallbacks, callback)
	}}
}

// WithErrorHandler sets the error handler.
func WithErrorHandler(f func(error) error) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithErrorHandler(%+v)", utils.GetFuncName(f)))

		o.ErrHandle = f
	}}
}

// WithBoundHandler adds remote.BoundHandler for client.
func WithBoundHandler(h remote.BoundHandler) Option {
	return Option{F: func(o *client.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("AddBoundHandler(%T)", h))

		o.RemoteOpt.AppendBoundHandler(h)
	}}
}