github.com/weaviate/weaviate@v1.24.6/usecases/sharding/config.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package sharding
    13  
    14  import (
    15  	"encoding/json"
    16  	"fmt"
    17  
    18  	"github.com/pkg/errors"
    19  )
    20  
    21  const (
    22  	DefaultVirtualPerPhysical = 128
    23  	DefaultKey                = "_id"
    24  	DefaultStrategy           = "hash"
    25  	DefaultFunction           = "murmur3"
    26  )
    27  
    28  type Config struct {
    29  	VirtualPerPhysical  int    `json:"virtualPerPhysical"`
    30  	DesiredCount        int    `json:"desiredCount"`
    31  	ActualCount         int    `json:"actualCount"`
    32  	DesiredVirtualCount int    `json:"desiredVirtualCount"`
    33  	ActualVirtualCount  int    `json:"actualVirtualCount"`
    34  	Key                 string `json:"key"`
    35  	Strategy            string `json:"strategy"`
    36  	Function            string `json:"function"`
    37  }
    38  
    39  func (c *Config) setDefaults(nodeCount int) {
    40  	c.VirtualPerPhysical = DefaultVirtualPerPhysical
    41  	c.DesiredCount = nodeCount
    42  	c.DesiredVirtualCount = c.DesiredCount * c.VirtualPerPhysical
    43  	c.Function = DefaultFunction
    44  	c.Key = DefaultKey
    45  	c.Strategy = DefaultStrategy
    46  
    47  	// these will only differ once there is an async component through replication
    48  	// or dynamic scaling. For now they have to be the same
    49  	c.ActualCount = c.DesiredCount
    50  	c.ActualVirtualCount = c.DesiredVirtualCount
    51  }
    52  
    53  func (c *Config) validate() error {
    54  	if c.Key != "_id" {
    55  		return errors.Errorf("sharding only supported on key '_id' for now, "+
    56  			"got: %s", c.Key)
    57  	}
    58  
    59  	if c.Strategy != "hash" {
    60  		return errors.Errorf("sharding only supported with strategy 'hash' for now, "+
    61  			"got: %s", c.Strategy)
    62  	}
    63  
    64  	if c.Function != "murmur3" {
    65  		return errors.Errorf("sharding only supported with function 'murmur3' for now, "+
    66  			"got: %s", c.Function)
    67  	}
    68  
    69  	return nil
    70  }
    71  
    72  func ParseConfig(input interface{}, nodeCount int) (Config, error) {
    73  	out := Config{}
    74  	out.setDefaults(nodeCount)
    75  
    76  	if input == nil {
    77  		return out, nil
    78  	}
    79  
    80  	asMap, ok := input.(map[string]interface{})
    81  	if !ok || asMap == nil {
    82  		return out, fmt.Errorf("input must be a non-nil map")
    83  	}
    84  
    85  	if err := optionalIntFromMap(asMap, "virtualPerPhysical", func(v int) {
    86  		out.VirtualPerPhysical = v
    87  	}); err != nil {
    88  		return out, err
    89  	}
    90  
    91  	if err := optionalIntFromMap(asMap, "desiredCount", func(v int) {
    92  		out.DesiredCount = v
    93  	}); err != nil {
    94  		return out, err
    95  	}
    96  
    97  	out.DesiredVirtualCount = out.DesiredCount * out.VirtualPerPhysical
    98  
    99  	if err := optionalIntFromMap(asMap, "desiredCount", func(v int) {
   100  		out.DesiredCount = v
   101  	}); err != nil {
   102  		return out, err
   103  	}
   104  
   105  	if err := optionalStringFromMap(asMap, "key", func(v string) {
   106  		out.Key = v
   107  	}); err != nil {
   108  		return out, err
   109  	}
   110  
   111  	if err := optionalStringFromMap(asMap, "strategy", func(v string) {
   112  		out.Strategy = v
   113  	}); err != nil {
   114  		return out, err
   115  	}
   116  
   117  	if err := optionalStringFromMap(asMap, "function", func(v string) {
   118  		out.Function = v
   119  	}); err != nil {
   120  		return out, err
   121  	}
   122  
   123  	// these will only differ once there is an async component through replication
   124  	// or dynamic scaling. For now they have to be the same
   125  	out.ActualCount = out.DesiredCount
   126  	out.ActualVirtualCount = out.DesiredVirtualCount
   127  
   128  	if err := out.validate(); err != nil {
   129  		return out, err
   130  	}
   131  
   132  	return out, nil
   133  }
   134  
   135  func optionalIntFromMap(in map[string]interface{}, name string,
   136  	setFn func(v int),
   137  ) error {
   138  	value, ok := in[name]
   139  	if !ok {
   140  		return nil
   141  	}
   142  
   143  	var asInt64 int64
   144  	var err error
   145  
   146  	// depending on whether we get the results from disk or from the REST API,
   147  	// numbers may be represented slightly differently
   148  	switch typed := value.(type) {
   149  	case json.Number:
   150  		asInt64, err = typed.Int64()
   151  	case int:
   152  		asInt64 = int64(typed)
   153  	case float64:
   154  		asInt64 = int64(typed)
   155  	}
   156  	if err != nil {
   157  		return errors.Wrapf(err, "%s", name)
   158  	}
   159  
   160  	setFn(int(asInt64))
   161  	return nil
   162  }
   163  
   164  func optionalStringFromMap(in map[string]interface{}, name string,
   165  	setFn func(v string),
   166  ) error {
   167  	value, ok := in[name]
   168  	if !ok {
   169  		return nil
   170  	}
   171  
   172  	asString, ok := value.(string)
   173  	if !ok {
   174  		return errors.Errorf("field %q must be of type string, got: %T", name, value)
   175  	}
   176  
   177  	setFn(asString)
   178  	return nil
   179  }