Switch to stateless compression with klauspost/compress

coder · Feb 13, 2020 · 7dbe93c · 7dbe93c
1 parent 503b469
commit 7dbe93c
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 35 deletions.
diff --git a/compress_notjs.go b/compress_notjs.go
@@ -98,22 +98,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/go.mod b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/golang/protobuf v1.3.3
 	github.com/google/go-cmp v0.4.0
 	github.com/gorilla/websocket v1.4.1
+	github.com/klauspost/compress v1.10.0
 	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543
 )
diff --git a/go.sum b/go.sum
@@ -10,6 +10,8 @@ github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/klauspost/compress v1.10.0 h1:92XGj1AcYzA6UrVdd4qIIBrT8OroryvRvdmg/IfmC7Y=
+github.com/klauspost/compress v1.10.0/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=

diff --git a/write.go b/write.go
@@ -4,14 +4,14 @@ package websocket
 
 import (
 	"bufio"
-	"compress/flate"
 	"context"
 	"crypto/rand"
 	"encoding/binary"
 	"io"
 	"sync"
 	"time"
 
+	"github.com/klauspost/compress/flate"
 	"golang.org/x/xerrors"
 
 	"nhooyr.io/websocket/internal/errd"
@@ -59,8 +59,8 @@ type msgWriter struct {
 	closed bool
 	flate  bool
 
-	trimWriter  *trimLastFourBytesWriter
-	flateWriter *flate.Writer
+	trimWriter *trimLastFourBytesWriter
+	dict       slidingWindow
 }
 
 func newMsgWriter(c *Conn) *msgWriter {
@@ -78,10 +78,7 @@ func (mw *msgWriter) ensureFlate() {
 		}
 	}
 
-	if mw.flateWriter == nil {
-		mw.flateWriter = getFlateWriter(mw.trimWriter)
-	}
-
+	mw.dict.init(8192)
 	mw.flate = true
 }
 
@@ -138,13 +135,6 @@ func (mw *msgWriter) reset(ctx context.Context, typ MessageType) error {
 	return nil
 }
 
-func (mw *msgWriter) returnFlateWriter() {
-	if mw.flateWriter != nil {
-		putFlateWriter(mw.flateWriter)
-		mw.flateWriter = nil
-	}
-}
-
 // Write writes the given bytes to the WebSocket connection.
 func (mw *msgWriter) Write(p []byte) (_ int, err error) {
 	defer errd.Wrap(&err, "failed to write")
@@ -165,7 +155,9 @@ func (mw *msgWriter) Write(p []byte) (_ int, err error) {
 	}
 
 	if mw.flate {
-		return mw.flateWriter.Write(p)
+		err = flate.StatelessDeflate(mw.trimWriter, p, false, mw.dict.buf)
+		mw.dict.write(p)
+		return len(p), err
 	}
 
 	return mw.write(p)
@@ -192,9 +184,9 @@ func (mw *msgWriter) Close() (err error) {
 	}
 
 	if mw.flate {
-		err = mw.flateWriter.Flush()
+		err = flate.StatelessDeflate(mw.trimWriter, nil, true, nil)
 		if err != nil {
-			return xerrors.Errorf("failed to flush flate writer: %w", err)
+			return xerrors.Errorf("failed to flush flate: %w", err)
 		}
 	}
 
@@ -208,15 +200,16 @@ func (mw *msgWriter) Close() (err error) {
 	}
 
 	if mw.flate && !mw.flateContextTakeover() {
-		mw.returnFlateWriter()
+		mw.dict.close()
 	}
+
 	mw.mu.Unlock()
 	return nil
 }
 
 func (mw *msgWriter) close() {
 	mw.writeMu.Lock()
-	mw.returnFlateWriter()
+	mw.dict.close()
 }
 
 func (c *Conn) writeControl(ctx context.Context, opcode opcode, p []byte) error {