milvus/pkg/util/funcutil/func.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package funcutil

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/cockroachdb/errors"
	"google.golang.org/grpc/codes"
	grpcStatus "google.golang.org/grpc/status"

	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
	"github.com/milvus-io/milvus/pkg/util"
	"github.com/milvus-io/milvus/pkg/util/typeutil"
)

// CheckGrpcReady wait for context timeout, or wait 100ms then send nil to targetCh
func CheckGrpcReady(ctx context.Context, targetCh chan error) {
	timer := time.NewTimer(100 * time.Millisecond)
	defer timer.Stop()
	select {
	case <-timer.C:
		targetCh <- nil
	case <-ctx.Done():
		return
	}
}

// GetIP return the ip address
func GetIP(ip string) string {
	if len(ip) == 0 {
		return GetLocalIP()
	}
	return ip
}

// GetLocalIP return the local ip address
func GetLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err == nil {
		for _, addr := range addrs {
			ipaddr, ok := addr.(*net.IPNet)
			if ok && ipaddr.IP.IsGlobalUnicast() && ipaddr.IP.To4() != nil {
				return ipaddr.IP.String()
			}
		}
	}
	return "127.0.0.1"
}

// JSONToMap parse the jsonic index parameters to map
func JSONToMap(mStr string) (map[string]string, error) {
	buffer := make(map[string]any)
	err := json.Unmarshal([]byte(mStr), &buffer)
	if err != nil {
		return nil, fmt.Errorf("unmarshal params failed, %w", err)
	}
	ret := make(map[string]string)
	for key, value := range buffer {
		valueStr := fmt.Sprintf("%v", value)
		ret[key] = valueStr
	}
	return ret, nil
}

func MapToJSON(m map[string]string) []byte {
	// error won't happen here.
	bs, _ := json.Marshal(m)
	return bs
}

func JSONToRoleDetails(mStr string) (map[string](map[string]([](map[string]string))), error) {
	buffer := make(map[string](map[string]([](map[string]string))), 0)
	err := json.Unmarshal([]byte(mStr), &buffer)
	if err != nil {
		return nil, fmt.Errorf("unmarshal `builtinRoles.Roles` failed, %w", err)
	}
	ret := make(map[string](map[string]([](map[string]string))), 0)
	for role, privilegesJSON := range buffer {
		ret[role] = make(map[string]([](map[string]string)), 0)
		privilegesArray := make([]map[string]string, 0)
		for _, privileges := range privilegesJSON[util.RoleConfigPrivileges] {
			privilegesArray = append(privilegesArray, map[string]string{
				util.RoleConfigObjectType: privileges[util.RoleConfigObjectType],
				util.RoleConfigObjectName: privileges[util.RoleConfigObjectName],
				util.RoleConfigPrivilege:  privileges[util.RoleConfigPrivilege],
				util.RoleConfigDBName:     privileges[util.RoleConfigDBName],
			})
		}
		ret[role]["privileges"] = privilegesArray
	}
	return ret, nil
}

func RoleDetailsToJSON(m map[string](map[string]([](map[string]string)))) []byte {
	bs, _ := json.Marshal(m)
	return bs
}

const (
	// PulsarMaxMessageSizeKey is the key of config item
	PulsarMaxMessageSizeKey = "maxMessageSize"
)

// GetAttrByKeyFromRepeatedKV return the value corresponding to key in kv pair
func GetAttrByKeyFromRepeatedKV(key string, kvs []*commonpb.KeyValuePair) (string, error) {
	for _, kv := range kvs {
		if kv.Key == key {
			return kv.Value, nil
		}
	}

	return "", fmt.Errorf("key %s not found", key)
}

// CheckCtxValid check if the context is valid
func CheckCtxValid(ctx context.Context) bool {
	return ctx.Err() != context.DeadlineExceeded && ctx.Err() != context.Canceled
}

func GetVecFieldIDs(schema *schemapb.CollectionSchema) []int64 {
	var vecFieldIDs []int64
	for _, field := range schema.Fields {
		if field.DataType == schemapb.DataType_BinaryVector || field.DataType == schemapb.DataType_FloatVector || field.DataType == schemapb.DataType_Float16Vector || field.DataType == schemapb.DataType_SparseFloatVector {
			vecFieldIDs = append(vecFieldIDs, field.FieldID)
		}
	}

	return vecFieldIDs
}

func Map2KeyValuePair(datas map[string]string) []*commonpb.KeyValuePair {
	results := make([]*commonpb.KeyValuePair, len(datas))
	offset := 0
	for key, value := range datas {
		results[offset] = &commonpb.KeyValuePair{
			Key:   key,
			Value: value,
		}
		offset++
	}
	return results
}

func KeyValuePair2Map(datas []*commonpb.KeyValuePair) map[string]string {
	results := make(map[string]string)
	for _, pair := range datas {
		results[pair.Key] = pair.Value
	}

	return results
}

func ConvertToKeyValuePairPointer(datas []commonpb.KeyValuePair) []*commonpb.KeyValuePair {
	var kvs []*commonpb.KeyValuePair
	for i := 0; i < len(datas); i++ {
		kvs = append(kvs, &datas[i])
	}
	return kvs
}

// GenChannelSubName generate subName to watch channel
func GenChannelSubName(prefix string, collectionID int64, nodeID int64) string {
	return fmt.Sprintf("%s-%d-%d", prefix, collectionID, nodeID)
}

// CheckPortAvailable check if a port is available to be listened on
func CheckPortAvailable(port int) bool {
	addr := ":" + strconv.Itoa(port)
	listener, err := net.Listen("tcp", addr)
	if listener != nil {
		listener.Close()
	}
	return err == nil
}

// GetAvailablePort return an available port that can be listened on
func GetAvailablePort() int {
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		panic(err)
	}
	defer listener.Close()

	return listener.Addr().(*net.TCPAddr).Port
}

// ToPhysicalChannel get physical channel name from virtual channel name
func ToPhysicalChannel(vchannel string) string {
	index := strings.LastIndex(vchannel, "_")
	if index < 0 {
		return vchannel
	}
	return vchannel[:index]
}

// ConvertChannelName assembles channel name according to parameters.
func ConvertChannelName(chanName string, tokenFrom string, tokenTo string) (string, error) {
	if tokenFrom == "" {
		return "", fmt.Errorf("the tokenFrom is empty")
	}
	if !strings.Contains(chanName, tokenFrom) {
		return "", fmt.Errorf("cannot find token '%s' in '%s'", tokenFrom, chanName)
	}
	return strings.Replace(chanName, tokenFrom, tokenTo, 1), nil
}

func getNumRowsOfScalarField(datas interface{}) uint64 {
	realTypeDatas := reflect.ValueOf(datas)
	return uint64(realTypeDatas.Len())
}

func GetNumRowsOfFloatVectorField(fDatas []float32, dim int64) (uint64, error) {
	if dim <= 0 {
		return 0, fmt.Errorf("dim(%d) should be greater than 0", dim)
	}
	l := len(fDatas)
	if int64(l)%dim != 0 {
		return 0, fmt.Errorf("the length(%d) of float data should divide the dim(%d)", l, dim)
	}
	return uint64(int64(l) / dim), nil
}

func GetNumRowsOfBinaryVectorField(bDatas []byte, dim int64) (uint64, error) {
	if dim <= 0 {
		return 0, fmt.Errorf("dim(%d) should be greater than 0", dim)
	}
	if dim%8 != 0 {
		return 0, fmt.Errorf("dim(%d) should divide 8", dim)
	}
	l := len(bDatas)
	if (8*int64(l))%dim != 0 {
		return 0, fmt.Errorf("the num(%d) of all bits should divide the dim(%d)", 8*l, dim)
	}
	return uint64((8 * int64(l)) / dim), nil
}

func GetNumRowsOfFloat16VectorField(f16Datas []byte, dim int64) (uint64, error) {
	if dim <= 0 {
		return 0, fmt.Errorf("dim(%d) should be greater than 0", dim)
	}
	l := len(f16Datas)
	if int64(l)%dim != 0 {
		return 0, fmt.Errorf("the length(%d) of float data should divide the dim(%d)", l, dim)
	}
	return uint64((int64(l)) / dim / 2), nil
}

func GetNumRowsOfBFloat16VectorField(bf16Datas []byte, dim int64) (uint64, error) {
	if dim <= 0 {
		return 0, fmt.Errorf("dim(%d) should be greater than 0", dim)
	}
	l := len(bf16Datas)
	if int64(l)%dim != 0 {
		return 0, fmt.Errorf("the length(%d) of float data should divide the dim(%d)", l, dim)
	}
	return uint64((int64(l)) / dim / 2), nil
}

func GetNumRowOfFieldData(fieldData *schemapb.FieldData) (uint64, error) {
	var fieldNumRows uint64
	var err error
	switch fieldType := fieldData.Field.(type) {
	case *schemapb.FieldData_Scalars:
		scalarField := fieldData.GetScalars()
		switch scalarType := scalarField.Data.(type) {
		case *schemapb.ScalarField_BoolData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetBoolData().Data)
		case *schemapb.ScalarField_IntData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetIntData().Data)
		case *schemapb.ScalarField_LongData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetLongData().Data)
		case *schemapb.ScalarField_FloatData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetFloatData().Data)
		case *schemapb.ScalarField_DoubleData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetDoubleData().Data)
		case *schemapb.ScalarField_StringData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetStringData().Data)
		case *schemapb.ScalarField_ArrayData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetArrayData().Data)
		case *schemapb.ScalarField_JsonData:
			fieldNumRows = getNumRowsOfScalarField(scalarField.GetJsonData().Data)
		default:
			return 0, fmt.Errorf("%s is not supported now", scalarType)
		}
	case *schemapb.FieldData_Vectors:
		vectorField := fieldData.GetVectors()
		switch vectorFieldType := vectorField.Data.(type) {
		case *schemapb.VectorField_FloatVector:
			dim := vectorField.GetDim()
			fieldNumRows, err = GetNumRowsOfFloatVectorField(vectorField.GetFloatVector().Data, dim)
			if err != nil {
				return 0, err
			}
		case *schemapb.VectorField_BinaryVector:
			dim := vectorField.GetDim()
			fieldNumRows, err = GetNumRowsOfBinaryVectorField(vectorField.GetBinaryVector(), dim)
			if err != nil {
				return 0, err
			}
		case *schemapb.VectorField_Float16Vector:
			dim := vectorField.GetDim()
			fieldNumRows, err = GetNumRowsOfFloat16VectorField(vectorField.GetFloat16Vector(), dim)
			if err != nil {
				return 0, err
			}
		case *schemapb.VectorField_Bfloat16Vector:
			dim := vectorField.GetDim()
			fieldNumRows, err = GetNumRowsOfBFloat16VectorField(vectorField.GetBfloat16Vector(), dim)
			if err != nil {
				return 0, err
			}
		case *schemapb.VectorField_SparseFloatVector:
			fieldNumRows = uint64(len(vectorField.GetSparseFloatVector().GetContents()))
		default:
			return 0, fmt.Errorf("%s is not supported now", vectorFieldType)
		}
	default:
		return 0, fmt.Errorf("%s is not supported now", fieldType)
	}

	return fieldNumRows, nil
}

// ReadBinary read byte slice as receiver.
func ReadBinary(endian binary.ByteOrder, bs []byte, receiver interface{}) error {
	buf := bytes.NewReader(bs)
	return binary.Read(buf, endian, receiver)
}

// IsGrpcErr checks whether err is instance of grpc status error.
func IsGrpcErr(err error, targets ...codes.Code) bool {
	set := typeutil.NewSet[codes.Code](targets...)
	for {
		if err == nil {
			return false
		}
		s, ok := grpcStatus.FromError(err)
		if ok {
			return set.Len() == 0 || set.Contain(s.Code())
		}
		err = errors.Unwrap(err)
	}
}

func IsEmptyString(str string) bool {
	return strings.TrimSpace(str) == ""
}

func HandleTenantForEtcdKey(prefix string, tenant string, key string) string {
	res := prefix
	if tenant != "" {
		res += "/" + tenant
	}
	if key != "" {
		res += "/" + key
	}
	return res
}

func IsRevoke(operateType milvuspb.OperatePrivilegeType) bool {
	return operateType == milvuspb.OperatePrivilegeType_Revoke
}

func IsGrant(operateType milvuspb.OperatePrivilegeType) bool {
	return operateType == milvuspb.OperatePrivilegeType_Grant
}

func EncodeUserRoleCache(user string, role string) string {
	return fmt.Sprintf("%s/%s", user, role)
}

func DecodeUserRoleCache(cache string) (string, string, error) {
	index := strings.LastIndex(cache, "/")
	if index == -1 {
		return "", "", fmt.Errorf("invalid param, cache: [%s]", cache)
	}
	user := cache[:index]
	role := cache[index+1:]
	return user, role, nil
}