github.com/cloudwego/kitex@v0.9.0/pkg/rpcinfo/rpcinfo.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package rpcinfo
    18  
    19  import (
    20  	"os"
    21  	"sync"
    22  	"sync/atomic"
    23  
    24  	"github.com/cloudwego/kitex/internal"
    25  )
    26  
    27  var (
    28  	rpcInfoPool sync.Pool
    29  	enablePool  int32 = 1
    30  )
    31  
    32  func init() {
    33  	// allow disabling by env without modifying the code and recompiling
    34  	if os.Getenv("KITEX_DISABLE_RPCINFO_POOL") != "" {
    35  		EnablePool(false)
    36  	}
    37  }
    38  
    39  // EnablePool allows user to enable/disable rpcInfoPool.
    40  // It's enabled by default for performance, but may cause trouble due to misuses:
    41  //
    42  //	referencing RPCInfo in another goroutine other than the one running the handler.
    43  //
    44  // By turning off the pool, we can quickly confirm whether the concurrency issues is
    45  // caused by such cases, but do remember there's a PERFORMANCE LOSS.
    46  func EnablePool(enable bool) {
    47  	if enable {
    48  		atomic.StoreInt32(&enablePool, 1)
    49  	} else {
    50  		atomic.StoreInt32(&enablePool, 0)
    51  	}
    52  }
    53  
    54  // PoolEnabled returns true if rpcInfoPool is enabled.
    55  func PoolEnabled() bool {
    56  	return atomic.LoadInt32(&enablePool) == 1
    57  }
    58  
    59  type rpcInfo struct {
    60  	from       EndpointInfo
    61  	to         EndpointInfo
    62  	invocation Invocation
    63  	config     RPCConfig
    64  	stats      RPCStats
    65  }
    66  
    67  // From implements the RPCInfo interface.
    68  func (r *rpcInfo) From() EndpointInfo { return r.from }
    69  
    70  // To implements the RPCInfo interface.
    71  func (r *rpcInfo) To() EndpointInfo { return r.to }
    72  
    73  // Config implements the RPCInfo interface.
    74  func (r *rpcInfo) Config() RPCConfig { return r.config }
    75  
    76  // Invocation implements the RPCInfo interface.
    77  func (r *rpcInfo) Invocation() Invocation { return r.invocation }
    78  
    79  // Stats implements the RPCInfo interface.
    80  func (r *rpcInfo) Stats() RPCStats { return r.stats }
    81  
    82  func (r *rpcInfo) zero() {
    83  	r.from = nil
    84  	r.to = nil
    85  	r.invocation = nil
    86  	r.config = nil
    87  	r.stats = nil
    88  }
    89  
    90  // Recycle reuses the rpcInfo.
    91  func (r *rpcInfo) Recycle() {
    92  	if !PoolEnabled() {
    93  		return
    94  	}
    95  	if v, ok := r.from.(internal.Reusable); ok {
    96  		v.Recycle()
    97  	}
    98  	if v, ok := r.to.(internal.Reusable); ok {
    99  		v.Recycle()
   100  	}
   101  	if v, ok := r.invocation.(internal.Reusable); ok {
   102  		v.Recycle()
   103  	}
   104  	if v, ok := r.config.(internal.Reusable); ok {
   105  		v.Recycle()
   106  	}
   107  	if v, ok := r.stats.(internal.Reusable); ok {
   108  		v.Recycle()
   109  	}
   110  	r.zero()
   111  	rpcInfoPool.Put(r)
   112  }
   113  
   114  func init() {
   115  	rpcInfoPool.New = newRPCInfo
   116  }
   117  
   118  // NewRPCInfo creates a new RPCInfo using the given information.
   119  func NewRPCInfo(from, to EndpointInfo, ink Invocation, config RPCConfig, stats RPCStats) RPCInfo {
   120  	r := rpcInfoPool.Get().(*rpcInfo)
   121  	r.from = from
   122  	r.to = to
   123  	r.invocation = ink
   124  	r.config = config
   125  	r.stats = stats
   126  	return r
   127  }
   128  
   129  func newRPCInfo() interface{} {
   130  	return &rpcInfo{}
   131  }