internal/quic/conn_streams.go - net - Git at Google

 // Copyright 2023 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 //go:build go1.21

 package quic

 import (
 	"context"
 	"sync"
 	"sync/atomic"
 	"time"
 )

 type streamsState struct {
 	queue queue[*Stream] // new, peer-created streams

 	streamsMu sync.Mutex
 	streams   map[streamID]*Stream

 	// Limits on the number of streams, indexed by streamType.
 	localLimit  [streamTypeCount]localStreamLimits
 	remoteLimit [streamTypeCount]remoteStreamLimits

 	// Peer configuration provided in transport parameters.
 	peerInitialMaxStreamDataRemote    [streamTypeCount]int64 // streams opened by us
 	peerInitialMaxStreamDataBidiLocal int64                  // streams opened by them

 	// Streams with frames to send are stored in a circular linked list.
 	// sendHead is the next stream to write, or nil if there are no streams
 	// with data to send. sendTail is the last stream to write.
 	needSend atomic.Bool
 	sendMu   sync.Mutex
 	sendHead *Stream
 	sendTail *Stream
 }

 func (c *Conn) streamsInit() {
 	c.streams.streams = make(map[streamID]*Stream)
 	c.streams.queue = newQueue[*Stream]()
 	c.streams.localLimit[bidiStream].init()
 	c.streams.localLimit[uniStream].init()
 	c.streams.remoteLimit[bidiStream].init(c.config.maxBidiRemoteStreams())
 	c.streams.remoteLimit[uniStream].init(c.config.maxUniRemoteStreams())
 }

 // AcceptStream waits for and returns the next stream created by the peer.
 func (c *Conn) AcceptStream(ctx context.Context) (*Stream, error) {
 	return c.streams.queue.getWithHooks(ctx, c.testHooks)
 }

 // NewStream creates a stream.
 //
 // If the peer's maximum stream limit for the connection has been reached,
 // NewStream blocks until the limit is increased or the context expires.
 func (c *Conn) NewStream(ctx context.Context) (*Stream, error) {
 	return c.newLocalStream(ctx, bidiStream)
 }

 // NewSendOnlyStream creates a unidirectional, send-only stream.
 //
 // If the peer's maximum stream limit for the connection has been reached,
 // NewSendOnlyStream blocks until the limit is increased or the context expires.
 func (c *Conn) NewSendOnlyStream(ctx context.Context) (*Stream, error) {
 	return c.newLocalStream(ctx, uniStream)
 }

 func (c *Conn) newLocalStream(ctx context.Context, styp streamType) (*Stream, error) {
 	c.streams.streamsMu.Lock()
 	defer c.streams.streamsMu.Unlock()

 	num, err := c.streams.localLimit[styp].open(ctx, c)
 	if err != nil {
 		return nil, err
 	}

 	s := newStream(c, newStreamID(c.side, styp, num))
 	s.outmaxbuf = c.config.streamWriteBufferSize()
 	s.outwin = c.streams.peerInitialMaxStreamDataRemote[styp]
 	if styp == bidiStream {
 		s.inmaxbuf = c.config.streamReadBufferSize()
 		s.inwin = c.config.streamReadBufferSize()
 	}
 	s.inUnlock()
 	s.outUnlock()

 	c.streams.streams[s.id] = s
 	return s, nil
 }

 // streamFrameType identifies which direction of a stream,
 // from the local perspective, a frame is associated with.
 //
 // For example, STREAM is a recvStream frame,
 // because it carries data from the peer to us.
 type streamFrameType uint8

 const (
 	sendStream = streamFrameType(iota) // for example, MAX_DATA
 	recvStream                         // for example, STREAM_DATA_BLOCKED
 )

 // streamForID returns the stream with the given id.
 // If the stream does not exist, it returns nil.
 func (c *Conn) streamForID(id streamID) *Stream {
 	c.streams.streamsMu.Lock()
 	defer c.streams.streamsMu.Unlock()
 	return c.streams.streams[id]
 }

 // streamForFrame returns the stream with the given id.
 // If the stream does not exist, it may be created.
 //
 // streamForFrame aborts the connection if the stream id, state, and frame type don't align.
 // For example, it aborts the connection with a STREAM_STATE error if a MAX_DATA frame
 // is received for a receive-only stream, or if the peer attempts to create a stream that
 // should be originated locally.
 //
 // streamForFrame returns nil if the stream no longer exists or if an error occurred.
 func (c *Conn) streamForFrame(now time.Time, id streamID, ftype streamFrameType) *Stream {
 	if id.streamType() == uniStream {
 		if (id.initiator() == c.side) != (ftype == sendStream) {
 			// Received an invalid frame for unidirectional stream.
 			// For example, a RESET_STREAM frame for a send-only stream.
 			c.abort(now, localTransportError(errStreamState))
 			return nil
 		}
 	}

 	c.streams.streamsMu.Lock()
 	defer c.streams.streamsMu.Unlock()
 	s, isOpen := c.streams.streams[id]
 	if s != nil {
 		return s
 	}

 	num := id.num()
 	styp := id.streamType()
 	if id.initiator() == c.side {
 		if num < c.streams.localLimit[styp].opened {
 			// This stream was created by us, and has been closed.
 			return nil
 		}
 		// Received a frame for a stream that should be originated by us,
 		// but which we never created.
 		c.abort(now, localTransportError(errStreamState))
 		return nil
 	} else {
 		// if isOpen, this is a stream that was implicitly opened by a
 		// previous frame for a larger-numbered stream, but we haven't
 		// actually created it yet.
 		if !isOpen && num < c.streams.remoteLimit[styp].opened {
 			// This stream was created by the peer, and has been closed.
 			return nil
 		}
 	}

 	prevOpened := c.streams.remoteLimit[styp].opened
 	if err := c.streams.remoteLimit[styp].open(id); err != nil {
 		c.abort(now, err)
 		return nil
 	}

 	// Receiving a frame for a stream implicitly creates all streams
 	// with the same initiator and type and a lower number.
 	// Add a nil entry to the streams map for each implicitly created stream.
 	for n := newStreamID(id.initiator(), id.streamType(), prevOpened); n < id; n += 4 {
 		c.streams.streams[n] = nil
 	}

 	s = newStream(c, id)
 	s.inmaxbuf = c.config.streamReadBufferSize()
 	s.inwin = c.config.streamReadBufferSize()
 	if id.streamType() == bidiStream {
 		s.outmaxbuf = c.config.streamWriteBufferSize()
 		s.outwin = c.streams.peerInitialMaxStreamDataBidiLocal
 	}
 	s.inUnlock()
 	s.outUnlock()

 	c.streams.streams[id] = s
 	c.streams.queue.put(s)
 	return s
 }

 // queueStreamForSend marks a stream as containing frames that need sending.
 func (c *Conn) queueStreamForSend(s *Stream) {
 	c.streams.sendMu.Lock()
 	defer c.streams.sendMu.Unlock()
 	if s.next != nil {
 		// Already in the queue.
 		return
 	}
 	if c.streams.sendHead == nil {
 		// The queue was empty.
 		c.streams.sendHead = s
 		c.streams.sendTail = s
 		s.next = s
 	} else {
 		// Insert this stream at the end of the queue.
 		c.streams.sendTail.next = s
 		c.streams.sendTail = s
 		s.next = c.streams.sendHead
 	}
 	c.streams.needSend.Store(true)
 	c.wake()
 }

 // appendStreamFrames writes stream-related frames to the current packet.
 //
 // It returns true if no more frames need appending,
 // false if not everything fit in the current packet.
 func (c *Conn) appendStreamFrames(w *packetWriter, pnum packetNumber, pto bool) bool {
 	c.streams.remoteLimit[uniStream].appendFrame(w, uniStream, pnum, pto)
 	c.streams.remoteLimit[bidiStream].appendFrame(w, bidiStream, pnum, pto)
 	if pto {
 		return c.appendStreamFramesPTO(w, pnum)
 	}
 	if !c.streams.needSend.Load() {
 		return true
 	}
 	c.streams.sendMu.Lock()
 	defer c.streams.sendMu.Unlock()
 	for {
 		s := c.streams.sendHead
 		const pto = false

 		state := s.state.load()
 		if state&streamInSend != 0 {
 			s.ingate.lock()
 			ok := s.appendInFramesLocked(w, pnum, pto)
 			state = s.inUnlockNoQueue()
 			if !ok {
 				return false
 			}
 		}

 		if state&streamOutSend != 0 {
 			avail := w.avail()
 			s.outgate.lock()
 			ok := s.appendOutFramesLocked(w, pnum, pto)
 			state = s.outUnlockNoQueue()
 			if !ok {
 				// We've sent some data for this stream, but it still has more to send.
 				// If the stream got a reasonable chance to put data in a packet,
 				// advance sendHead to the next stream in line, to avoid starvation.
 				// We'll come back to this stream after going through the others.
 				//
 				// If the packet was already mostly out of space, leave sendHead alone
 				// and come back to this stream again on the next packet.
 				if avail > 512 {
 					c.streams.sendHead = s.next
 					c.streams.sendTail = s
 				}
 				return false
 			}
 		}

 		if state == streamInDone|streamOutDone {
 			// Stream is finished, remove it from the conn.
 			s.state.set(streamConnRemoved, streamConnRemoved)
 			delete(c.streams.streams, s.id)

 			// Record finalization of remote streams, to know when
 			// to extend the peer's stream limit.
 			if s.id.initiator() != c.side {
 				c.streams.remoteLimit[s.id.streamType()].close()
 			}
 		}

 		next := s.next
 		s.next = nil
 		if (next == s) != (s == c.streams.sendTail) {
 			panic("BUG: sendable stream list state is inconsistent")
 		}
 		if s == c.streams.sendTail {
 			// This was the last stream.
 			c.streams.sendHead = nil
 			c.streams.sendTail = nil
 			c.streams.needSend.Store(false)
 			return true
 		}
 		// We've sent all data for this stream, so remove it from the list.
 		c.streams.sendTail.next = next
 		c.streams.sendHead = next
 	}
 }

 // appendStreamFramesPTO writes stream-related frames to the current packet
 // for a PTO probe.
 //
 // It returns true if no more frames need appending,
 // false if not everything fit in the current packet.
 func (c *Conn) appendStreamFramesPTO(w *packetWriter, pnum packetNumber) bool {
 	c.streams.sendMu.Lock()
 	defer c.streams.sendMu.Unlock()
 	const pto = true
 	for _, s := range c.streams.streams {
 		const pto = true
 		s.ingate.lock()
 		inOK := s.appendInFramesLocked(w, pnum, pto)
 		s.inUnlockNoQueue()
 		if !inOK {
 			return false
 		}

 		s.outgate.lock()
 		outOK := s.appendOutFramesLocked(w, pnum, pto)
 		s.outUnlockNoQueue()
 		if !outOK {
 			return false
 		}
 	}
 	return true
 }
	// Copyright 2023 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	//go:build go1.21

	package quic

	import (
	"context"
	"sync"
	"sync/atomic"
	"time"
	)

	type streamsState struct {
	queue queue[*Stream] // new, peer-created streams

	streamsMu sync.Mutex
	streams map[streamID]*Stream

	// Limits on the number of streams, indexed by streamType.
	localLimit [streamTypeCount]localStreamLimits
	remoteLimit [streamTypeCount]remoteStreamLimits

	// Peer configuration provided in transport parameters.
	peerInitialMaxStreamDataRemote [streamTypeCount]int64 // streams opened by us
	peerInitialMaxStreamDataBidiLocal int64 // streams opened by them

	// Streams with frames to send are stored in a circular linked list.
	// sendHead is the next stream to write, or nil if there are no streams
	// with data to send. sendTail is the last stream to write.
	needSend atomic.Bool
	sendMu sync.Mutex
	sendHead *Stream
	sendTail *Stream
	}

	func (c *Conn) streamsInit() {
	c.streams.streams = make(map[streamID]*Stream)
	c.streams.queue = newQueue[*Stream]()
	c.streams.localLimit[bidiStream].init()
	c.streams.localLimit[uniStream].init()
	c.streams.remoteLimit[bidiStream].init(c.config.maxBidiRemoteStreams())
	c.streams.remoteLimit[uniStream].init(c.config.maxUniRemoteStreams())
	}

	// AcceptStream waits for and returns the next stream created by the peer.
	func (c Conn) AcceptStream(ctx context.Context) (Stream, error) {
	return c.streams.queue.getWithHooks(ctx, c.testHooks)
	}

	// NewStream creates a stream.
	//
	// If the peer's maximum stream limit for the connection has been reached,
	// NewStream blocks until the limit is increased or the context expires.
	func (c Conn) NewStream(ctx context.Context) (Stream, error) {
	return c.newLocalStream(ctx, bidiStream)
	}

	// NewSendOnlyStream creates a unidirectional, send-only stream.
	//
	// If the peer's maximum stream limit for the connection has been reached,
	// NewSendOnlyStream blocks until the limit is increased or the context expires.
	func (c Conn) NewSendOnlyStream(ctx context.Context) (Stream, error) {
	return c.newLocalStream(ctx, uniStream)
	}

	func (c Conn) newLocalStream(ctx context.Context, styp streamType) (Stream, error) {
	c.streams.streamsMu.Lock()
	defer c.streams.streamsMu.Unlock()

	num, err := c.streams.localLimit[styp].open(ctx, c)
	if err != nil {
	return nil, err
	}

	s := newStream(c, newStreamID(c.side, styp, num))
	s.outmaxbuf = c.config.streamWriteBufferSize()
	s.outwin = c.streams.peerInitialMaxStreamDataRemote[styp]
	if styp == bidiStream {
	s.inmaxbuf = c.config.streamReadBufferSize()
	s.inwin = c.config.streamReadBufferSize()
	}
	s.inUnlock()
	s.outUnlock()

	c.streams.streams[s.id] = s
	return s, nil
	}

	// streamFrameType identifies which direction of a stream,
	// from the local perspective, a frame is associated with.
	//
	// For example, STREAM is a recvStream frame,
	// because it carries data from the peer to us.
	type streamFrameType uint8

	const (
	sendStream = streamFrameType(iota) // for example, MAX_DATA
	recvStream // for example, STREAM_DATA_BLOCKED
	)

	// streamForID returns the stream with the given id.
	// If the stream does not exist, it returns nil.
	func (c Conn) streamForID(id streamID) Stream {
	c.streams.streamsMu.Lock()
	defer c.streams.streamsMu.Unlock()
	return c.streams.streams[id]
	}

	// streamForFrame returns the stream with the given id.
	// If the stream does not exist, it may be created.
	//
	// streamForFrame aborts the connection if the stream id, state, and frame type don't align.
	// For example, it aborts the connection with a STREAM_STATE error if a MAX_DATA frame
	// is received for a receive-only stream, or if the peer attempts to create a stream that
	// should be originated locally.
	//
	// streamForFrame returns nil if the stream no longer exists or if an error occurred.
	func (c Conn) streamForFrame(now time.Time, id streamID, ftype streamFrameType) Stream {
	if id.streamType() == uniStream {
	if (id.initiator() == c.side) != (ftype == sendStream) {
	// Received an invalid frame for unidirectional stream.
	// For example, a RESET_STREAM frame for a send-only stream.
	c.abort(now, localTransportError(errStreamState))
	return nil
	}
	}

	c.streams.streamsMu.Lock()
	defer c.streams.streamsMu.Unlock()
	s, isOpen := c.streams.streams[id]
	if s != nil {
	return s
	}

	num := id.num()
	styp := id.streamType()
	if id.initiator() == c.side {
	if num < c.streams.localLimit[styp].opened {
	// This stream was created by us, and has been closed.
	return nil
	}
	// Received a frame for a stream that should be originated by us,
	// but which we never created.
	c.abort(now, localTransportError(errStreamState))
	return nil
	} else {
	// if isOpen, this is a stream that was implicitly opened by a
	// previous frame for a larger-numbered stream, but we haven't
	// actually created it yet.
	if !isOpen && num < c.streams.remoteLimit[styp].opened {
	// This stream was created by the peer, and has been closed.
	return nil
	}
	}

	prevOpened := c.streams.remoteLimit[styp].opened
	if err := c.streams.remoteLimit[styp].open(id); err != nil {
	c.abort(now, err)
	return nil
	}

	// Receiving a frame for a stream implicitly creates all streams
	// with the same initiator and type and a lower number.
	// Add a nil entry to the streams map for each implicitly created stream.
	for n := newStreamID(id.initiator(), id.streamType(), prevOpened); n < id; n += 4 {
	c.streams.streams[n] = nil
	}

	s = newStream(c, id)
	s.inmaxbuf = c.config.streamReadBufferSize()
	s.inwin = c.config.streamReadBufferSize()
	if id.streamType() == bidiStream {
	s.outmaxbuf = c.config.streamWriteBufferSize()
	s.outwin = c.streams.peerInitialMaxStreamDataBidiLocal
	}
	s.inUnlock()
	s.outUnlock()

	c.streams.streams[id] = s
	c.streams.queue.put(s)
	return s
	}

	// queueStreamForSend marks a stream as containing frames that need sending.
	func (c Conn) queueStreamForSend(s Stream) {
	c.streams.sendMu.Lock()
	defer c.streams.sendMu.Unlock()
	if s.next != nil {
	// Already in the queue.
	return
	}
	if c.streams.sendHead == nil {
	// The queue was empty.
	c.streams.sendHead = s
	c.streams.sendTail = s
	s.next = s
	} else {
	// Insert this stream at the end of the queue.
	c.streams.sendTail.next = s
	c.streams.sendTail = s
	s.next = c.streams.sendHead
	}
	c.streams.needSend.Store(true)
	c.wake()
	}

	// appendStreamFrames writes stream-related frames to the current packet.
	//
	// It returns true if no more frames need appending,
	// false if not everything fit in the current packet.
	func (c Conn) appendStreamFrames(w packetWriter, pnum packetNumber, pto bool) bool {
	c.streams.remoteLimit[uniStream].appendFrame(w, uniStream, pnum, pto)
	c.streams.remoteLimit[bidiStream].appendFrame(w, bidiStream, pnum, pto)
	if pto {
	return c.appendStreamFramesPTO(w, pnum)
	}
	if !c.streams.needSend.Load() {
	return true
	}
	c.streams.sendMu.Lock()
	defer c.streams.sendMu.Unlock()
	for {
	s := c.streams.sendHead
	const pto = false

	state := s.state.load()
	if state&streamInSend != 0 {
	s.ingate.lock()
	ok := s.appendInFramesLocked(w, pnum, pto)
	state = s.inUnlockNoQueue()
	if !ok {
	return false
	}
	}

	if state&streamOutSend != 0 {
	avail := w.avail()
	s.outgate.lock()
	ok := s.appendOutFramesLocked(w, pnum, pto)
	state = s.outUnlockNoQueue()
	if !ok {
	// We've sent some data for this stream, but it still has more to send.
	// If the stream got a reasonable chance to put data in a packet,
	// advance sendHead to the next stream in line, to avoid starvation.
	// We'll come back to this stream after going through the others.
	//
	// If the packet was already mostly out of space, leave sendHead alone
	// and come back to this stream again on the next packet.
	if avail > 512 {
	c.streams.sendHead = s.next
	c.streams.sendTail = s
	}
	return false
	}
	}

	if state == streamInDone\|streamOutDone {
	// Stream is finished, remove it from the conn.
	s.state.set(streamConnRemoved, streamConnRemoved)
	delete(c.streams.streams, s.id)

	// Record finalization of remote streams, to know when
	// to extend the peer's stream limit.
	if s.id.initiator() != c.side {
	c.streams.remoteLimit[s.id.streamType()].close()
	}
	}

	next := s.next
	s.next = nil
	if (next == s) != (s == c.streams.sendTail) {
	panic("BUG: sendable stream list state is inconsistent")
	}
	if s == c.streams.sendTail {
	// This was the last stream.
	c.streams.sendHead = nil
	c.streams.sendTail = nil
	c.streams.needSend.Store(false)
	return true
	}
	// We've sent all data for this stream, so remove it from the list.
	c.streams.sendTail.next = next
	c.streams.sendHead = next
	}
	}

	// appendStreamFramesPTO writes stream-related frames to the current packet
	// for a PTO probe.
	//
	// It returns true if no more frames need appending,
	// false if not everything fit in the current packet.
	func (c Conn) appendStreamFramesPTO(w packetWriter, pnum packetNumber) bool {
	c.streams.sendMu.Lock()
	defer c.streams.sendMu.Unlock()
	const pto = true
	for _, s := range c.streams.streams {
	const pto = true
	s.ingate.lock()
	inOK := s.appendInFramesLocked(w, pnum, pto)
	s.inUnlockNoQueue()
	if !inOK {
	return false
	}

	s.outgate.lock()
	outOK := s.appendOutFramesLocked(w, pnum, pto)
	s.outUnlockNoQueue()
	if !outOK {
	return false
	}
	}
	return true
	}