github.com/MerlinKodo/quic-go@v0.39.2/datagram_queue.go (about)

     1  package quic
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/MerlinKodo/quic-go/internal/protocol"
     8  	"github.com/MerlinKodo/quic-go/internal/utils"
     9  	"github.com/MerlinKodo/quic-go/internal/wire"
    10  )
    11  
    12  const DatagramFrameMaxPeekTimes = 10
    13  
    14  type datagramQueue struct {
    15  	sendQueue chan *wire.DatagramFrame
    16  	nextFrame *wire.DatagramFrame
    17  	peekTimes int
    18  
    19  	rcvMx    sync.Mutex
    20  	rcvQueue [][]byte
    21  	rcvd     chan struct{} // used to notify Receive that a new datagram was received
    22  
    23  	closeErr error
    24  	closed   chan struct{}
    25  
    26  	hasData func()
    27  
    28  	dequeued chan error
    29  
    30  	logger utils.Logger
    31  }
    32  
    33  func newDatagramQueue(hasData func(), logger utils.Logger) *datagramQueue {
    34  	return &datagramQueue{
    35  		hasData:   hasData,
    36  		sendQueue: make(chan *wire.DatagramFrame, 1),
    37  		rcvd:      make(chan struct{}, 1),
    38  		dequeued:  make(chan error),
    39  		closed:    make(chan struct{}),
    40  		logger:    logger,
    41  	}
    42  }
    43  
    44  // AddAndWait queues a new DATAGRAM frame for sending.
    45  // It blocks until the frame has been dequeued.
    46  func (h *datagramQueue) AddAndWait(f *wire.DatagramFrame) error {
    47  	select {
    48  	case h.sendQueue <- f:
    49  		h.hasData()
    50  	case <-h.closed:
    51  		return h.closeErr
    52  	}
    53  
    54  	select {
    55  	case err := <-h.dequeued:
    56  		return err
    57  	case <-h.closed:
    58  		return h.closeErr
    59  	}
    60  }
    61  
    62  // Peek gets the next DATAGRAM frame for sending.
    63  // If actually sent out, Pop needs to be called before the next call to Peek.
    64  func (h *datagramQueue) Peek() *wire.DatagramFrame {
    65  	if h.nextFrame != nil {
    66  		h.peekTimes++
    67  		if h.peekTimes > DatagramFrameMaxPeekTimes {
    68  			if h.logger.Debug() {
    69  				h.logger.Debugf("Discarded DATAGRAM frame (%d bytes payload)", len(h.nextFrame.Data))
    70  			}
    71  			h.nextFrame = nil
    72  		} else {
    73  			return h.nextFrame
    74  		}
    75  	}
    76  	select {
    77  	case h.nextFrame = <-h.sendQueue:
    78  	default:
    79  		return nil
    80  	}
    81  	h.peekTimes = 0
    82  	return h.nextFrame
    83  }
    84  
    85  func (h *datagramQueue) Pop() {
    86  	if h.nextFrame == nil {
    87  		panic("datagramQueue BUG: Pop called for nil frame")
    88  	}
    89  	h.nextFrame = nil
    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) < protocol.DatagramRcvQueueLen {
    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 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  }