Report how efficient compression is in BenchmarkConn

nhooyr · nhooyr · commit d57b25304679 · 2020-02-15T21:16:58.000-05:00
diff --git a/ci/test.mk b/ci/test.mk
@@ -1,4 +1,4 @@
-test: gotest ci/out/coverage.html
+test: ci/out/coverage.html
 ifdef CI
 test: coveralls
 endif
diff --git a/compress_notjs.go b/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
 }
diff --git a/conn_test.go b/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)
 		})
diff --git a/export_test.go b/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
+}
diff --git a/internal/test/assert/assert.go b/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)
 	}
 }
diff --git a/write.go b/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.

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-test: gotest ci/out/coverage.html`
	`1`	`+test: ci/out/coverage.html`
`2`	`2`	`ifdef CI`
`3`	`3`	`test: coveralls`
`4`	`4`	`endif`
Original file line number	Diff line number	Diff line change
`@@ -39,8 +39,8 @@ func Error(t testing.TB, err error) {`
`39`	`39`	`func Contains(t testing.TB, v interface{}, sub string) {`
`40`	`40`	`t.Helper()`
`41`	`41`
`42`		`- vstr := fmt.Sprint(v)`
`43`		`- if !strings.Contains(vstr, sub) {`
`44`		`- t.Fatalf("expected %q to contain %q", vstr, sub)`
	`42`	`+ s := fmt.Sprint(v)`
	`43`	`+ if !strings.Contains(s, sub) {`
	`44`	`+ t.Fatalf("expected %q to contain %q", s, sub)`
`45`	`45`	`}`
`46`	`46`	`}`