github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/model/v1beta1/utils.go (about)

     1  package v1beta1
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"reflect"
     7  	"strings"
     8  
     9  	"github.com/c2h5oh/datasize"
    10  	"k8s.io/apimachinery/pkg/labels"
    11  	"sigs.k8s.io/yaml"
    12  )
    13  
    14  type KeyString string
    15  type KeyInt int
    16  
    17  const MaxSerializedStringInput = int(10 * datasize.MB)
    18  
    19  // Arbitrarily choosing 1000 jobs to serialize - this is a pretty high
    20  const MaxNumberOfObjectsToSerialize = 1000
    21  
    22  const ShortIDLength = 8
    23  
    24  func equal(a, b string) bool {
    25  	a = strings.TrimSpace(a)
    26  	b = strings.TrimSpace(b)
    27  	return strings.EqualFold(a, b)
    28  }
    29  
    30  const (
    31  	jsonMarshal = iota
    32  	jsonMarshalIndent
    33  	yamlMarshal
    34  	jsonUnmarshal
    35  	yamlUnmarshal
    36  )
    37  
    38  func JSONMarshalWithMax[T any](t T) ([]byte, error) {
    39  	return genericMarshalWithMax(t, jsonMarshal, 0)
    40  }
    41  
    42  func JSONMarshalIndentWithMax[T any](t T, indentSpaces int) ([]byte, error) {
    43  	return genericMarshalWithMax(t, jsonMarshalIndent, indentSpaces)
    44  }
    45  
    46  func YAMLMarshalWithMax[T any](t T) ([]byte, error) {
    47  	return genericMarshalWithMax(t, yamlMarshal, -1)
    48  }
    49  
    50  // Create function to take generic and marshall func and return []byte and error
    51  func genericMarshalWithMax[T any](t T, marshalType int, indentSpaces int) ([]byte, error) {
    52  	err := ConfirmMaxSliceSize(t, MaxNumberOfObjectsToSerialize)
    53  	if err != nil {
    54  		return nil, fmt.Errorf("cannot serialize more than %d %s",
    55  			MaxNumberOfObjectsToSerialize,
    56  			reflect.TypeOf(t).String())
    57  	}
    58  	if marshalType == jsonMarshal {
    59  		return json.Marshal(t)
    60  	} else if marshalType == jsonMarshalIndent {
    61  		return json.MarshalIndent(t, "", strings.Repeat(" ", indentSpaces))
    62  	} else if marshalType == yamlMarshal {
    63  		return yaml.Marshal(t)
    64  	}
    65  
    66  	return nil, fmt.Errorf("unknown marshal type %d", marshalType)
    67  }
    68  
    69  func JSONUnmarshalWithMax[T any](b []byte, t *T) error {
    70  	return genericUnmarshalWithMax(b, t, jsonUnmarshal)
    71  }
    72  
    73  func YAMLUnmarshalWithMax[T any](b []byte, t *T) error {
    74  	return genericUnmarshalWithMax(b, t, yamlUnmarshal)
    75  }
    76  
    77  func genericUnmarshalWithMax[T any](b []byte, t *T, unmarshalType int) error {
    78  	if len(b) > MaxSerializedStringInput {
    79  		return fmt.Errorf("size of bytes to unmarshal (%d) larger than maximum allowed (%d)",
    80  			len(b),
    81  			MaxSerializedStringInput)
    82  	}
    83  	if unmarshalType == jsonUnmarshal {
    84  		return json.Unmarshal(b, t)
    85  	} else if unmarshalType == yamlUnmarshal {
    86  		// Our format requires that we use the 	"sigs.k8s.io/yaml" library
    87  		return yaml.Unmarshal(b, t)
    88  	}
    89  	return fmt.Errorf("unknown unmarshal type")
    90  }
    91  
    92  func ConfirmMaxSliceSize[T any](t T, maxSize int) error {
    93  	if _, isSlice := any(t).([]T); isSlice {
    94  		tt := any(t).([]T)
    95  		if len(tt) > maxSize {
    96  			return fmt.Errorf("number of objects (%d) more than max (%d)", len(tt), maxSize)
    97  		}
    98  	}
    99  	return nil
   100  }
   101  
   102  func GetShardID(jobID string, shardIndex int) string {
   103  	return fmt.Sprintf("%s:%d", jobID, shardIndex)
   104  }
   105  
   106  func ToLabelSelectorRequirements(requirements ...labels.Requirement) []LabelSelectorRequirement {
   107  	var labelSelectorRequirements []LabelSelectorRequirement
   108  	for _, requirement := range requirements {
   109  		labelSelectorRequirements = append(labelSelectorRequirements, LabelSelectorRequirement{
   110  			Key:      requirement.Key(),
   111  			Operator: requirement.Operator(),
   112  			Values:   requirement.Values().List(),
   113  		})
   114  	}
   115  	return labelSelectorRequirements
   116  }
   117  
   118  func FromLabelSelectorRequirements(requirements ...LabelSelectorRequirement) ([]labels.Requirement, error) {
   119  	var labelSelectorRequirements []labels.Requirement
   120  	for _, requirement := range requirements {
   121  		req, err := labels.NewRequirement(requirement.Key, requirement.Operator, requirement.Values)
   122  		if err != nil {
   123  			return nil, err
   124  		}
   125  		labelSelectorRequirements = append(labelSelectorRequirements, *req)
   126  	}
   127  	return labelSelectorRequirements, nil
   128  }