github.com/cloudwego/hertz@v0.9.3/pkg/app/middlewares/client/sd/options.go (about)

     1  /*
     2   * Copyright 2022 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 sd
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"net"
    23  	"strings"
    24  
    25  	"github.com/cloudwego/hertz/pkg/app/client/discovery"
    26  	"github.com/cloudwego/hertz/pkg/app/client/loadbalance"
    27  	"github.com/cloudwego/hertz/pkg/app/server/registry"
    28  )
    29  
    30  // ServiceDiscoveryOptions service discovery option for client
    31  type ServiceDiscoveryOptions struct {
    32  	// Resolver is used to client discovery
    33  	Resolver discovery.Resolver
    34  
    35  	// Balancer is used to client load balance
    36  	Balancer loadbalance.Loadbalancer
    37  
    38  	// LbOpts LoadBalance option
    39  	LbOpts loadbalance.Options
    40  }
    41  
    42  func (o *ServiceDiscoveryOptions) Apply(opts []ServiceDiscoveryOption) {
    43  	for _, op := range opts {
    44  		op.F(o)
    45  	}
    46  }
    47  
    48  type ServiceDiscoveryOption struct {
    49  	F func(o *ServiceDiscoveryOptions)
    50  }
    51  
    52  // WithCustomizedAddrs specifies the target instance addresses when doing service discovery.
    53  // It overwrites the results from the Resolver
    54  func WithCustomizedAddrs(addrs ...string) ServiceDiscoveryOption {
    55  	return ServiceDiscoveryOption{
    56  		F: func(o *ServiceDiscoveryOptions) {
    57  			var ins []discovery.Instance
    58  			for _, addr := range addrs {
    59  				if _, err := net.ResolveTCPAddr("tcp", addr); err == nil {
    60  					ins = append(ins, discovery.NewInstance("tcp", addr, registry.DefaultWeight, nil))
    61  					continue
    62  				}
    63  				if _, err := net.ResolveUnixAddr("unix", addr); err == nil {
    64  					ins = append(ins, discovery.NewInstance("unix", addr, registry.DefaultWeight, nil))
    65  					continue
    66  				}
    67  				panic(fmt.Errorf("WithCustomizedAddrs: invalid '%s'", addr))
    68  			}
    69  			if len(ins) == 0 {
    70  				panic("WithCustomizedAddrs() requires at least one argument")
    71  			}
    72  
    73  			targets := strings.Join(addrs, ",")
    74  			o.Resolver = &discovery.SynthesizedResolver{
    75  				ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
    76  					return discovery.Result{
    77  						CacheKey:  "fixed",
    78  						Instances: ins,
    79  					}, nil
    80  				},
    81  				NameFunc: func() string { return targets },
    82  				TargetFunc: func(ctx context.Context, target *discovery.TargetInfo) string {
    83  					return targets
    84  				},
    85  			}
    86  		},
    87  	}
    88  }
    89  
    90  // WithLoadBalanceOptions  sets Loadbalancer and loadbalance options for hertz client
    91  func WithLoadBalanceOptions(lb loadbalance.Loadbalancer, options loadbalance.Options) ServiceDiscoveryOption {
    92  	return ServiceDiscoveryOption{F: func(o *ServiceDiscoveryOptions) {
    93  		o.LbOpts = options
    94  		o.Balancer = lb
    95  	}}
    96  }