google.golang.org/grpc@v1.62.1/internal/balancer/stub/stub.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 stub implements a balancer for testing purposes.
    20  package stub
    21  
    22  import (
    23  	"encoding/json"
    24  	"fmt"
    25  
    26  	"google.golang.org/grpc/balancer"
    27  	"google.golang.org/grpc/serviceconfig"
    28  )
    29  
    30  // BalancerFuncs contains all balancer.Balancer functions with a preceding
    31  // *BalancerData parameter for passing additional instance information.  Any
    32  // nil functions will never be called.
    33  type BalancerFuncs struct {
    34  	// Init is called after ClientConn and BuildOptions are set in
    35  	// BalancerData.  It may be used to initialize BalancerData.Data.
    36  	Init func(*BalancerData)
    37  	// ParseConfig is used for parsing LB configs, if specified.
    38  	ParseConfig func(json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
    39  
    40  	UpdateClientConnState func(*BalancerData, balancer.ClientConnState) error
    41  	ResolverError         func(*BalancerData, error)
    42  	Close                 func(*BalancerData)
    43  	ExitIdle              func(*BalancerData)
    44  }
    45  
    46  // BalancerData contains data relevant to a stub balancer.
    47  type BalancerData struct {
    48  	// ClientConn is set by the builder.
    49  	ClientConn balancer.ClientConn
    50  	// BuildOptions is set by the builder.
    51  	BuildOptions balancer.BuildOptions
    52  	// Data may be used to store arbitrary user data.
    53  	Data any
    54  }
    55  
    56  type bal struct {
    57  	bf BalancerFuncs
    58  	bd *BalancerData
    59  }
    60  
    61  func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error {
    62  	if b.bf.UpdateClientConnState != nil {
    63  		return b.bf.UpdateClientConnState(b.bd, c)
    64  	}
    65  	return nil
    66  }
    67  
    68  func (b *bal) ResolverError(e error) {
    69  	if b.bf.ResolverError != nil {
    70  		b.bf.ResolverError(b.bd, e)
    71  	}
    72  }
    73  
    74  func (b *bal) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
    75  	panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, scs))
    76  }
    77  
    78  func (b *bal) Close() {
    79  	if b.bf.Close != nil {
    80  		b.bf.Close(b.bd)
    81  	}
    82  }
    83  
    84  func (b *bal) ExitIdle() {
    85  	if b.bf.ExitIdle != nil {
    86  		b.bf.ExitIdle(b.bd)
    87  	}
    88  }
    89  
    90  type bb struct {
    91  	name string
    92  	bf   BalancerFuncs
    93  }
    94  
    95  func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
    96  	b := &bal{bf: bb.bf, bd: &BalancerData{ClientConn: cc, BuildOptions: opts}}
    97  	if b.bf.Init != nil {
    98  		b.bf.Init(b.bd)
    99  	}
   100  	return b
   101  }
   102  
   103  func (bb bb) Name() string { return bb.name }
   104  
   105  func (bb bb) ParseConfig(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
   106  	if bb.bf.ParseConfig != nil {
   107  		return bb.bf.ParseConfig(lbCfg)
   108  	}
   109  	return nil, nil
   110  }
   111  
   112  // Register registers a stub balancer builder which will call the provided
   113  // functions.  The name used should be unique.
   114  func Register(name string, bf BalancerFuncs) {
   115  	balancer.Register(bb{name: name, bf: bf})
   116  }