github.com/sagernet/quic-go@v0.43.1-beta.1/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) (isDone bool) {
    31  	d.mx.Lock()
    32  	defer d.mx.Unlock()
    33  
    34  	d.receiveErr = err
    35  	d.signalHasData()
    36  	return d.sendErr != nil
    37  }
    38  
    39  func (d *datagrammer) SetSendError(err error) (isDone bool) {
    40  	d.mx.Lock()
    41  	defer d.mx.Unlock()
    42  
    43  	d.sendErr = err
    44  	return d.receiveErr != nil
    45  }
    46  
    47  func (d *datagrammer) Send(b []byte) error {
    48  	d.mx.Lock()
    49  	sendErr := d.sendErr
    50  	d.mx.Unlock()
    51  	if sendErr != nil {
    52  		return sendErr
    53  	}
    54  
    55  	return d.sendDatagram(b)
    56  }
    57  
    58  func (d *datagrammer) signalHasData() {
    59  	select {
    60  	case d.hasData <- struct{}{}:
    61  	default:
    62  	}
    63  }
    64  
    65  func (d *datagrammer) enqueue(data []byte) {
    66  	d.mx.Lock()
    67  	defer d.mx.Unlock()
    68  
    69  	if d.receiveErr != nil {
    70  		return
    71  	}
    72  	if len(d.queue) >= streamDatagramQueueLen {
    73  		return
    74  	}
    75  	d.queue = append(d.queue, data)
    76  	d.signalHasData()
    77  }
    78  
    79  func (d *datagrammer) Receive(ctx context.Context) ([]byte, error) {
    80  start:
    81  	d.mx.Lock()
    82  	if len(d.queue) >= 1 {
    83  		data := d.queue[0]
    84  		d.queue = d.queue[1:]
    85  		d.mx.Unlock()
    86  		return data, nil
    87  	}
    88  	if d.receiveErr != nil {
    89  		d.mx.Unlock()
    90  		return nil, d.receiveErr
    91  	}
    92  	d.mx.Unlock()
    93  
    94  	select {
    95  	case <-ctx.Done():
    96  		return nil, context.Cause(ctx)
    97  	case <-d.hasData:
    98  	}
    99  	goto start
   100  }