github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/balancers/config.go (about)

     1  package balancers
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  
     7  	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
     8  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
     9  )
    10  
    11  type balancerType string
    12  
    13  const (
    14  	typeRoundRobin   = balancerType("round_robin")
    15  	typeRandomChoice = balancerType("random_choice")
    16  	typeSingle       = balancerType("single")
    17  	typeDisable      = balancerType("disable")
    18  )
    19  
    20  type preferType string
    21  
    22  const (
    23  	preferTypeLocalDC   = preferType("local_dc")
    24  	preferTypeLocations = preferType("locations")
    25  )
    26  
    27  type balancersConfig struct {
    28  	Type      balancerType `json:"type"`
    29  	Prefer    preferType   `json:"prefer,omitempty"`
    30  	Fallback  bool         `json:"fallback,omitempty"`
    31  	Locations []string     `json:"locations,omitempty"`
    32  }
    33  
    34  type fromConfigOptionsHolder struct {
    35  	fallbackBalancer *balancerConfig.Config
    36  	errorHandler     func(error)
    37  }
    38  
    39  type fromConfigOption func(h *fromConfigOptionsHolder)
    40  
    41  func WithParseErrorFallbackBalancer(b *balancerConfig.Config) fromConfigOption {
    42  	return func(h *fromConfigOptionsHolder) {
    43  		h.fallbackBalancer = b
    44  	}
    45  }
    46  
    47  func WithParseErrorHandler(errorHandler func(error)) fromConfigOption {
    48  	return func(h *fromConfigOptionsHolder) {
    49  		h.errorHandler = errorHandler
    50  	}
    51  }
    52  
    53  func createByType(t balancerType) (*balancerConfig.Config, error) {
    54  	switch t {
    55  	case typeDisable:
    56  		return SingleConn(), nil
    57  	case typeSingle:
    58  		return SingleConn(), nil
    59  	case typeRandomChoice:
    60  		return RandomChoice(), nil
    61  	case typeRoundRobin:
    62  		return RoundRobin(), nil
    63  	default:
    64  		return nil, xerrors.WithStackTrace(fmt.Errorf("unknown type of balancer: %s", t))
    65  	}
    66  }
    67  
    68  func CreateFromConfig(s string) (*balancerConfig.Config, error) {
    69  	// try to parse s as identifier of balancer
    70  	if c, err := createByType(balancerType(s)); err == nil {
    71  		return c, nil
    72  	}
    73  
    74  	var (
    75  		b   *balancerConfig.Config
    76  		err error
    77  		c   balancersConfig
    78  	)
    79  
    80  	// try to parse s as json
    81  	if err = json.Unmarshal([]byte(s), &c); err != nil {
    82  		return nil, xerrors.WithStackTrace(err)
    83  	}
    84  
    85  	b, err = createByType(c.Type)
    86  	if err != nil {
    87  		return nil, xerrors.WithStackTrace(err)
    88  	}
    89  
    90  	switch c.Prefer {
    91  	case preferTypeLocalDC:
    92  		if c.Fallback {
    93  			return PreferLocalDCWithFallBack(b), nil
    94  		}
    95  
    96  		return PreferLocalDC(b), nil
    97  	case preferTypeLocations:
    98  		if len(c.Locations) == 0 {
    99  			return nil, xerrors.WithStackTrace(fmt.Errorf("empty locations list in balancer '%s' config", c.Type))
   100  		}
   101  		if c.Fallback {
   102  			return PreferLocationsWithFallback(b, c.Locations...), nil
   103  		}
   104  
   105  		return PreferLocations(b, c.Locations...), nil
   106  	default:
   107  		return b, nil
   108  	}
   109  }
   110  
   111  func FromConfig(config string, opts ...fromConfigOption) *balancerConfig.Config {
   112  	var (
   113  		h = fromConfigOptionsHolder{
   114  			fallbackBalancer: Default(),
   115  		}
   116  		b   *balancerConfig.Config
   117  		err error
   118  	)
   119  	for _, o := range opts {
   120  		if o != nil {
   121  			o(&h)
   122  		}
   123  	}
   124  
   125  	b, err = CreateFromConfig(config)
   126  	if err != nil {
   127  		if h.errorHandler != nil {
   128  			h.errorHandler(err)
   129  		}
   130  
   131  		return h.fallbackBalancer
   132  	}
   133  
   134  	return b
   135  }