github.com/metacubex/quic-go@v0.44.1-0.20240520163451-20b689a59136/http3/datagram.go (about)

     1  package http3
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  const maxQuarterStreamID = 1<<60 - 1
     9  
    10  const streamDatagramQueueLen = 32
    11  
    12  type datagrammer struct {
    13  	sendDatagram func([]byte) error
    14  
    15  	hasData chan struct{}
    16  	queue   [][]byte // TODO: use a ring buffer
    17  
    18  	mx         sync.Mutex
    19  	sendErr    error
    20  	receiveErr error
    21  }
    22  
    23  func newDatagrammer(sendDatagram func([]byte) error) *datagrammer {
    24  	return &datagrammer{
    25  		sendDatagram: sendDatagram,
    26  		hasData:      make(chan struct{}, 1),
    27  	}
    28  }
    29  
    30  func (d *datagrammer) SetReceiveError(err error) {
    31  	d.mx.Lock()
    32  	defer d.mx.Unlock()
    33  
    34  	d.receiveErr = err
    35  	d.signalHasData()
    36  }
    37  
    38  func (d *datagrammer) SetSendError(err error) {
    39  	d.mx.Lock()
    40  	defer d.mx.Unlock()
    41  
    42  	d.sendErr = err
    43  }
    44  
    45  func (d *datagrammer) Send(b []byte) error {
    46  	d.mx.Lock()
    47  	sendErr := d.sendErr
    48  	d.mx.Unlock()
    49  	if sendErr != nil {
    50  		return sendErr
    51  	}
    52  
    53  	return d.sendDatagram(b)
    54  }
    55  
    56  func (d *datagrammer) signalHasData() {
    57  	select {
    58  	case d.hasData <- struct{}{}:
    59  	default:
    60  	}
    61  }
    62  
    63  func (d *datagrammer) enqueue(data []byte) {
    64  	d.mx.Lock()
    65  	defer d.mx.Unlock()
    66  
    67  	if d.receiveErr != nil {
    68  		return
    69  	}
    70  	if len(d.queue) >= streamDatagramQueueLen {
    71  		return
    72  	}
    73  	d.queue = append(d.queue, data)
    74  	d.signalHasData()
    75  }
    76  
    77  func (d *datagrammer) Receive(ctx context.Context) ([]byte, error) {
    78  start:
    79  	d.mx.Lock()
    80  	if len(d.queue) >= 1 {
    81  		data := d.queue[0]
    82  		d.queue = d.queue[1:]
    83  		d.mx.Unlock()
    84  		return data, nil
    85  	}
    86  	if receiveErr := d.receiveErr; receiveErr != nil {
    87  		d.mx.Unlock()
    88  		return nil, receiveErr
    89  	}
    90  	d.mx.Unlock()
    91  
    92  	select {
    93  	case <-ctx.Done():
    94  		return nil, context.Cause(ctx)
    95  	case <-d.hasData:
    96  	}
    97  	goto start
    98  }