tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/mux/util_window.go (about)

     1  package mux
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // window represents the buffer available to clients
     9  // wishing to write to a channel.
    10  type window struct {
    11  	*sync.Cond
    12  	win          uint32 // RFC 4254 5.2 says the window size can grow to 2^32-1
    13  	writeWaiters int
    14  	closed       bool
    15  }
    16  
    17  // add adds win to the amount of window available
    18  // for consumers.
    19  func (w *window) add(win uint32) bool {
    20  	// a zero sized window adjust is a noop.
    21  	if win == 0 {
    22  		return true
    23  	}
    24  	w.L.Lock()
    25  	if w.win+win < win {
    26  		w.L.Unlock()
    27  		return false
    28  	}
    29  	w.win += win
    30  	// It is unusual that multiple goroutines would be attempting to reserve
    31  	// window space, but not guaranteed. Use broadcast to notify all waiters
    32  	// that additional window is available.
    33  	w.Broadcast()
    34  	w.L.Unlock()
    35  	return true
    36  }
    37  
    38  // close sets the window to closed, so all reservations fail
    39  // immediately.
    40  func (w *window) close() {
    41  	w.L.Lock()
    42  	w.closed = true
    43  	w.Broadcast()
    44  	w.L.Unlock()
    45  }
    46  
    47  // reserve reserves win from the available window capacity.
    48  // If no capacity remains, reserve will block. reserve may
    49  // return less than requested.
    50  func (w *window) reserve(win uint32) (uint32, error) {
    51  	var err error
    52  	w.L.Lock()
    53  	w.writeWaiters++
    54  	w.Broadcast()
    55  	for w.win == 0 && !w.closed {
    56  		w.Wait()
    57  	}
    58  	w.writeWaiters--
    59  	if w.win < win {
    60  		win = w.win
    61  	}
    62  	w.win -= win
    63  	if w.closed {
    64  		err = io.EOF
    65  	}
    66  	w.L.Unlock()
    67  	return win, err
    68  }
    69  
    70  // waitWriterBlocked waits until some goroutine is blocked for further
    71  // writes. It is used in tests only.
    72  func (w *window) waitWriterBlocked() {
    73  	w.Cond.L.Lock()
    74  	for w.writeWaiters == 0 {
    75  		w.Cond.Wait()
    76  	}
    77  	w.Cond.L.Unlock()
    78  }