prefix_endpoints.go raw

   1  package endpoints
   2  
   3  import (
   4  	"context"
   5  	"sync"
   6  
   7  	"google.golang.org/grpc"
   8  	"google.golang.org/protobuf/reflect/protoreflect"
   9  
  10  	"github.com/yandex-cloud/go-sdk/v2/pkg/errors"
  11  )
  12  
  13  // PrefixToEndpoint maps a protobuf FullName prefix to its EndpointParams.
  14  type PrefixToEndpoint map[protoreflect.FullName]*EndpointParams
  15  
  16  // Get returns the EndpointParams for the longest matching prefix of p.
  17  // It walks up the namespace hierarchy until it finds a match or returns nil.
  18  func (p2e PrefixToEndpoint) Get(p protoreflect.FullName) *EndpointParams {
  19  	for name := p; name != ""; name = name.Parent() {
  20  		if ep, ok := p2e[name]; ok {
  21  			return ep
  22  		}
  23  	}
  24  	return nil
  25  }
  26  
  27  // Clone returns a shallow copy of the PrefixToEndpoint map.
  28  func (p2e PrefixToEndpoint) Clone() PrefixToEndpoint {
  29  	out := make(PrefixToEndpoint, len(p2e))
  30  	for k, v := range p2e {
  31  		out[k] = v
  32  	}
  33  	return out
  34  }
  35  
  36  // Merge returns a new PrefixToEndpoint containing:
  37  // 1) All entries from the original map that are not overridden by src (including any parent overrides),
  38  // 2) All entries from src (overriding or adding as needed).
  39  func (p2e PrefixToEndpoint) Merge(src PrefixToEndpoint) PrefixToEndpoint {
  40  	result := make(PrefixToEndpoint, len(p2e)+len(src))
  41  	// Copy entries from the original that are not covered by src.
  42  	for k, v := range p2e {
  43  		if src.Get(k) == nil {
  44  			result[k] = v
  45  		}
  46  	}
  47  	// Add or override with all entries from src.
  48  	for k, v := range src {
  49  		result[k] = v
  50  	}
  51  	return result
  52  }
  53  
  54  type PrefixEndpointsResolver struct {
  55  	base             PrefixToEndpoint
  56  	prefixToEndpoint map[protoreflect.FullName]*Endpoint
  57  	cache            sync.Map // map[protoreflect.FullName]*Endpoint
  58  }
  59  
  60  func NewPrefixEndpointsResolver(p2e PrefixToEndpoint) *PrefixEndpointsResolver {
  61  	m := make(map[protoreflect.FullName]*Endpoint, len(p2e))
  62  
  63  	for p, params := range p2e {
  64  		m[p] = params.Build()
  65  	}
  66  	return &PrefixEndpointsResolver{prefixToEndpoint: m, base: p2e}
  67  }
  68  
  69  func (r *PrefixEndpointsResolver) Endpoint(_ context.Context, method protoreflect.FullName, _ ...grpc.CallOption) (*Endpoint, error) {
  70  	if v, ok := r.cache.Load(method); ok {
  71  		return v.(*Endpoint), nil
  72  	}
  73  
  74  	for name := method; name != ""; name = name.Parent() {
  75  		if e, ok := r.prefixToEndpoint[name]; ok {
  76  			r.cache.Store(method, e)
  77  			return e, nil
  78  		}
  79  	}
  80  
  81  	return nil, &errors.EndpointNotFoundError{Method: method}
  82  }
  83  
  84  func (r *PrefixEndpointsResolver) PrefixToEndpoint() PrefixToEndpoint {
  85  	return r.base.Clone()
  86  }
  87