server/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 server

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/cloudwego/localsession/backup"

	internal_server "github.com/cloudwego/kitex/internal/server"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/limit"
	"github.com/cloudwego/kitex/pkg/limiter"
	"github.com/cloudwego/kitex/pkg/registry"
	"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/trans/netpollmux"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/streaming"
	"github.com/cloudwego/kitex/pkg/utils"
)

// Option is the only way to config server.
type Option = internal_server.Option

// Options is used to initialize the server.
type Options = internal_server.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
}

// WithSuite adds an option suite for server.
func WithSuite(suite Suite) Option {
	return Option{F: func(o *internal_server.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)
	}}
}

// WithMuxTransport specifies the transport type to be mux.
func WithMuxTransport() Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push("WithMuxTransport()")
		o.RemoteOpt.SvrHandlerFactory = netpollmux.NewSvrTransHandlerFactory()
		// set limit options
		o.Limit.QPSLimitPostDecode = true
	}}
}

// WithMiddleware adds middleware for server to handle request.
func WithMiddleware(mw endpoint.Middleware) Option {
	mwb := func(ctx context.Context) endpoint.Middleware {
		return mw
	}
	return Option{F: func(o *internal_server.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 server to handle request
func WithMiddlewareBuilder(mwb endpoint.MiddlewareBuilder, funcName ...string) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMiddlewareBuilder(%+v)", utils.GetFuncName(mwb)))
		o.MWBs = append(o.MWBs, mwb)
	}}
}

// WithReadWriteTimeout sets the read/write timeout on network.
// IMPORTANT: this option is not stable, and will be changed or removed in the future!!!
// We don't promise compatibility for this option in future versions!!!
func WithReadWriteTimeout(d time.Duration) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithReadWriteTimeout(%v)", d))

		rpcinfo.AsMutableRPCConfig(o.Configs).SetReadWriteTimeout(d)
		o.LockBits |= rpcinfo.BitReadWriteTimeout
	}}
}

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

// WithExitWaitTime sets the wait duration for graceful shutdown.
func WithExitWaitTime(timeout time.Duration) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithExitWaitTime(%+v)", timeout))
		if timeout > 0 {
			o.RemoteOpt.ExitWaitTime = timeout
		}
	}}
}

// WithMaxConnIdleTime sets the max idle time on connection from clients.
func WithMaxConnIdleTime(timeout time.Duration) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithMaxConnIdleTime(%+v)", timeout))
		if timeout > 0 {
			o.RemoteOpt.MaxConnectionIdleTime = timeout
		}
	}}
}

// WithLimit sets the limitation of concurrent connections or max QPS.
// IMPORTANT: this option is not stable, and will be changed or removed in the future!!!
// We don't promise compatibility for this option in future versions!!!
func WithLimit(lim *limit.Option) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithLimit(%+v)", lim))

		o.Limit.Limits = lim
	}}
}

// WithConnectionLimiter sets the limiter of connections.
// If both WithLimit and WithConnectionLimiter are called, only the latter will take effect.
func WithConnectionLimiter(conLimit limiter.ConcurrencyLimiter) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithConnectionLimiter(%T{%+v})", conLimit, conLimit))

		o.Limit.ConLimit = conLimit
	}}
}

// WithQPSLimiter sets the limiter of max QPS.
// If both WithLimit and WithQPSLimiter are called, only the latter will take effect.
func WithQPSLimiter(qpsLimit limiter.RateLimiter) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithQPSLimiter(%T{%+v})", qpsLimit, qpsLimit))

		o.Limit.QPSLimit = qpsLimit
	}}
}

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

		if o.TracerCtl == nil {
			o.TracerCtl = &rpcinfo.TraceController{}
		}
		o.TracerCtl.Append(c)
	}}
}

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

// WithServiceAddr sets the listen address for server.
func WithServiceAddr(addr net.Addr) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithServiceAddr(%+v)", addr))

		o.RemoteOpt.Address = addr
	}}
}

// 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 *internal_server.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 *internal_server.Options, di *utils.Slice) {
		if thrift.IsThriftCodec(c) {
			// default thriftCodec has been registered,
			// if using NewThriftCodecWithConfig to set codec mode, just replace the registered one
			di.Push(fmt.Sprintf("ResetThriftPayloadCodec(%+v)", c))
			remote.PutPayloadCode(serviceinfo.Thrift, c)
		} else if protobuf.IsProtobufCodec(c) {
			di.Push(fmt.Sprintf("ResetProtobufPayloadCodec(%+v)", c))
			remote.PutPayloadCode(serviceinfo.Protobuf, c)
		} else {
			di.Push(fmt.Sprintf("WithPayloadCodec(%+v)", c))
			// if specify RemoteOpt.PayloadCodec, then the priority is highest, all payload decode will use this one
			o.RemoteOpt.PayloadCodec = c
		}
	}}
}

// WithRegistry to set a Registry to register service
func WithRegistry(r registry.Registry) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithRegistry(%T)", r))

		o.Registry = r
	}}
}

// WithRegistryInfo to set Registry Info if needed.
// It is used to add customized info and is suggested to use with WithRegistry option.
func WithRegistryInfo(info *registry.Info) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithRegistryInfo(%+v)", info))

		o.RegistryInfo = info
	}}
}

// 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 WriteBufferSize ServerOption of gRPC.
func WithGRPCWriteBufferSize(s uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCWriteBufferSize(%+v)", s))

		o.RemoteOpt.GRPCCfg.WriteBufferSize = s
	}}
}

// WithGRPCReadBufferSize lets you set the size of read buffer, this determines how much data can be read at most
// for one 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 ReadBufferSize ServerOption of gRPC.
func WithGRPCReadBufferSize(s uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCReadBufferSize(%+v)", s))

		o.RemoteOpt.GRPCCfg.ReadBufferSize = s
	}}
}

// WithGRPCInitialWindowSize returns a Option that sets window size for stream.
// The lower bound for window size is 64K and any value smaller than that will be ignored.
// It corresponds to the InitialWindowSize ServerOption of gRPC.
func WithGRPCInitialWindowSize(s uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithRegistryInfo(%+v)", s))

		o.RemoteOpt.GRPCCfg.InitialWindowSize = s
	}}
}

// WithGRPCInitialConnWindowSize returns an Option that sets window size for a connection.
// The lower bound for window size is 64K and any value smaller than that will be ignored.
// It corresponds to the InitialConnWindowSize ServerOption of gRPC.
func WithGRPCInitialConnWindowSize(s uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCInitialConnWindowSize(%+v)", s))

		o.RemoteOpt.GRPCCfg.InitialConnWindowSize = s
	}}
}

// WithGRPCKeepaliveParams returns an Option that sets keepalive and max-age parameters for the server.
// It corresponds to the KeepaliveParams ServerOption of gRPC.
func WithGRPCKeepaliveParams(kp grpc.ServerKeepalive) Option {
	if kp.Time > 0 && kp.Time < time.Second {
		kp.Time = time.Second
	}

	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCKeepaliveParams(%+v)", kp))

		o.RemoteOpt.GRPCCfg.KeepaliveParams = kp
	}}
}

// WithGRPCKeepaliveEnforcementPolicy returns an Option that sets keepalive enforcement policy for the server.
// It corresponds to the KeepaliveEnforcementPolicy ServerOption of gRPC.
func WithGRPCKeepaliveEnforcementPolicy(kep grpc.EnforcementPolicy) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCKeepaliveEnforcementPolicy(%+v)", kep))

		o.RemoteOpt.GRPCCfg.KeepaliveEnforcementPolicy = kep
	}}
}

// WithGRPCMaxConcurrentStreams returns an Option that will apply a limit on the number
// of concurrent streams to each ServerTransport.
// It corresponds to the MaxConcurrentStreams ServerOption of gRPC.
func WithGRPCMaxConcurrentStreams(n uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCMaxConcurrentStreams(%+v)", n))

		o.RemoteOpt.GRPCCfg.MaxStreams = n
	}}
}

// WithGRPCMaxHeaderListSize returns a ServerOption that sets the max (uncompressed) size
// of header list that the server is prepared to accept.
// It corresponds to the MaxHeaderListSize ServerOption of gRPC.
func WithGRPCMaxHeaderListSize(s uint32) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCMaxHeaderListSize(%+v)", s))

		o.RemoteOpt.GRPCCfg.MaxHeaderListSize = &s
	}}
}

func WithGRPCUnknownServiceHandler(f func(ctx context.Context, methodName string, stream streaming.Stream) error) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithGRPCUnknownServiceHandler(%+v)", utils.GetFuncName(f)))
		o.RemoteOpt.GRPCUnknownServiceHandler = f
	}}
}

// Deprecated: Use WithConnectionLimiter instead.
func WithConcurrencyLimiter(conLimit limiter.ConcurrencyLimiter) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithConcurrencyLimiter(%T{%+v})", conLimit, conLimit))

		o.Limit.ConLimit = conLimit
	}}
}

// WithContextBackup enables local-session to backup kitex server's context,
// in case of user don't correctly pass context into next RPC call.
//   - enable means enable context backup
//   - async means backup session will propagate to children goroutines, otherwise it only works on the same goroutine handler
func WithContextBackup(enable, async bool) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		di.Push(fmt.Sprintf("WithLocalSession({%+v, %+v})", enable, async))
		o.BackupOpt = backup.DefaultOptions()
		o.BackupOpt.Enable = enable
		o.BackupOpt.EnableImplicitlyTransmitAsync = async
	}}
}

// WithRefuseTrafficWithoutServiceName returns an Option that only accepts traffics with service name.
// This is used for a server with multi services and is one of the options to avoid a server startup error
// when having conflicting method names between services without specifying a fallback service for the method.
func WithRefuseTrafficWithoutServiceName() Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.RefuseTrafficWithoutServiceName = true
	}}
}

// WithEnableContextTimeout enables handler timeout.
// Available since Kitex >= v0.9.0
// If enabled, a timeout middleware will be added to the beginning of the middleware chain.
// The timeout value will be read from RPCInfo.Config().RPCTimeout(), which can be set by a custom MetaHandler
// NOTE:
// If there's an error (excluding BizStatusError) returned by server handler or any middleware, cancel will be
// called automatically.
//
// For an opensource Kitex user, TTHeader has builtin support of timeout-passing (not enabled by default):
//   - Client side: add the following NewClient options for enabling TTHeader and setting the timeout to TTHeader
//     client.WithTransportProtocol(transport.TTHeader),
//     client.WithMetaHandler(transmeta.ClientTTHeaderHandler),
//   - Server side: add the following NewServer options for reading from TTHeader and enable timeout control
//     server.WithMetaHandler(transmeta.ServerTTHeaderHandler)
//     server.WithEnableContextTimeout(true)
//
// For requests on GRPC transport, a deadline will be added to the context if the header 'grpc-timeout' is positive,
// so there's no need to use this option.
func WithEnableContextTimeout(enable bool) Option {
	return Option{F: func(o *internal_server.Options, di *utils.Slice) {
		o.EnableContextTimeout = enable
	}}
}