add a context to Connection.ReceiveMessage (#3926)

* add context to ReceiveMessage

* add newlines

---------

Co-authored-by: Marten Seemann <martenseemann@gmail.com>

connection.go

@@ -2302,11 +2302,11 @@ func (s *connection) SendMessage(p []byte) error {
 	return s.datagramQueue.AddAndWait(f)
 }
 
-func (s *connection) ReceiveMessage() ([]byte, error) {
+func (s *connection) ReceiveMessage(ctx context.Context) ([]byte, error) {
 	if !s.config.EnableDatagrams {
 		return nil, errors.New("datagram support disabled")
 	}
-	return s.datagramQueue.Receive()
+	return s.datagramQueue.Receive(ctx)
 }
 
 func (s *connection) LocalAddr() net.Addr {

datagram_queue.go

@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"context"
 	"sync"
 
 	"github.com/quic-go/quic-go/internal/protocol"
@@ -98,7 +99,7 @@ func (h *datagramQueue) HandleDatagramFrame(f *wire.DatagramFrame) {
 }
 
 // Receive gets a received DATAGRAM frame.
-func (h *datagramQueue) Receive() ([]byte, error) {
+func (h *datagramQueue) Receive(ctx context.Context) ([]byte, error) {
 	for {
 		h.rcvMx.Lock()
 		if len(h.rcvQueue) > 0 {
@@ -113,6 +114,8 @@ func (h *datagramQueue) Receive() ([]byte, error) {
 			continue
 		case <-h.closed:
 			return nil, h.closeErr
+		case <-ctx.Done():
+			return nil, ctx.Err()
 		}
 	}
 }

datagram_queue_test.go

@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"context"
 	"errors"
 
 	"github.com/quic-go/quic-go/internal/utils"
@@ -81,10 +82,10 @@ var _ = Describe("Datagram Queue", func() {
 		It("receives DATAGRAM frames", func() {
 			queue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("foo")})
 			queue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("bar")})
-			data, err := queue.Receive()
+			data, err := queue.Receive(context.Background())
 			Expect(err).ToNot(HaveOccurred())
 			Expect(data).To(Equal([]byte("foo")))
-			data, err = queue.Receive()
+			data, err = queue.Receive(context.Background())
 			Expect(err).ToNot(HaveOccurred())
 			Expect(data).To(Equal([]byte("bar")))
 		})
@@ -93,7 +94,7 @@ var _ = Describe("Datagram Queue", func() {
 			c := make(chan []byte, 1)
 			go func() {
 				defer GinkgoRecover()
-				data, err := queue.Receive()
+				data, err := queue.Receive(context.Background())
 				Expect(err).ToNot(HaveOccurred())
 				c <- data
 			}()
@@ -103,11 +104,25 @@ var _ = Describe("Datagram Queue", func() {
 			Eventually(c).Should(Receive(Equal([]byte("foobar"))))
 		})
 
+		It("blocks until context is done", func() {
+			ctx, cancel := context.WithCancel(context.Background())
+			errChan := make(chan error)
+			go func() {
+				defer GinkgoRecover()
+				_, err := queue.Receive(ctx)
+				errChan <- err
+			}()
+
+			Consistently(errChan).ShouldNot(Receive())
+			cancel()
+			Eventually(errChan).Should(Receive(Equal(context.Canceled)))
+		})
+
 		It("closes", func() {
 			errChan := make(chan error, 1)
 			go func() {
 				defer GinkgoRecover()
-				_, err := queue.Receive()
+				_, err := queue.Receive(context.Background())
 				errChan <- err
 			}()
 

integrationtests/self/datagram_test.go

@@ -120,7 +120,7 @@ var _ = Describe("Datagram test", func() {
 		for {
 			// Close the connection if no message is received for 100 ms.
 			timer := time.AfterFunc(scaleDuration(100*time.Millisecond), func() { conn.CloseWithError(0, "") })
-			if _, err := conn.ReceiveMessage(); err != nil {
+			if _, err := conn.ReceiveMessage(context.Background()); err != nil {
 				break
 			}
 			timer.Stop()

interface.go

@@ -185,7 +185,7 @@ type Connection interface {
 	// SendMessage sends a message as a datagram, as specified in RFC 9221.
 	SendMessage([]byte) error
 	// ReceiveMessage gets a message received in a datagram, as specified in RFC 9221.
-	ReceiveMessage() ([]byte, error)
+	ReceiveMessage(context.Context) ([]byte, error)
 }
 
 // An EarlyConnection is a connection that is handshaking.