rpcinfo.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 (
	"os"
	"sync"
	"sync/atomic"

	"github.com/cloudwego/kitex/internal"
)

var (
	rpcInfoPool sync.Pool
	enablePool  int32 = 1
)

func init() {
	// allow disabling by env without modifying the code and recompiling
	if os.Getenv("KITEX_DISABLE_RPCINFO_POOL") != "" {
		EnablePool(false)
	}
}

// EnablePool allows user to enable/disable rpcInfoPool.
// It's enabled by default for performance, but may cause trouble due to misuses:
//
//	referencing RPCInfo in another goroutine other than the one running the handler.
//
// By turning off the pool, we can quickly confirm whether the concurrency issues is
// caused by such cases, but do remember there's a PERFORMANCE LOSS.
func EnablePool(enable bool) {
	if enable {
		atomic.StoreInt32(&enablePool, 1)
	} else {
		atomic.StoreInt32(&enablePool, 0)
	}
}

// PoolEnabled returns true if rpcInfoPool is enabled.
func PoolEnabled() bool {
	return atomic.LoadInt32(&enablePool) == 1
}

type rpcInfo struct {
	from       EndpointInfo
	to         EndpointInfo
	invocation Invocation
	config     RPCConfig
	stats      RPCStats
}

// From implements the RPCInfo interface.
func (r *rpcInfo) From() EndpointInfo { return r.from }

// To implements the RPCInfo interface.
func (r *rpcInfo) To() EndpointInfo { return r.to }

// Config implements the RPCInfo interface.
func (r *rpcInfo) Config() RPCConfig { return r.config }

// Invocation implements the RPCInfo interface.
func (r *rpcInfo) Invocation() Invocation { return r.invocation }

// Stats implements the RPCInfo interface.
func (r *rpcInfo) Stats() RPCStats { return r.stats }

func (r *rpcInfo) zero() {
	r.from = nil
	r.to = nil
	r.invocation = nil
	r.config = nil
	r.stats = nil
}

// Recycle reuses the rpcInfo.
func (r *rpcInfo) Recycle() {
	if !PoolEnabled() {
		return
	}
	if v, ok := r.from.(internal.Reusable); ok {
		v.Recycle()
	}
	if v, ok := r.to.(internal.Reusable); ok {
		v.Recycle()
	}
	if v, ok := r.invocation.(internal.Reusable); ok {
		v.Recycle()
	}
	if v, ok := r.config.(internal.Reusable); ok {
		v.Recycle()
	}
	if v, ok := r.stats.(internal.Reusable); ok {
		v.Recycle()
	}
	r.zero()
	rpcInfoPool.Put(r)
}

func init() {
	rpcInfoPool.New = newRPCInfo
}

// NewRPCInfo creates a new RPCInfo using the given information.
func NewRPCInfo(from, to EndpointInfo, ink Invocation, config RPCConfig, stats RPCStats) RPCInfo {
	r := rpcInfoPool.Get().(*rpcInfo)
	r.from = from
	r.to = to
	r.invocation = ink
	r.config = config
	r.stats = stats
	return r
}

func newRPCInfo() interface{} {
	return &rpcInfo{}
}