trpc.group/trpc-go/trpc-go@v1.0.3/stream/flow_control.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  	"reflect"
    19  	"sync/atomic"
    20  )
    21  
    22  // sendControl is the behavior control of the sender.
    23  type sendControl struct {
    24  	// The largest window is defined as uint32, but there may be negative numbers, so it is int64.
    25  	window int64
    26  	// When sending window to less than or equal to 0, ch will block, when receiving window update, ch will unblock.
    27  	ch chan struct{}
    28  	// waits wait for data to arrive or the stream to end.
    29  	waits []reflect.SelectCase
    30  }
    31  
    32  // feedback is the feedback type.
    33  type feedback func(uint32) error
    34  
    35  func newSendControl(window uint32, dones ...<-chan struct{}) *sendControl {
    36  	s := &sendControl{
    37  		window: int64(window),
    38  		ch:     make(chan struct{}, 1),
    39  	}
    40  	s.waits = []reflect.SelectCase{{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(s.ch)}}
    41  	for _, done := range dones {
    42  		s.waits = append(s.waits, reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(done)})
    43  	}
    44  	return s
    45  }
    46  
    47  // GetWindow gets the sending window of a certain size, if it can't get it, it will block.
    48  // precision is not guaranteed, may be negative.
    49  func (s *sendControl) GetWindow(w uint32) error {
    50  	for w := int64(w); ; {
    51  		// First determine the currently available port, if the available window is <= 0, wait for the window to update.
    52  		// If it is greater than 0, subtract this window and return.
    53  		// Note that it may become a negative number after subtraction.
    54  		if atomic.LoadInt64(&s.window) > 0 {
    55  			atomic.AddInt64(&s.window, -w)
    56  			return nil
    57  		}
    58  		if chosen, _, _ := reflect.Select(s.waits); chosen == 0 {
    59  			// received data
    60  			continue
    61  		}
    62  		// stream closed
    63  		return errors.New(streamClosed)
    64  	}
    65  }
    66  
    67  // UpdateWindow is used to update the window after receiving the feedback.
    68  func (s *sendControl) UpdateWindow(increment uint32) {
    69  	updatedWindow := atomic.AddInt64(&s.window, int64(increment))
    70  	if !checkUpdate(updatedWindow, int64(increment)) {
    71  		return
    72  	}
    73  	select {
    74  	// Signal a blocked get window request
    75  	case s.ch <- struct{}{}:
    76  	default:
    77  	}
    78  }
    79  
    80  func checkUpdate(updatedWindow, increment int64) bool {
    81  	return (updatedWindow-increment <= 0) && (updatedWindow > 0)
    82  }
    83  
    84  // receiveControl represents the flow control statistics from the perspective of the receiving end.
    85  type receiveControl struct {
    86  	buffer    uint32   // upper limit.
    87  	unUpdated uint32   // Consumed, no window update sent.
    88  	fb        feedback // function for feedback.
    89  }
    90  
    91  func newReceiveControl(buffer uint32, fb feedback) *receiveControl {
    92  	return &receiveControl{
    93  		buffer: buffer,
    94  		fb:     fb,
    95  	}
    96  }
    97  
    98  // OnRecv application is called when data is received, and the window is updated.
    99  func (r *receiveControl) OnRecv(n uint32) error {
   100  	r.unUpdated += n
   101  	if r.unUpdated >= r.buffer/4 {
   102  		increment := r.unUpdated
   103  		r.unUpdated = 0
   104  		if r.fb != nil {
   105  			return r.fb(increment)
   106  		}
   107  		return nil
   108  	}
   109  	return nil
   110  }