github.com/daeuniverse/quic-go@v0.0.0-20240413031024-943f218e0810/integrationtests/self/stream_test.go (about)

     1  package self_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"sync"
     9  
    10  	. "github.com/onsi/ginkgo/v2"
    11  	. "github.com/onsi/gomega"
    12  	"github.com/daeuniverse/quic-go"
    13  )
    14  
    15  var _ = Describe("Bidirectional streams", func() {
    16  	const numStreams = 300
    17  
    18  	var (
    19  		server     *quic.Listener
    20  		serverAddr string
    21  	)
    22  
    23  	BeforeEach(func() {
    24  		var err error
    25  		server, err = quic.ListenAddr("localhost:0", getTLSConfig(), getQuicConfig(nil))
    26  		Expect(err).ToNot(HaveOccurred())
    27  		serverAddr = fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port)
    28  	})
    29  
    30  	AfterEach(func() {
    31  		server.Close()
    32  	})
    33  
    34  	runSendingPeer := func(conn quic.Connection) {
    35  		var wg sync.WaitGroup
    36  		wg.Add(numStreams)
    37  		for i := 0; i < numStreams; i++ {
    38  			str, err := conn.OpenStreamSync(context.Background())
    39  			Expect(err).ToNot(HaveOccurred())
    40  			data := GeneratePRData(25 * i)
    41  			go func() {
    42  				defer GinkgoRecover()
    43  				_, err := str.Write(data)
    44  				Expect(err).ToNot(HaveOccurred())
    45  				Expect(str.Close()).To(Succeed())
    46  			}()
    47  			go func() {
    48  				defer GinkgoRecover()
    49  				defer wg.Done()
    50  				dataRead, err := io.ReadAll(str)
    51  				Expect(err).ToNot(HaveOccurred())
    52  				Expect(dataRead).To(Equal(data))
    53  			}()
    54  		}
    55  		wg.Wait()
    56  	}
    57  
    58  	runReceivingPeer := func(conn quic.Connection) {
    59  		var wg sync.WaitGroup
    60  		wg.Add(numStreams)
    61  		for i := 0; i < numStreams; i++ {
    62  			str, err := conn.AcceptStream(context.Background())
    63  			Expect(err).ToNot(HaveOccurred())
    64  			go func() {
    65  				defer GinkgoRecover()
    66  				defer wg.Done()
    67  				// shouldn't use io.Copy here
    68  				// we should read from the stream as early as possible, to free flow control credit
    69  				data, err := io.ReadAll(str)
    70  				Expect(err).ToNot(HaveOccurred())
    71  				_, err = str.Write(data)
    72  				Expect(err).ToNot(HaveOccurred())
    73  				Expect(str.Close()).To(Succeed())
    74  			}()
    75  		}
    76  		wg.Wait()
    77  	}
    78  
    79  	It(fmt.Sprintf("client opening %d streams to a server", numStreams), func() {
    80  		var conn quic.Connection
    81  		go func() {
    82  			defer GinkgoRecover()
    83  			var err error
    84  			conn, err = server.Accept(context.Background())
    85  			Expect(err).ToNot(HaveOccurred())
    86  			runReceivingPeer(conn)
    87  		}()
    88  
    89  		client, err := quic.DialAddr(
    90  			context.Background(),
    91  			serverAddr,
    92  			getTLSClientConfig(),
    93  			getQuicConfig(nil),
    94  		)
    95  		Expect(err).ToNot(HaveOccurred())
    96  		runSendingPeer(client)
    97  		client.CloseWithError(0, "")
    98  		<-conn.Context().Done()
    99  	})
   100  
   101  	It(fmt.Sprintf("server opening %d streams to a client", numStreams), func() {
   102  		go func() {
   103  			defer GinkgoRecover()
   104  			conn, err := server.Accept(context.Background())
   105  			Expect(err).ToNot(HaveOccurred())
   106  			runSendingPeer(conn)
   107  			conn.CloseWithError(0, "")
   108  		}()
   109  
   110  		client, err := quic.DialAddr(
   111  			context.Background(),
   112  			serverAddr,
   113  			getTLSClientConfig(),
   114  			getQuicConfig(nil),
   115  		)
   116  		Expect(err).ToNot(HaveOccurred())
   117  		runReceivingPeer(client)
   118  		Eventually(client.Context().Done()).Should(BeClosed())
   119  	})
   120  
   121  	It(fmt.Sprintf("client and server opening %d each and sending data to the peer", numStreams), func() {
   122  		done1 := make(chan struct{})
   123  		go func() {
   124  			defer GinkgoRecover()
   125  			conn, err := server.Accept(context.Background())
   126  			Expect(err).ToNot(HaveOccurred())
   127  			done := make(chan struct{})
   128  			go func() {
   129  				defer GinkgoRecover()
   130  				runReceivingPeer(conn)
   131  				close(done)
   132  			}()
   133  			runSendingPeer(conn)
   134  			<-done
   135  			close(done1)
   136  		}()
   137  
   138  		client, err := quic.DialAddr(
   139  			context.Background(),
   140  			serverAddr,
   141  			getTLSClientConfig(),
   142  			getQuicConfig(nil),
   143  		)
   144  		Expect(err).ToNot(HaveOccurred())
   145  		done2 := make(chan struct{})
   146  		go func() {
   147  			defer GinkgoRecover()
   148  			runSendingPeer(client)
   149  			close(done2)
   150  		}()
   151  		runReceivingPeer(client)
   152  		<-done1
   153  		<-done2
   154  		client.CloseWithError(0, "")
   155  	})
   156  })