google.golang.org/grpc@v1.62.1/balancer/rls/internal/test/e2e/rls_child_policy.go (about)

     1  /*
     2   *
     3   * Copyright 2021 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 e2e
    20  
    21  import (
    22  	"encoding/json"
    23  	"errors"
    24  	"fmt"
    25  
    26  	"google.golang.org/grpc"
    27  	"google.golang.org/grpc/balancer"
    28  	"google.golang.org/grpc/internal/grpcsync"
    29  	"google.golang.org/grpc/resolver"
    30  	"google.golang.org/grpc/serviceconfig"
    31  )
    32  
    33  const (
    34  	// RLSChildPolicyTargetNameField is a top-level field name to add to the child
    35  	// policy's config, whose value is set to the target for the child policy.
    36  	RLSChildPolicyTargetNameField = "Backend"
    37  	// RLSChildPolicyBadTarget is a value which is considered a bad target by the
    38  	// child policy. This is useful to test bad child policy configuration.
    39  	RLSChildPolicyBadTarget = "bad-target"
    40  )
    41  
    42  // ErrParseConfigBadTarget is the error returned from ParseConfig when the
    43  // backend field is set to RLSChildPolicyBadTarget.
    44  var ErrParseConfigBadTarget = errors.New("backend field set to RLSChildPolicyBadTarget")
    45  
    46  // BalancerFuncs is a set of callbacks which get invoked when the corresponding
    47  // method on the child policy is invoked.
    48  type BalancerFuncs struct {
    49  	UpdateClientConnState func(cfg *RLSChildPolicyConfig) error
    50  	Close                 func()
    51  }
    52  
    53  // RegisterRLSChildPolicy registers a balancer builder with the given name, to
    54  // be used as a child policy for the RLS LB policy.
    55  //
    56  // The child policy uses a pickfirst balancer under the hood to send all traffic
    57  // to the single backend specified by the `RLSChildPolicyTargetNameField` field
    58  // in its configuration which looks like: {"Backend": "Backend-address"}.
    59  func RegisterRLSChildPolicy(name string, bf *BalancerFuncs) {
    60  	balancer.Register(bb{name: name, bf: bf})
    61  }
    62  
    63  type bb struct {
    64  	name string
    65  	bf   *BalancerFuncs
    66  }
    67  
    68  func (bb bb) Name() string { return bb.name }
    69  
    70  func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
    71  	pf := balancer.Get(grpc.PickFirstBalancerName)
    72  	b := &bal{
    73  		Balancer: pf.Build(cc, opts),
    74  		bf:       bb.bf,
    75  		done:     grpcsync.NewEvent(),
    76  	}
    77  	go b.run()
    78  	return b
    79  }
    80  
    81  func (bb bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
    82  	cfg := &RLSChildPolicyConfig{}
    83  	if err := json.Unmarshal(c, cfg); err != nil {
    84  		return nil, err
    85  	}
    86  	if cfg.Backend == RLSChildPolicyBadTarget {
    87  		return nil, ErrParseConfigBadTarget
    88  	}
    89  	return cfg, nil
    90  }
    91  
    92  type bal struct {
    93  	balancer.Balancer
    94  	bf   *BalancerFuncs
    95  	done *grpcsync.Event
    96  }
    97  
    98  // RLSChildPolicyConfig is the LB config for the test child policy.
    99  type RLSChildPolicyConfig struct {
   100  	serviceconfig.LoadBalancingConfig
   101  	Backend string // The target for which this child policy was created.
   102  	Random  string // A random field to test child policy config changes.
   103  }
   104  
   105  func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error {
   106  	cfg, ok := c.BalancerConfig.(*RLSChildPolicyConfig)
   107  	if !ok {
   108  		return fmt.Errorf("received balancer config of type %T, want %T", c.BalancerConfig, &RLSChildPolicyConfig{})
   109  	}
   110  	if b.bf != nil && b.bf.UpdateClientConnState != nil {
   111  		b.bf.UpdateClientConnState(cfg)
   112  	}
   113  	return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
   114  		ResolverState: resolver.State{Addresses: []resolver.Address{{Addr: cfg.Backend}}},
   115  	})
   116  }
   117  
   118  func (b *bal) Close() {
   119  	b.Balancer.Close()
   120  	if b.bf != nil && b.bf.Close != nil {
   121  		b.bf.Close()
   122  	}
   123  	b.done.Fire()
   124  }
   125  
   126  // run is a dummy goroutine to make sure that child policies are closed at the
   127  // end of tests. If they are not closed, these goroutines will be picked up by
   128  // the leakcheker and tests will fail.
   129  func (b *bal) run() {
   130  	<-b.done.Done()
   131  }