github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/plan/function/builtin/unary/sleep.go (about)

     1  // Copyright 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package unary
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    19  	"github.com/matrixorigin/matrixone/pkg/container/types"
    20  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    21  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    22  	"time"
    23  )
    24  
    25  type number interface {
    26  	uint64 | float64
    27  }
    28  
    29  func Sleep[T number](vs []*vector.Vector, proc *process.Process) (rs *vector.Vector, err error) {
    30  	defer func() {
    31  		if err != nil && rs != nil {
    32  			rs.Free(proc.Mp())
    33  		}
    34  	}()
    35  	resultType := types.T_uint8.ToType()
    36  	inputs := vs[0]
    37  	if inputs.Nsp.Any() {
    38  		err = moerr.NewInvalidArg(proc.Ctx, "sleep", "input contains null")
    39  		return
    40  	}
    41  	sleepSlice := vector.MustTCols[T](inputs)
    42  	if checkNegative(sleepSlice) {
    43  		err = moerr.NewInvalidArg(proc.Ctx, "sleep", "input contains negative")
    44  		return
    45  	}
    46  	if inputs.IsScalar() {
    47  		sleepSeconds := sleepSlice[0]
    48  		sleepNano := time.Nanosecond * time.Duration(sleepSeconds*1e9)
    49  		length := int64(inputs.Length())
    50  		if length == 1 {
    51  			rs = proc.AllocScalarVector(resultType)
    52  			result := vector.MustTCols[uint8](rs)
    53  			select {
    54  			case <-time.After(sleepNano):
    55  				result[0] = 0
    56  			case <-proc.Ctx.Done(): //query aborted
    57  				result[0] = 1
    58  			}
    59  			return
    60  		}
    61  		rs, err = proc.AllocVectorOfRows(resultType, length, nil)
    62  		if err != nil {
    63  			return
    64  		}
    65  		result := vector.MustTCols[uint8](rs)
    66  		for i := int64(0); i < length; i++ {
    67  			select {
    68  			case <-time.After(sleepNano):
    69  				result[i] = 0
    70  			case <-proc.Ctx.Done(): //query aborted
    71  				for ; i < length; i++ {
    72  					result[i] = 1
    73  				}
    74  				return
    75  			}
    76  		}
    77  		return
    78  	}
    79  	rs, err = proc.AllocVectorOfRows(resultType, int64(len(sleepSlice)), inputs.Nsp)
    80  	if err != nil {
    81  		return
    82  	}
    83  	result := vector.MustTCols[uint8](rs)
    84  	for i, sleepSeconds := range sleepSlice {
    85  		sleepNano := time.Nanosecond * time.Duration(sleepSeconds*1e9)
    86  		select {
    87  		case <-time.After(sleepNano):
    88  			result[i] = 0
    89  		case <-proc.Ctx.Done(): //query aborted
    90  			for ; i < len(sleepSlice); i++ {
    91  				result[i] = 1
    92  			}
    93  			return
    94  		}
    95  	}
    96  	return
    97  }
    98  
    99  func checkNegative[T number](rs []T) bool {
   100  	for _, v := range rs {
   101  		if v < 0 {
   102  			return true
   103  		}
   104  	}
   105  	return false
   106  }