github.com/cloudwego/kitex@v0.9.0/pkg/rpcinfo/rpcconfig.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 rpcinfo

import (
	"sync"
	"time"

	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/transport"
)

var (
	_             MutableRPCConfig = &rpcConfig{}
	_             RPCConfig        = &rpcConfig{}
	rpcConfigPool sync.Pool
)

// default values.
var (
	defaultRPCTimeout       = time.Duration(0)
	defaultConnectTimeout   = time.Millisecond * 50
	defaultReadWriteTimeout = time.Second * 5
	defaultBufferSize       = 4096
	defaultInteractionMode  = PingPong
)

// Mask bits.
const (
	BitRPCTimeout = 1 << iota
	BitConnectTimeout
	BitReadWriteTimeout
	BitIOBufferSize
)

type InteractionMode int32

const (
	PingPong  InteractionMode = 0
	Oneway    InteractionMode = 1
	Streaming InteractionMode = 2
)

// rpcConfig is a set of configurations used during RPC calls.
type rpcConfig struct {
	readOnlyMask      int
	rpcTimeout        time.Duration
	connectTimeout    time.Duration
	readWriteTimeout  time.Duration
	ioBufferSize      int
	transportProtocol transport.Protocol
	interactionMode   InteractionMode
	payloadCodec      serviceinfo.PayloadCodec
}

func init() {
	rpcConfigPool.New = newRPCConfig
}

func newRPCConfig() interface{} {
	c := &rpcConfig{}
	c.initialize()
	return c
}

// LockConfig sets the bits of readonly mask to prevent sequential modification on certain configs.
func (r *rpcConfig) LockConfig(bits int) {
	r.readOnlyMask |= bits
}

// SetRPCTimeout implements MutableRPCConfig.
func (r *rpcConfig) SetRPCTimeout(to time.Duration) error {
	if !r.IsRPCTimeoutLocked() {
		r.rpcTimeout = to
		return nil
	}
	return kerrors.ErrNotSupported
}

// IsRPCTimeoutLocked implements the MutableRPCConfig interface.
func (r *rpcConfig) IsRPCTimeoutLocked() bool {
	return r.readOnlyMask&BitRPCTimeout != 0
}

// SetConnectTimeout implements MutableRPCConfig interface.
func (r *rpcConfig) SetConnectTimeout(to time.Duration) error {
	if !r.IsConnectTimeoutLocked() {
		r.connectTimeout = to
		return nil
	}
	return kerrors.ErrNotSupported
}

// IsConnectTimeoutLocked implements the MutableRPCConfig interface.
func (r *rpcConfig) IsConnectTimeoutLocked() bool {
	return r.readOnlyMask&BitConnectTimeout != 0
}

// SetReadWriteTimeout implements MutableRPCConfig interface.
func (r *rpcConfig) SetReadWriteTimeout(to time.Duration) error {
	if !r.IsReadWriteTimeoutLocked() {
		r.readWriteTimeout = to
		return nil
	}
	return kerrors.ErrNotSupported
}

// IsReadWriteTimeoutLocked implements the MutableRPCConfig interface.
func (r *rpcConfig) IsReadWriteTimeoutLocked() bool {
	return r.readOnlyMask&BitReadWriteTimeout != 0
}

// SetIOBufferSize implements MutableRPCConfig interface.
func (r *rpcConfig) SetIOBufferSize(sz int) error {
	if (r.readOnlyMask & BitIOBufferSize) == 0 {
		r.ioBufferSize = sz
		return nil
	}
	return kerrors.ErrNotSupported
}

// ImmutableView implements MutableRPCConfig interface.
func (r *rpcConfig) ImmutableView() RPCConfig {
	return r
}

// RPCTimeout implements RPCConfig interface.
func (r *rpcConfig) RPCTimeout() time.Duration {
	return r.rpcTimeout
}

// ConnectTimeout implements RPCConfig interface.
func (r *rpcConfig) ConnectTimeout() time.Duration {
	return r.connectTimeout
}

// ReadWriteTimeout implements RPCConfig interface.
func (r *rpcConfig) ReadWriteTimeout() time.Duration {
	return r.readWriteTimeout
}

// IOBufferSize implements RPCConfig interface.
func (r *rpcConfig) IOBufferSize() int {
	return r.ioBufferSize
}

// TransportProtocol implements RPCConfig interface. It is only useful for client.
func (r *rpcConfig) TransportProtocol() transport.Protocol {
	return r.transportProtocol
}

// SetTransportProtocol implements MutableRPCConfig interface.
func (r *rpcConfig) SetTransportProtocol(tp transport.Protocol) error {
	r.transportProtocol |= tp
	return nil
}

func (r *rpcConfig) SetInteractionMode(mode InteractionMode) error {
	r.interactionMode = mode
	return nil
}

func (r *rpcConfig) InteractionMode() InteractionMode {
	return r.interactionMode
}

func (r *rpcConfig) SetPayloadCodec(codec serviceinfo.PayloadCodec) {
	r.payloadCodec = codec
}

func (r *rpcConfig) PayloadCodec() serviceinfo.PayloadCodec {
	return r.payloadCodec
}

// Clone returns a copy of the current rpcConfig.
func (r *rpcConfig) Clone() MutableRPCConfig {
	r2 := rpcConfigPool.Get().(*rpcConfig)
	*r2 = *r
	return r2
}

func (r *rpcConfig) CopyFrom(from RPCConfig) {
	f := from.(*rpcConfig)
	*r = *f
}

func (r *rpcConfig) initialize() {
	r.readOnlyMask = 0
	r.rpcTimeout = defaultRPCTimeout
	r.connectTimeout = defaultConnectTimeout
	r.readWriteTimeout = defaultReadWriteTimeout
	r.ioBufferSize = defaultBufferSize
	r.transportProtocol = 0
	r.interactionMode = defaultInteractionMode
}

// Recycle reuses the rpcConfig.
func (r *rpcConfig) Recycle() {
	r.initialize()
	rpcConfigPool.Put(r)
}

// NewRPCConfig creates a default RPCConfig.
func NewRPCConfig() RPCConfig {
	r := rpcConfigPool.Get().(*rpcConfig)
	return r
}