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

     1  package quic
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/metacubex/quic-go/internal/utils"
     9  	"github.com/metacubex/quic-go/internal/wire"
    10  
    11  	. "github.com/onsi/ginkgo/v2"
    12  	. "github.com/onsi/gomega"
    13  )
    14  
    15  var _ = Describe("Datagram Queue", func() {
    16  	var queue *datagramQueue
    17  	var queued chan struct{}
    18  
    19  	BeforeEach(func() {
    20  		queued = make(chan struct{}, 100)
    21  		queue = newDatagramQueue(func() { queued <- struct{}{} }, utils.DefaultLogger)
    22  	})
    23  
    24  	Context("sending", func() {
    25  		It("returns nil when there's no datagram to send", func() {
    26  			Expect(queue.Peek()).To(BeNil())
    27  		})
    28  
    29  		It("queues a datagram", func() {
    30  			frame := &wire.DatagramFrame{Data: []byte("foobar")}
    31  			Expect(queue.Add(frame)).To(Succeed())
    32  			Expect(queued).To(HaveLen(1))
    33  			f := queue.Peek()
    34  			Expect(f.Data).To(Equal([]byte("foobar")))
    35  			queue.Pop()
    36  			Expect(queue.Peek()).To(BeNil())
    37  		})
    38  
    39  		It("blocks when the maximum number of datagrams have been queued", func() {
    40  			for i := 0; i < maxDatagramSendQueueLen; i++ {
    41  				Expect(queue.Add(&wire.DatagramFrame{Data: []byte{0}})).To(Succeed())
    42  			}
    43  			errChan := make(chan error, 1)
    44  			go func() {
    45  				defer GinkgoRecover()
    46  				errChan <- queue.Add(&wire.DatagramFrame{Data: []byte("foobar")})
    47  			}()
    48  			Consistently(errChan, 50*time.Millisecond).ShouldNot(Receive())
    49  			Expect(queue.Peek()).ToNot(BeNil())
    50  			Consistently(errChan, 50*time.Millisecond).ShouldNot(Receive())
    51  			queue.Pop()
    52  			Eventually(errChan).Should(Receive(BeNil()))
    53  			for i := 1; i < maxDatagramSendQueueLen; i++ {
    54  				queue.Pop()
    55  			}
    56  			f := queue.Peek()
    57  			Expect(f).ToNot(BeNil())
    58  			Expect(f.Data).To(Equal([]byte("foobar")))
    59  		})
    60  
    61  		It("returns the same datagram multiple times, when Pop isn't called", func() {
    62  			Expect(queue.Add(&wire.DatagramFrame{Data: []byte("foo")})).To(Succeed())
    63  			Expect(queue.Add(&wire.DatagramFrame{Data: []byte("bar")})).To(Succeed())
    64  
    65  			Eventually(queued).Should(HaveLen(2))
    66  			f := queue.Peek()
    67  			Expect(f.Data).To(Equal([]byte("foo")))
    68  			Expect(queue.Peek()).To(Equal(f))
    69  			Expect(queue.Peek()).To(Equal(f))
    70  			queue.Pop()
    71  			f = queue.Peek()
    72  			Expect(f).ToNot(BeNil())
    73  			Expect(f.Data).To(Equal([]byte("bar")))
    74  		})
    75  
    76  		It("closes", func() {
    77  			for i := 0; i < maxDatagramSendQueueLen; i++ {
    78  				Expect(queue.Add(&wire.DatagramFrame{Data: []byte("foo")})).To(Succeed())
    79  			}
    80  			errChan := make(chan error, 1)
    81  			go func() {
    82  				defer GinkgoRecover()
    83  				errChan <- queue.Add(&wire.DatagramFrame{Data: []byte("foo")})
    84  			}()
    85  			Consistently(errChan, 25*time.Millisecond).ShouldNot(Receive())
    86  			testErr := errors.New("test error")
    87  			queue.CloseWithError(testErr)
    88  			Eventually(errChan).Should(Receive(MatchError(testErr)))
    89  		})
    90  	})
    91  
    92  	Context("receiving", func() {
    93  		It("receives DATAGRAM frames", func() {
    94  			queue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("foo")})
    95  			queue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("bar")})
    96  			data, err := queue.Receive(context.Background())
    97  			Expect(err).ToNot(HaveOccurred())
    98  			Expect(data).To(Equal([]byte("foo")))
    99  			data, err = queue.Receive(context.Background())
   100  			Expect(err).ToNot(HaveOccurred())
   101  			Expect(data).To(Equal([]byte("bar")))
   102  		})
   103  
   104  		It("blocks until a frame is received", func() {
   105  			c := make(chan []byte, 1)
   106  			go func() {
   107  				defer GinkgoRecover()
   108  				data, err := queue.Receive(context.Background())
   109  				Expect(err).ToNot(HaveOccurred())
   110  				c <- data
   111  			}()
   112  
   113  			Consistently(c).ShouldNot(Receive())
   114  			queue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("foobar")})
   115  			Eventually(c).Should(Receive(Equal([]byte("foobar"))))
   116  		})
   117  
   118  		It("blocks until context is done", func() {
   119  			ctx, cancel := context.WithCancel(context.Background())
   120  			errChan := make(chan error)
   121  			go func() {
   122  				defer GinkgoRecover()
   123  				_, err := queue.Receive(ctx)
   124  				errChan <- err
   125  			}()
   126  
   127  			Consistently(errChan).ShouldNot(Receive())
   128  			cancel()
   129  			Eventually(errChan).Should(Receive(Equal(context.Canceled)))
   130  		})
   131  
   132  		It("closes", func() {
   133  			errChan := make(chan error, 1)
   134  			go func() {
   135  				defer GinkgoRecover()
   136  				_, err := queue.Receive(context.Background())
   137  				errChan <- err
   138  			}()
   139  
   140  			Consistently(errChan).ShouldNot(Receive())
   141  			queue.CloseWithError(errors.New("test error"))
   142  			Eventually(errChan).Should(Receive(MatchError("test error")))
   143  		})
   144  	})
   145  })