quic: add Stream.ReadByte, Stream.WriteByte

Currently unoptimized and slow.
Adding along with a benchmark to compare to the fast-path followup.

For golang/go#58547

Change-Id: If02b65e6e7cfc770d3f949e5fb9fbb9d8a765a90
Reviewed-on: https://go-review.googlesource.com/c/net/+/564477
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

internal/quic/bench_test.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"sync"
 	"testing"
 )
 
@@ -72,6 +73,76 @@ func throughput(b *testing.B, totalBytes int64) {
 	}
 }
 
+func BenchmarkReadByte(b *testing.B) {
+	cli, srv := newLocalConnPair(b, &Config{}, &Config{})
+
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		buf := make([]byte, 1<<20)
+		sconn, err := srv.AcceptStream(context.Background())
+		if err != nil {
+			panic(fmt.Errorf("AcceptStream: %v", err))
+		}
+		for {
+			if _, err := sconn.Write(buf); err != nil {
+				break
+			}
+			sconn.Flush()
+		}
+	}()
+
+	b.SetBytes(1)
+	cconn, err := cli.NewStream(context.Background())
+	if err != nil {
+		b.Fatalf("NewStream: %v", err)
+	}
+	cconn.Flush()
+	for i := 0; i < b.N; i++ {
+		_, err := cconn.ReadByte()
+		if err != nil {
+			b.Fatalf("ReadByte: %v", err)
+		}
+	}
+	cconn.Close()
+}
+
+func BenchmarkWriteByte(b *testing.B) {
+	cli, srv := newLocalConnPair(b, &Config{}, &Config{})
+
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		sconn, err := srv.AcceptStream(context.Background())
+		if err != nil {
+			panic(fmt.Errorf("AcceptStream: %v", err))
+		}
+		n, err := io.Copy(io.Discard, sconn)
+		if n != int64(b.N) || err != nil {
+			b.Errorf("server io.Copy() = %v, %v; want %v, nil", n, err, b.N)
+		}
+	}()
+
+	b.SetBytes(1)
+	cconn, err := cli.NewStream(context.Background())
+	if err != nil {
+		b.Fatalf("NewStream: %v", err)
+	}
+	cconn.Flush()
+	for i := 0; i < b.N; i++ {
+		if err := cconn.WriteByte(0); err != nil {
+			b.Fatalf("WriteByte: %v", err)
+		}
+	}
+	cconn.Close()
+}
+
 func BenchmarkStreamCreation(b *testing.B) {
 	cli, srv := newLocalConnPair(b, &Config{}, &Config{})
 

internal/quic/stream.go

@@ -245,6 +245,13 @@ func (s *Stream) Read(b []byte) (n int, err error) {
 	return len(b), nil
 }
 
+// ReadByte reads and returns a single byte from the stream.
+func (s *Stream) ReadByte() (byte, error) {
+	var b [1]byte
+	_, err := s.Read(b[:])
+	return b[0], err
+}
+
 // shouldUpdateFlowControl determines whether to send a flow control window update.
 //
 // We want to balance keeping the peer well-supplied with flow control with not sending
@@ -326,6 +333,13 @@ func (s *Stream) Write(b []byte) (n int, err error) {
 	return n, nil
 }
 
+// WriteBytes writes a single byte to the stream.
+func (s *Stream) WriteByte(c byte) error {
+	b := [1]byte{c}
+	_, err := s.Write(b[:])
+	return err
+}
+
 // Flush flushes data written to the stream.
 // It does not wait for the peer to acknowledge receipt of the data.
 // Use Close to wait for the peer's acknowledgement.