github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/funcs/fn_sample.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package funcs
     7  
     8  import (
     9  	"fmt"
    10  	"time"
    11  
    12  	"github.com/GuanceCloud/platypus/pkg/ast"
    13  	"github.com/GuanceCloud/platypus/pkg/engine/runtime"
    14  	"github.com/GuanceCloud/platypus/pkg/errchain"
    15  )
    16  
    17  func SampleChecking(ctx *runtime.Task, funcExpr *ast.CallExpr) *errchain.PlError {
    18  	if len(funcExpr.Param) != 1 {
    19  		return runtime.NewRunError(ctx, fmt.Sprintf(
    20  			"func %s expects 1 arg", funcExpr.Name), funcExpr.NamePos)
    21  	}
    22  
    23  	switch funcExpr.Param[0].NodeType { //nolint:exhaustive
    24  	case ast.TypeIntegerLiteral, ast.TypeFloatLiteral, ast.TypeArithmeticExpr:
    25  	default:
    26  		return runtime.NewRunError(ctx, fmt.Sprintf(
    27  			"expect NumberLiteral or ArithmeticExpr, got %s", funcExpr.Param[0].NodeType),
    28  			funcExpr.Param[0].StartPos())
    29  	}
    30  
    31  	return nil
    32  }
    33  
    34  func Sample(ctx *runtime.Task, funcExpr *ast.CallExpr) *errchain.PlError {
    35  	if len(funcExpr.Param) != 1 {
    36  		return runtime.NewRunError(ctx, fmt.Sprintf(
    37  			"func %s expects 1 arg", funcExpr.Name), funcExpr.NamePos)
    38  	}
    39  
    40  	var probability float64
    41  
    42  	switch funcExpr.Param[0].NodeType { //nolint:exhaustive
    43  	case ast.TypeArithmeticExpr:
    44  		res, dtype, err := runtime.RunArithmeticExpr(ctx, funcExpr.Param[0].ArithmeticExpr())
    45  		if err != nil {
    46  			return err
    47  		}
    48  		if dtype != ast.Float && dtype != ast.Int {
    49  			return runtime.NewRunError(ctx,
    50  				"data type of the result of arithmetic expression is neither float nor int",
    51  				funcExpr.Param[0].StartPos())
    52  		}
    53  		p, ok := res.(float64)
    54  		if !ok {
    55  			return runtime.NewRunError(ctx, "failed to convert expression evaluation result to float",
    56  				funcExpr.Param[0].StartPos())
    57  		}
    58  		probability = p
    59  
    60  	case ast.TypeFloatLiteral:
    61  		probability = funcExpr.Param[0].FloatLiteral().Val
    62  	case ast.TypeIntegerLiteral:
    63  		probability = float64(funcExpr.Param[0].IntegerLiteral().Val)
    64  	default:
    65  		return runtime.NewRunError(ctx, fmt.Sprintf(
    66  			"expect NumberLiteral or ArithmeticExpr, got %s", funcExpr.Param[0].NodeType),
    67  			funcExpr.Param[0].StartPos())
    68  	}
    69  
    70  	if probability < 0 || probability > 1 {
    71  		l.Error("sampling probability should be in the range [0, 1]")
    72  		return nil
    73  		// return runtime.NewRunError(ctx,
    74  		// 	"sampling probability should be in the range [0, 1]", funcExpr.NamePos)
    75  	}
    76  
    77  	res := time.Now().UnixMicro()%100 <= int64(probability*100)
    78  	ctx.Regs.ReturnAppend(res, ast.Bool)
    79  	return nil
    80  }