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

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