gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/internal/hierarchy/hierarchy.go (about)

     1  /*
     2   *
     3   * Copyright 2020 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package hierarchy contains functions to set and get hierarchy string from
    20  // addresses.
    21  //
    22  // This package is experimental.
    23  package hierarchy
    24  
    25  import (
    26  	"gitee.com/ks-custle/core-gm/grpc/resolver"
    27  )
    28  
    29  type pathKeyType string
    30  
    31  const pathKey = pathKeyType("grpc.internal.address.hierarchical_path")
    32  
    33  type pathValue []string
    34  
    35  func (p pathValue) Equal(o interface{}) bool {
    36  	op, ok := o.(pathValue)
    37  	if !ok {
    38  		return false
    39  	}
    40  	if len(op) != len(p) {
    41  		return false
    42  	}
    43  	for i, v := range p {
    44  		if v != op[i] {
    45  			return false
    46  		}
    47  	}
    48  	return true
    49  }
    50  
    51  // Get returns the hierarchical path of addr.
    52  func Get(addr resolver.Address) []string {
    53  	attrs := addr.BalancerAttributes
    54  	if attrs == nil {
    55  		return nil
    56  	}
    57  	path, _ := attrs.Value(pathKey).(pathValue)
    58  	return path
    59  }
    60  
    61  // Set overrides the hierarchical path in addr with path.
    62  func Set(addr resolver.Address, path []string) resolver.Address {
    63  	addr.BalancerAttributes = addr.BalancerAttributes.WithValue(pathKey, pathValue(path))
    64  	return addr
    65  }
    66  
    67  // Group splits a slice of addresses into groups based on
    68  // the first hierarchy path. The first hierarchy path will be removed from the
    69  // result.
    70  //
    71  // Input:
    72  // [
    73  //
    74  //	{addr0, path: [p0, wt0]}
    75  //	{addr1, path: [p0, wt1]}
    76  //	{addr2, path: [p1, wt2]}
    77  //	{addr3, path: [p1, wt3]}
    78  //
    79  // ]
    80  //
    81  // Addresses will be split into p0/p1, and the p0/p1 will be removed from the
    82  // path.
    83  //
    84  // Output:
    85  //
    86  //	{
    87  //	  p0: [
    88  //	    {addr0, path: [wt0]},
    89  //	    {addr1, path: [wt1]},
    90  //	  ],
    91  //	  p1: [
    92  //	    {addr2, path: [wt2]},
    93  //	    {addr3, path: [wt3]},
    94  //	  ],
    95  //	}
    96  //
    97  // If hierarchical path is not set, or has no path in it, the address is
    98  // dropped.
    99  func Group(addrs []resolver.Address) map[string][]resolver.Address {
   100  	ret := make(map[string][]resolver.Address)
   101  	for _, addr := range addrs {
   102  		oldPath := Get(addr)
   103  		if len(oldPath) == 0 {
   104  			continue
   105  		}
   106  		curPath := oldPath[0]
   107  		newPath := oldPath[1:]
   108  		newAddr := Set(addr, newPath)
   109  		ret[curPath] = append(ret[curPath], newAddr)
   110  	}
   111  	return ret
   112  }