trpc.group/trpc-go/trpc-go@v1.0.2/stream/flow_control_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 stream
    15  
    16  import (
    17  	"errors"
    18  	"sync/atomic"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/stretchr/testify/assert"
    23  )
    24  
    25  // TestSendControl test flow control sender related implementation.
    26  func TestSendControl(t *testing.T) {
    27  	done := make(chan struct{})
    28  	sc := newSendControl(defaultInitWindowSize, done)
    29  	err := sc.GetWindow(100)
    30  	assert.Nil(t, err)
    31  
    32  	// Available window drops less than 0.
    33  	err = sc.GetWindow(uint32(defaultInitWindowSize - 99))
    34  	assert.Nil(t, err)
    35  
    36  	// block
    37  	t1 := time.Now()
    38  	go func() {
    39  		time.Sleep(500 * time.Millisecond)
    40  		sc.UpdateWindow(201)
    41  	}()
    42  	err = sc.GetWindow(200)
    43  	assert.Nil(t, err)
    44  	t2 := int64(time.Now().Sub(t1))
    45  	assert.GreaterOrEqual(t, t2, int64(500*time.Millisecond))
    46  }
    47  
    48  // TestReceiveControl test.
    49  func TestReceiveControl(t *testing.T) {
    50  	fb := func(uint32) error {
    51  		return nil
    52  	}
    53  	rc := newReceiveControl(defaultInitWindowSize, fb)
    54  	err := rc.OnRecv(100)
    55  	assert.Nil(t, err)
    56  
    57  	n := atomic.LoadUint32(&rc.left)
    58  	assert.Equal(t, defaultInitWindowSize-uint32(100), n)
    59  
    60  	// need to send updates.
    61  	err = rc.OnRecv(defaultInitWindowSize / 4)
    62  	assert.Nil(t, err)
    63  
    64  	// test for feedback errors.
    65  	fb = func(uint32) error {
    66  		return errors.New("feedback error")
    67  	}
    68  	err = rc.OnRecv(100)
    69  	assert.Nil(t, err)
    70  
    71  	rc = newReceiveControl(defaultInitWindowSize, fb)
    72  	err = rc.OnRecv(defaultInitWindowSize / 4)
    73  	assert.NotNil(t, err)
    74  }