Report how efficient compression is in BenchmarkConn

Author: nhooyr

ci/test.mk

@@ -1,4 +1,4 @@
-test: gotest ci/out/coverage.html
+test: ci/out/coverage.html
 ifdef CI
 test: coveralls
 endif

compress_notjs.go

@@ -108,22 +108,6 @@ func putFlateReader(fr io.Reader) {
 	flateReaderPool.Put(fr)
 }
 
-var flateWriterPool sync.Pool
-
-func getFlateWriter(w io.Writer) *flate.Writer {
-	fw, ok := flateWriterPool.Get().(*flate.Writer)
-	if !ok {
-		fw, _ = flate.NewWriter(w, flate.BestSpeed)
-		return fw
-	}
-	fw.Reset(w)
-	return fw
-}
-
-func putFlateWriter(w *flate.Writer) {
-	flateWriterPool.Put(w)
-}
-
 type slidingWindow struct {
 	buf []byte
 }

conn_test.go

@@ -402,6 +402,9 @@ func BenchmarkConn(b *testing.B) {
 
 			bb.goEchoLoop(c2)
 
+			bytesWritten := c1.RecordBytesWritten()
+			bytesRead := c1.RecordBytesRead()
+
 			msg := []byte(strings.Repeat("1234", 128))
 			readBuf := make([]byte, len(msg))
 			writes := make(chan struct{})
@@ -451,6 +454,9 @@ func BenchmarkConn(b *testing.B) {
 			}
 			b.StopTimer()
 
+			b.ReportMetric(float64(*bytesWritten/b.N), "written/op")
+			b.ReportMetric(float64(*bytesRead/b.N), "read/op")
+
 			err := c1.Close(websocket.StatusNormalClosure, "")
 			assert.Success(b, err)
 		})

export_test.go

@@ -0,0 +1,22 @@
+// +build !js
+
+package websocket
+
+func (c *Conn) RecordBytesWritten() *int {
+	var bytesWritten int
+	c.bw.Reset(writerFunc(func(p []byte) (int, error) {
+		bytesWritten += len(p)
+		return c.rwc.Write(p)
+	}))
+	return &bytesWritten
+}
+
+func (c *Conn) RecordBytesRead() *int {
+	var bytesRead int
+	c.br.Reset(readerFunc(func(p []byte) (int, error) {
+		n, err := c.rwc.Read(p)
+		bytesRead += n
+		return n, err
+	}))
+	return &bytesRead
+}

internal/test/assert/assert.go

@@ -39,8 +39,8 @@ func Error(t testing.TB, err error) {
 func Contains(t testing.TB, v interface{}, sub string) {
 	t.Helper()
 
-	vstr := fmt.Sprint(v)
-	if !strings.Contains(vstr, sub) {
-		t.Fatalf("expected %q to contain %q", vstr, sub)
+	s := fmt.Sprint(v)
+	if !strings.Contains(s, sub) {
+		t.Fatalf("expected %q to contain %q", s, sub)
 	}
 }

write.go

@@ -12,7 +12,6 @@ import (
 	"time"
 
 	"github.com/klauspost/compress/flate"
-	kflate "github.com/klauspost/compress/flate"
 	"golang.org/x/xerrors"
 
 	"nhooyr.io/websocket/internal/errd"
@@ -77,9 +76,8 @@ type msgWriterState struct {
 	opcode opcode
 	flate  bool
 
-	trimWriter  *trimLastFourBytesWriter
-	flateWriter *flate.Writer
-	dict        slidingWindow
+	trimWriter *trimLastFourBytesWriter
+	dict       slidingWindow
 }
 
 func newMsgWriterState(c *Conn) *msgWriterState {
@@ -90,23 +88,14 @@ func newMsgWriterState(c *Conn) *msgWriterState {
 	return mw
 }
 
-const stateless = true
-
 func (mw *msgWriterState) ensureFlate() {
 	if mw.trimWriter == nil {
 		mw.trimWriter = &trimLastFourBytesWriter{
 			w: writerFunc(mw.write),
 		}
 	}
 
-	if stateless {
-		mw.dict.init(8192)
-	} else {
-		if mw.flateWriter == nil {
-			mw.flateWriter = getFlateWriter(mw.trimWriter)
-		}
-	}
-
+	mw.dict.init(8192)
 	mw.flate = true
 }
 
@@ -163,13 +152,6 @@ func (mw *msgWriterState) reset(ctx context.Context, typ MessageType) error {
 	return nil
 }
 
-func (mw *msgWriterState) putFlateWriter() {
-	if mw.flateWriter != nil {
-		putFlateWriter(mw.flateWriter)
-		mw.flateWriter = nil
-	}
-}
-
 // Write writes the given bytes to the WebSocket connection.
 func (mw *msgWriterState) Write(p []byte) (_ int, err error) {
 	defer errd.Wrap(&err, "failed to write")
@@ -186,15 +168,12 @@ func (mw *msgWriterState) Write(p []byte) (_ int, err error) {
 	}
 
 	if mw.flate {
-		if stateless {
-			err = kflate.StatelessDeflate(mw.trimWriter, p, false, mw.dict.buf)
-			if err != nil {
-				return 0, err
-			}
-			mw.dict.write(p)
-			return len(p), nil
+		err = flate.StatelessDeflate(mw.trimWriter, p, false, mw.dict.buf)
+		if err != nil {
+			return 0, err
 		}
-		return mw.flateWriter.Write(p)
+		mw.dict.write(p)
+		return len(p), nil
 	}
 
 	return mw.write(p)
@@ -216,29 +195,20 @@ func (mw *msgWriterState) Close() (err error) {
 	mw.writeMu.Lock()
 	defer mw.writeMu.Unlock()
 
-	if mw.flate && !stateless {
-		err = mw.flateWriter.Flush()
-		if err != nil {
-			return xerrors.Errorf("failed to flush flate: %w", err)
-		}
-	}
-
 	_, err = mw.c.writeFrame(mw.ctx, true, mw.flate, mw.opcode, nil)
 	if err != nil {
 		return xerrors.Errorf("failed to write fin frame: %w", err)
 	}
 
 	if mw.flate && !mw.flateContextTakeover() {
 		mw.dict.close()
-		mw.putFlateWriter()
 	}
 	mw.mu.Unlock()
 	return nil
 }
 
 func (mw *msgWriterState) close() {
 	mw.writeMu.Lock()
-	mw.putFlateWriter()
 	mw.dict.close()
 }
 
@@ -311,14 +281,13 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 	return n, nil
 }
 
-func (c *Conn) writeFramePayload(p []byte) (_ int, err error) {
+func (c *Conn) writeFramePayload(p []byte) (n int, err error) {
 	defer errd.Wrap(&err, "failed to write frame payload")
 
 	if !c.writeHeader.masked {
 		return c.bw.Write(p)
 	}
 
-	var n int
 	maskKey := c.writeHeader.maskKey
 	for len(p) > 0 {
 		// If the buffer is full, we need to flush.