github.com/tumi8/quic-go@v0.37.4-tum/datagram_queue.go (about)

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