server/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 server defines the Options of server
package server

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

import (
	"context"
	"fmt"
	"net"
	"reflect"

	internal_server "github.com/cloudwego/kitex/internal/server"
	"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/limiter"
	"github.com/cloudwego/kitex/pkg/profiler"
	"github.com/cloudwego/kitex/pkg/proxy"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

// WithServerBasicInfo provides initial information for client endpoint in RPCInfo.
func WithServerBasicInfo(ebi *rpcinfo.EndpointBasicInfo) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithServerBasicInfo(%+v)", ebi))

		if ebi != nil {
			o.Svr = ebi
		}
	}}
}

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

		o.DebugService = ds
	}}
}

// WithACLRules sets the ACL rules.
func WithACLRules(rules ...acl.RejectFunc) Option {
	return Option{F: func(o *internal_server.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...)
	}}
}

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

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

// WithProxy sets the backward Proxy for server.
func WithProxy(p proxy.ReverseProxy) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		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 server.
func WithTransHandlerFactory(f remote.ServerTransHandlerFactory) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithTransHandlerFactory(%T)", f))

		o.RemoteOpt.SvrHandlerFactory = f
	}}
}

// WithTransServerFactory sets the TransServerFactory for server.
func WithTransServerFactory(f remote.TransServerFactory) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithTransServerFactory(%T)", f))

		o.RemoteOpt.TransServerFactory = f
	}}
}

// WithLimitReporter do report when server limit happen
func WithLimitReporter(r limiter.LimitReporter) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.Once.OnceOrPanic()
		di.Push(fmt.Sprintf("WithLimitReporter(%T)", r))

		o.Limit.LimitReporter = r
	}}
}

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

		if g == nil {
			panic("invalid Generic: nil")
		}
		o.RemoteOpt.PayloadCodec = g.PayloadCodec()
	}}
}

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

		o.ErrHandle = f
	}}
}

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

		exist := false
		switch handler := h.(type) {
		case remote.InboundHandler:
			for _, inboundHandler := range o.RemoteOpt.Inbounds {
				if reflect.DeepEqual(inboundHandler, handler) {
					exist = true
					break
				}
			}
		case remote.OutboundHandler:
			for _, outboundHandler := range o.RemoteOpt.Outbounds {
				if reflect.DeepEqual(outboundHandler, handler) {
					exist = true
					break
				}
			}
		}
		// prevent duplication
		if !exist {
			doAddBoundHandler(h, o.RemoteOpt)
		} else {
			klog.Warnf("KITEX: BoundHandler already exists, BoundHandler=%v", h)
		}
	}}
}

// WithExitSignal adds ExitSignal for server.
func WithExitSignal(f func() <-chan error) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("AddExitSignal(%+v)", utils.GetFuncName(f)))
		o.ExitSignal = f
	}}
}

// WithListener sets the listener for server, the priority is higher than WithServiceAddr
func WithListener(ln net.Listener) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithListener(%+v)", ln))

		o.RemoteOpt.Listener = ln
	}}
}

// WithReusePort sets SO_REUSEPORT on listener, it is only used with Option `WithServiceAddr`.
// It won't take effect when listener is specified by WithListener.
func WithReusePort(reuse bool) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithReusePort(%+v)", reuse))

		o.RemoteOpt.ReusePort = reuse
	}}
}

// WithSupportedTransportsFunc sets a function which converts supported transports from server option.
func WithSupportedTransportsFunc(f func(option remote.ServerOption) []string) Option {
	return Option{
		F: func(o *internal_server.Options, di *utils.Slice) {
			di.Push("WithSupportedTransportsFunc()")
			o.SupportedTransportsFunc = f
		},
	}
}

// WithProfiler set a profiler to server.
func WithProfiler(pc profiler.Profiler) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithProfiler(%T{%+v})", pc, pc))
		o.RemoteOpt.Profiler = pc
	}}
}

// WithProfilerTransInfoTagging set transinfo tagging function to profiler
// TransInfoTagging extracting tags after TransInfo decoded but before message decoded.
// At this stage, we can only get msg.TransInfo() and the real message payload is not decoded yet.
// If upstream is not use TTHeader protocol, we can get nothing here.
// So if you don't very care about the accuracy of statistics, we recommend to use WithProfilerMessageTagging to extract your custom tags.
func WithProfilerTransInfoTagging(tagging remote.TransInfoTagging) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithProfilerTransInfoTagging(%+v)", utils.GetFuncName(tagging)))
		interTagging := o.RemoteOpt.ProfilerTransInfoTagging
		if interTagging == nil {
			o.RemoteOpt.ProfilerTransInfoTagging = tagging
			return
		}
		o.RemoteOpt.ProfilerTransInfoTagging = func(ctx context.Context, msg remote.Message) (context.Context, []string) {
			c, t := tagging(ctx, msg)
			c2, t2 := interTagging(c, msg)
			return c2, append(t, t2...)
		}
	}}
}

// WithProfilerMessageTagging set message tagging function to profiler
// MessageTagging extracting tags after whole decode process finished.
// At this stage, we can get the rpcInfo from ctx, and full complete message.
func WithProfilerMessageTagging(tagging remote.MessageTagging) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithProfilerMessageTagging(%+v)", utils.GetFuncName(tagging)))
		interTagging := o.RemoteOpt.ProfilerMessageTagging
		if interTagging == nil {
			o.RemoteOpt.ProfilerMessageTagging = tagging
			return
		}
		o.RemoteOpt.ProfilerMessageTagging = func(ctx context.Context, msg remote.Message) (context.Context, []string) {
			c, t := tagging(ctx, msg)
			c2, t2 := interTagging(c, msg)
			return c2, append(t, t2...)
		}
	}}
}