github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/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  	"github.com/hxx258456/ccgo/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 ([]string)(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  //   {addr0, path: [p0, wt0]}
    74  //   {addr1, path: [p0, wt1]}
    75  //   {addr2, path: [p1, wt2]}
    76  //   {addr3, path: [p1, wt3]}
    77  // ]
    78  //
    79  // Addresses will be split into p0/p1, and the p0/p1 will be removed from the
    80  // path.
    81  //
    82  // Output:
    83  // {
    84  //   p0: [
    85  //     {addr0, path: [wt0]},
    86  //     {addr1, path: [wt1]},
    87  //   ],
    88  //   p1: [
    89  //     {addr2, path: [wt2]},
    90  //     {addr3, path: [wt3]},
    91  //   ],
    92  // }
    93  //
    94  // If hierarchical path is not set, or has no path in it, the address is
    95  // dropped.
    96  func Group(addrs []resolver.Address) map[string][]resolver.Address {
    97  	ret := make(map[string][]resolver.Address)
    98  	for _, addr := range addrs {
    99  		oldPath := Get(addr)
   100  		if len(oldPath) == 0 {
   101  			continue
   102  		}
   103  		curPath := oldPath[0]
   104  		newPath := oldPath[1:]
   105  		newAddr := Set(addr, newPath)
   106  		ret[curPath] = append(ret[curPath], newAddr)
   107  	}
   108  	return ret
   109  }