github.com/kubewharf/katalyst-core@v0.5.3/pkg/util/asyncworker/async_workers_test.go (about)

     1  /*
     2  Copyright 2022 The Katalyst Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package asyncworker
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"testing"
    23  	"time"
    24  
    25  	"github.com/stretchr/testify/require"
    26  
    27  	"github.com/kubewharf/katalyst-core/pkg/metrics"
    28  )
    29  
    30  func TestAsyncWorkers(t *testing.T) {
    31  	t.Parallel()
    32  
    33  	rt := require.New(t)
    34  
    35  	asw := NewAsyncWorkers("test", metrics.DummyMetrics{})
    36  
    37  	result, a, b, c, d, e, f, g, h := 0, 1, 2, 3, 4, 5, 6, 7, 8
    38  
    39  	timeoutSeconds := 100 * time.Millisecond
    40  	fn := func(ctx context.Context, params ...interface{}) error {
    41  		if len(params) != 2 {
    42  			return fmt.Errorf("invalid params")
    43  		}
    44  
    45  		time.Sleep(5 * time.Millisecond)
    46  		p1Int := params[0].(int)
    47  		p2Int := params[1].(int)
    48  		result = p1Int + p2Int
    49  		_ = EmitAsyncedMetrics(ctx, metrics.MetricTag{})
    50  		return nil
    51  	}
    52  
    53  	work1Name := "work1"
    54  	work1DeliveredAt := time.Now()
    55  	work1 := &Work{
    56  		Fn:          fn,
    57  		Params:      []interface{}{a, b},
    58  		DeliveredAt: work1DeliveredAt,
    59  	}
    60  
    61  	err := asw.AddWork(work1Name, work1, DuplicateWorkPolicyOverride)
    62  	rt.Nil(err)
    63  	asw.workLock.Lock()
    64  	rt.NotNil(asw.workStatuses[work1Name])
    65  	asw.workLock.Unlock()
    66  
    67  	asw.workLock.Lock()
    68  	for asw.workStatuses[work1Name].working {
    69  		asw.workLock.Unlock()
    70  		time.Sleep(10 * time.Millisecond)
    71  
    72  		if time.Now().Sub(work1DeliveredAt) > timeoutSeconds {
    73  			rt.Failf("%s timeout", work1Name)
    74  		}
    75  
    76  		asw.workLock.Lock()
    77  	}
    78  	asw.workLock.Unlock()
    79  
    80  	rt.Equal(result, a+b)
    81  
    82  	work2Name := "work2"
    83  	work2DeliveredAt := time.Now()
    84  	work2 := &Work{
    85  		Fn:          fn,
    86  		Params:      []interface{}{c, d},
    87  		DeliveredAt: work2DeliveredAt,
    88  	}
    89  
    90  	err = asw.AddWork(work2Name, work2, DuplicateWorkPolicyOverride)
    91  	rt.Nil(err)
    92  	asw.workLock.Lock()
    93  	rt.NotNil(asw.workStatuses[work2Name])
    94  	rt.Nil(asw.lastUndeliveredWork[work2Name])
    95  	asw.workLock.Unlock()
    96  
    97  	work3DeliveredAt := time.Now()
    98  	work3 := &Work{
    99  		Fn:          fn,
   100  		Params:      []interface{}{e, f},
   101  		DeliveredAt: work3DeliveredAt,
   102  	}
   103  
   104  	err = asw.AddWork(work2Name, work3, DuplicateWorkPolicyOverride)
   105  	rt.Nil(err)
   106  	asw.workLock.Lock()
   107  	rt.NotNil(asw.lastUndeliveredWork[work2Name])
   108  	asw.workLock.Unlock()
   109  
   110  	work4DeliveredAt := time.Now()
   111  	work4 := &Work{
   112  		Fn:          fn,
   113  		Params:      []interface{}{g, h},
   114  		DeliveredAt: work4DeliveredAt,
   115  	}
   116  
   117  	err = asw.AddWork(work2Name, work4, DuplicateWorkPolicyDiscard)
   118  	rt.Nil(err)
   119  	asw.workLock.Lock()
   120  	rt.Equal(work3, asw.lastUndeliveredWork[work2Name])
   121  	asw.workLock.Unlock()
   122  
   123  	asw.workLock.Lock()
   124  	for asw.workStatuses[work2Name].working {
   125  		asw.workLock.Unlock()
   126  		time.Sleep(10 * time.Millisecond)
   127  
   128  		if time.Now().Sub(work1DeliveredAt) > 3*timeoutSeconds {
   129  			rt.Failf("%s timeout", work2Name)
   130  		}
   131  
   132  		asw.workLock.Lock()
   133  	}
   134  	asw.workLock.Unlock()
   135  
   136  	rt.Equal(result, e+f)
   137  }