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

     1  /*
     2   * Copyright 2023 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 rpctimeout
    18  
    19  import (
    20  	"sync/atomic"
    21  	"time"
    22  
    23  	"github.com/cloudwego/configmanager/iface"
    24  	"github.com/cloudwego/configmanager/util"
    25  
    26  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    27  )
    28  
    29  var (
    30  	_ iface.ConfigValueItem   = (*RPCTimeout)(nil)
    31  	_ rpcinfo.Timeouts        = (*RPCTimeout)(nil)
    32  	_ rpcinfo.TimeoutProvider = (*Container)(nil)
    33  )
    34  
    35  // TypeRPCTimeout is used as itemKey in ConfigValueImpl
    36  const TypeRPCTimeout iface.ItemType = "rpc_timeout"
    37  
    38  const wildcardMethod = "*"
    39  
    40  var defaultRPCTimeout = &RPCTimeout{
    41  	RPCTimeoutMS:  1000,
    42  	ConnTimeoutMS: 50,
    43  }
    44  
    45  // RPCTimeout is used as itemValue in ConfigValueImpl
    46  type RPCTimeout struct {
    47  	RPCTimeoutMS  int `json:"rpc_timeout_ms"`
    48  	ConnTimeoutMS int `json:"conn_timeout_ms"`
    49  }
    50  
    51  // NewRPCTimeout is a function decoding json bytes to a RPCTimeout object
    52  var NewRPCTimeout = util.JsonInitializer(func() iface.ConfigValueItem {
    53  	return &RPCTimeout{}
    54  })
    55  
    56  // CopyDefaultRPCTimeout returns a copy of defaultRPCTimeout, thus avoiding default values changed by business
    57  func CopyDefaultRPCTimeout() iface.ConfigValueItem {
    58  	return defaultRPCTimeout.DeepCopy()
    59  }
    60  
    61  // DeepCopy returns a copy of the current RPCTimeout
    62  func (r *RPCTimeout) DeepCopy() iface.ConfigValueItem {
    63  	result := &RPCTimeout{
    64  		RPCTimeoutMS:  r.RPCTimeoutMS,
    65  		ConnTimeoutMS: r.ConnTimeoutMS,
    66  	}
    67  	return result
    68  }
    69  
    70  // EqualsTo returns true if the current RPCTimeout equals to the other RPCTimeout
    71  func (r *RPCTimeout) EqualsTo(other iface.ConfigValueItem) bool {
    72  	o := other.(*RPCTimeout)
    73  	return r.ConnTimeoutMS == o.ConnTimeoutMS &&
    74  		r.RPCTimeoutMS == o.RPCTimeoutMS
    75  }
    76  
    77  // RPCTimeout implements rpcinfo.Timeouts
    78  func (r *RPCTimeout) RPCTimeout() time.Duration {
    79  	return time.Duration(r.RPCTimeoutMS) * time.Millisecond
    80  }
    81  
    82  // ConnectTimeout implements rpcinfo.Timeouts
    83  func (r *RPCTimeout) ConnectTimeout() time.Duration {
    84  	return time.Duration(r.ConnTimeoutMS) * time.Millisecond
    85  }
    86  
    87  // ReadWriteTimeout implements rpcinfo.Timeouts
    88  func (r *RPCTimeout) ReadWriteTimeout() time.Duration {
    89  	return time.Duration(r.RPCTimeoutMS) * time.Millisecond
    90  }
    91  
    92  // NewContainer build Container for timeout provider.
    93  func NewContainer() *Container {
    94  	c := &Container{}
    95  	rtc := &rpcTimeoutConfig{
    96  		configs:      map[string]*RPCTimeout{},
    97  		globalConfig: CopyDefaultRPCTimeout().(*RPCTimeout),
    98  	}
    99  	c.config.Store(rtc)
   100  	return c
   101  }
   102  
   103  type rpcTimeoutConfig struct {
   104  	configs      map[string]*RPCTimeout
   105  	globalConfig *RPCTimeout
   106  }
   107  
   108  // Container is the implementation of rpcinfo.TimeoutProvider.
   109  // Provide the ability to dynamically configure the rpctimeout
   110  // config on the method hierarchy.
   111  type Container struct {
   112  	config atomic.Value
   113  }
   114  
   115  // NotifyPolicyChange to receive policy when it changes
   116  func (c *Container) NotifyPolicyChange(configs map[string]*RPCTimeout) {
   117  	rtc := &rpcTimeoutConfig{
   118  		configs:      map[string]*RPCTimeout{},
   119  		globalConfig: CopyDefaultRPCTimeout().(*RPCTimeout),
   120  	}
   121  
   122  	for method, cfg := range configs {
   123  		// The configs may be modified by the container invoker,
   124  		// copy them to avoid data race.
   125  		rt := cfg.DeepCopy().(*RPCTimeout)
   126  
   127  		// if has wildcardMethod rpctimeout config, overwrite the default one.
   128  		if method == wildcardMethod {
   129  			rtc.globalConfig = rt
   130  		}
   131  		rtc.configs[method] = rt
   132  	}
   133  
   134  	c.config.Store(rtc)
   135  }
   136  
   137  // Timeouts return the rpc timeout config by the method name of rpc info.
   138  func (c *Container) Timeouts(ri rpcinfo.RPCInfo) rpcinfo.Timeouts {
   139  	rtc := c.config.Load().(*rpcTimeoutConfig)
   140  	if config, ok := rtc.configs[ri.Invocation().MethodName()]; ok {
   141  		return config
   142  	}
   143  	return rtc.globalConfig
   144  }