github.com/polarismesh/polaris@v1.17.8/common/model/ratelimit.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package model
    19  
    20  import (
    21  	"encoding/json"
    22  	"time"
    23  
    24  	apimodel "github.com/polarismesh/specification/source/go/api/v1/model"
    25  	apitraffic "github.com/polarismesh/specification/source/go/api/v1/traffic_manage"
    26  )
    27  
    28  // RateLimit 限流规则
    29  type RateLimit struct {
    30  	Proto     *apitraffic.Rule
    31  	ID        string
    32  	ServiceID string
    33  	Name      string
    34  	Method    string
    35  	// Labels for old compatible, will be removed later
    36  	Labels     string
    37  	Priority   uint32
    38  	Rule       string
    39  	Revision   string
    40  	Disable    bool
    41  	Valid      bool
    42  	CreateTime time.Time
    43  	ModifyTime time.Time
    44  	EnableTime time.Time
    45  }
    46  
    47  // Labels2Arguments 适配老的标签到新的参数列表
    48  func (r *RateLimit) Labels2Arguments() (map[string]*apimodel.MatchString, error) {
    49  	if len(r.Proto.Arguments) == 0 && len(r.Labels) > 0 {
    50  		var labels = make(map[string]*apimodel.MatchString)
    51  		if err := json.Unmarshal([]byte(r.Labels), &labels); err != nil {
    52  			return nil, err
    53  		}
    54  		for key, value := range labels {
    55  			r.Proto.Arguments = append(r.Proto.Arguments, &apitraffic.MatchArgument{
    56  				Type:  apitraffic.MatchArgument_CUSTOM,
    57  				Key:   key,
    58  				Value: value,
    59  			})
    60  		}
    61  		return labels, nil
    62  	}
    63  	return nil, nil
    64  }
    65  
    66  const (
    67  	LabelKeyPath          = "$path"
    68  	LabelKeyMethod        = "$method"
    69  	LabelKeyHeader        = "$header"
    70  	LabelKeyQuery         = "$query"
    71  	LabelKeyCallerService = "$caller_service"
    72  	LabelKeyCallerIP      = "$caller_ip"
    73  )
    74  
    75  // Arguments2Labels 将参数列表适配成旧的标签模型
    76  func Arguments2Labels(arguments []*apitraffic.MatchArgument) map[string]*apimodel.MatchString {
    77  	if len(arguments) > 0 {
    78  		var labels = make(map[string]*apimodel.MatchString)
    79  		for _, argument := range arguments {
    80  			key := BuildArgumentKey(argument.Type, argument.Key)
    81  			labels[key] = argument.Value
    82  		}
    83  		return labels
    84  	}
    85  	return nil
    86  }
    87  
    88  func BuildArgumentKey(argumentType apitraffic.MatchArgument_Type, key string) string {
    89  	switch argumentType {
    90  	case apitraffic.MatchArgument_HEADER:
    91  		return LabelKeyHeader + "." + key
    92  	case apitraffic.MatchArgument_QUERY:
    93  		return LabelKeyQuery + "." + key
    94  	case apitraffic.MatchArgument_CALLER_SERVICE:
    95  		return LabelKeyCallerService + "." + key
    96  	case apitraffic.MatchArgument_CALLER_IP:
    97  		return LabelKeyCallerIP
    98  	case apitraffic.MatchArgument_CUSTOM:
    99  		return key
   100  	case apitraffic.MatchArgument_METHOD:
   101  		return LabelKeyMethod
   102  	default:
   103  		return key
   104  	}
   105  }
   106  
   107  // AdaptArgumentsAndLabels 对存量标签进行兼容，同时将argument适配成标签
   108  func (r *RateLimit) AdaptArgumentsAndLabels() error {
   109  	// 新的限流规则，需要适配老的SDK使用场景
   110  	labels := Arguments2Labels(r.Proto.GetArguments())
   111  	if len(labels) > 0 {
   112  		r.Proto.Labels = labels
   113  	} else {
   114  		var err error
   115  		// 存量限流规则，需要适配成新的规则
   116  		labels, err = r.Labels2Arguments()
   117  		if nil != err {
   118  			return err
   119  		}
   120  		r.Proto.Labels = labels
   121  	}
   122  	return nil
   123  }
   124  
   125  // AdaptLabels 对存量标签进行兼容，对存量labels进行清空
   126  func (r *RateLimit) AdaptLabels() error {
   127  	// 存量限流规则，需要适配成新的规则
   128  	_, err := r.Labels2Arguments()
   129  	if nil != err {
   130  		return err
   131  	}
   132  	r.Proto.Labels = nil
   133  	return nil
   134  }
   135  
   136  // ExtendRateLimit 包含服务信息的限流规则
   137  type ExtendRateLimit struct {
   138  	ServiceName   string
   139  	NamespaceName string
   140  	RateLimit     *RateLimit
   141  }
   142  
   143  // RateLimitRevision 包含最新版本号的限流规则
   144  type RateLimitRevision struct {
   145  	ServiceID    string
   146  	LastRevision string
   147  	ModifyTime   time.Time
   148  }