github.com/sagernet/quic-go@v0.43.1-beta.1/datagram_queue.go (about)

     1  package quic
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/sagernet/quic-go/internal/utils"
     8  	"github.com/sagernet/quic-go/internal/utils/ringbuffer"
     9  	"github.com/sagernet/quic-go/internal/wire"
    10  )
    11  
    12  const (
    13  	maxDatagramSendQueueLen = 32
    14  	maxDatagramRcvQueueLen  = 128
    15  )
    16  
    17  type datagramQueue struct {
    18  	sendMx    sync.Mutex
    19  	sendQueue ringbuffer.RingBuffer[*wire.DatagramFrame]
    20  	sent      chan struct{} // used to notify Add that a datagram was dequeued
    21  
    22  	rcvMx    sync.Mutex
    23  	rcvQueue [][]byte
    24  	rcvd     chan struct{} // used to notify Receive that a new datagram was received
    25  
    26  	closeErr error
    27  	closed   chan struct{}
    28  
    29  	hasData func()
    30  
    31  	logger utils.Logger
    32  }
    33  
    34  func newDatagramQueue(hasData func(), logger utils.Logger) *datagramQueue {
    35  	return &datagramQueue{
    36  		hasData: hasData,
    37  		rcvd:    make(chan struct{}, 1),
    38  		sent:    make(chan struct{}, 1),
    39  		closed:  make(chan struct{}),
    40  		logger:  logger,
    41  	}
    42  }
    43  
    44  // Add queues a new DATAGRAM frame for sending.
    45  // Up to 32 DATAGRAM frames will be queued.
    46  // Once that limit is reached, Add blocks until the queue size has reduced.
    47  func (h *datagramQueue) Add(f *wire.DatagramFrame) error {
    48  	h.sendMx.Lock()
    49  
    50  	for {
    51  		if h.sendQueue.Len() < maxDatagramSendQueueLen {
    52  			h.sendQueue.PushBack(f)
    53  			h.sendMx.Unlock()
    54  			h.hasData()
    55  			return nil
    56  		}
    57  		select {
    58  		case <-h.sent: // drain the queue so we don't loop immediately
    59  		default:
    60  		}
    61  		h.sendMx.Unlock()
    62  		select {
    63  		case <-h.closed:
    64  			return h.closeErr
    65  		case <-h.sent:
    66  		}
    67  		h.sendMx.Lock()
    68  	}
    69  }
    70  
    71  // Peek gets the next DATAGRAM frame for sending.
    72  // If actually sent out, Pop needs to be called before the next call to Peek.
    73  func (h *datagramQueue) Peek() *wire.DatagramFrame {
    74  	h.sendMx.Lock()
    75  	defer h.sendMx.Unlock()
    76  	if h.sendQueue.Empty() {
    77  		return nil
    78  	}
    79  	return h.sendQueue.PeekFront()
    80  }
    81  
    82  func (h *datagramQueue) Pop() {
    83  	h.sendMx.Lock()
    84  	defer h.sendMx.Unlock()
    85  	_ = h.sendQueue.PopFront()
    86  	select {
    87  	case h.sent <- struct{}{}:
    88  	default:
    89  	}
    90  }
    91  
    92  // HandleDatagramFrame handles a received DATAGRAM frame.
    93  func (h *datagramQueue) HandleDatagramFrame(f *wire.DatagramFrame) {
    94  	data := make([]byte, len(f.Data))
    95  	copy(data, f.Data)
    96  	var queued bool
    97  	h.rcvMx.Lock()
    98  	if len(h.rcvQueue) < maxDatagramRcvQueueLen {
    99  		h.rcvQueue = append(h.rcvQueue, data)
   100  		queued = true
   101  		select {
   102  		case h.rcvd <- struct{}{}:
   103  		default:
   104  		}
   105  	}
   106  	h.rcvMx.Unlock()
   107  	if !queued && h.logger.Debug() {
   108  		h.logger.Debugf("Discarding received DATAGRAM frame (%d bytes payload)", len(f.Data))
   109  	}
   110  }
   111  
   112  // Receive gets a received DATAGRAM frame.
   113  func (h *datagramQueue) Receive(ctx context.Context) ([]byte, error) {
   114  	for {
   115  		h.rcvMx.Lock()
   116  		if len(h.rcvQueue) > 0 {
   117  			data := h.rcvQueue[0]
   118  			h.rcvQueue = h.rcvQueue[1:]
   119  			h.rcvMx.Unlock()
   120  			return data, nil
   121  		}
   122  		h.rcvMx.Unlock()
   123  		select {
   124  		case <-h.rcvd:
   125  			continue
   126  		case <-h.closed:
   127  			return nil, h.closeErr
   128  		case <-ctx.Done():
   129  			return nil, ctx.Err()
   130  		}
   131  	}
   132  }
   133  
   134  func (h *datagramQueue) CloseWithError(e error) {
   135  	h.closeErr = e
   136  	close(h.closed)
   137  }