trpc.group/trpc-go/trpc-go@v1.0.3/internal/queue/queue_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package queue
    15  
    16  import (
    17  	"sync"
    18  	"testing"
    19  	"time"
    20  
    21  	"github.com/stretchr/testify/assert"
    22  )
    23  
    24  // TestQueue test queue related functions
    25  func TestQueue(t *testing.T) {
    26  	done := make(chan struct{})
    27  	q := New[[]byte](done)
    28  
    29  	// Get it normally
    30  	v := []byte("hello world")
    31  	q.Put(v)
    32  	e, ok := q.Get()
    33  	assert.True(t, ok)
    34  	assert.Equal(t, []byte("hello world"), e)
    35  
    36  	// no data blocking
    37  	t1 := time.Now()
    38  	go func() {
    39  		time.Sleep(500 * time.Millisecond)
    40  		q.Put(v)
    41  	}()
    42  	e, ok = q.Get()
    43  
    44  	assert.True(t, ok)
    45  	assert.Equal(t, []byte("hello world"), e)
    46  	t2 := int64(time.Now().Sub(t1))
    47  	assert.GreaterOrEqual(t, t2, int64(500*time.Millisecond))
    48  
    49  	// queue Done
    50  	time.Sleep(200 * time.Millisecond)
    51  	close(done)
    52  	e, ok = q.Get()
    53  	assert.False(t, ok)
    54  	assert.Nil(t, e)
    55  }
    56  
    57  // TestConcurrentQueue test queue concurrency
    58  func TestConcurrentQueue(t *testing.T) {
    59  	done := make(chan struct{})
    60  	q := New[[]byte](done)
    61  	wg := &sync.WaitGroup{}
    62  	wg.Add(3)
    63  
    64  	// a goroutine write
    65  	go func() {
    66  		defer wg.Done()
    67  		for i := 0; i < 5000; i++ {
    68  			v := []byte("hello world")
    69  			q.Put(v)
    70  		}
    71  	}()
    72  
    73  	// write another goroutine
    74  	go func() {
    75  		defer wg.Done()
    76  		for i := 0; i < 5000; i++ {
    77  			v := []byte("hello world")
    78  			q.Put(v)
    79  		}
    80  	}()
    81  
    82  	// a goroutine read
    83  	go func() {
    84  		defer wg.Done()
    85  		for i := 0; i < 10000; i++ {
    86  			e, ok := q.Get()
    87  			assert.True(t, ok)
    88  			assert.Equal(t, []byte("hello world"), e)
    89  		}
    90  	}()
    91  	wg.Wait()
    92  }