Simplify sliding window API

nhooyr · nhooyr · commit 0233644fe805 · 2020-02-13T01:57:19.000-05:00
diff --git a/compress_notjs.go b/compress_notjs.go
@@ -120,41 +120,47 @@ type slidingWindow struct {
 
 var swPool = map[int]*sync.Pool{}
 
-func newSlidingWindow(n int) *slidingWindow {
+func (sw *slidingWindow) init(n int) {
+	if sw.buf != nil {
+		return
+	}
+
 	p, ok := swPool[n]
 	if !ok {
 		p = &sync.Pool{}
 		swPool[n] = p
 	}
-	sw, ok := p.Get().(*slidingWindow)
+	buf, ok := p.Get().([]byte)
 	if ok {
-		return sw
-	}
-	return &slidingWindow{
-		buf: make([]byte, 0, n),
+		sw.buf = buf[:0]
+	} else {
+		sw.buf = make([]byte, 0, n)
 	}
 }
 
-func returnSlidingWindow(sw *slidingWindow) {
-	sw.buf = sw.buf[:0]
-	swPool[cap(sw.buf)].Put(sw)
+func (sw *slidingWindow) close() {
+	if sw.buf == nil {
+		return
+	}
+
+	swPool[cap(sw.buf)].Put(sw.buf)
 }
 
-func (w *slidingWindow) write(p []byte) {
-	if len(p) >= cap(w.buf) {
-		w.buf = w.buf[:cap(w.buf)]
-		p = p[len(p)-cap(w.buf):]
-		copy(w.buf, p)
+func (sw *slidingWindow) write(p []byte) {
+	if len(p) >= cap(sw.buf) {
+		sw.buf = sw.buf[:cap(sw.buf)]
+		p = p[len(p)-cap(sw.buf):]
+		copy(sw.buf, p)
 		return
 	}
 
-	left := cap(w.buf) - len(w.buf)
+	left := cap(sw.buf) - len(sw.buf)
 	if left < len(p) {
 		// We need to shift spaceNeeded bytes from the end to make room for p at the end.
 		spaceNeeded := len(p) - left
-		copy(w.buf, w.buf[spaceNeeded:])
-		w.buf = w.buf[:len(w.buf)-spaceNeeded]
+		copy(sw.buf, sw.buf[spaceNeeded:])
+		sw.buf = sw.buf[:len(sw.buf)-spaceNeeded]
 	}
 
-	w.buf = append(w.buf, p...)
+	sw.buf = append(sw.buf, p...)
 }
diff --git a/compress_test.go b/compress_test.go
@@ -21,12 +21,13 @@ func Test_slidingWindow(t *testing.T) {
 
 			input := xrand.String(maxWindow)
 			windowLength := xrand.Int(maxWindow)
-			r := newSlidingWindow(windowLength)
-			r.write([]byte(input))
+			var sw slidingWindow
+			sw.init(windowLength)
+			sw.write([]byte(input))
 
-			assert.Equal(t, "window length", windowLength, cap(r.buf))
-			if !strings.HasSuffix(input, string(r.buf)) {
-				t.Fatalf("r.buf is not a suffix of input: %q and %q", input, r.buf)
+			assert.Equal(t, "window length", windowLength, cap(sw.buf))
+			if !strings.HasSuffix(input, string(sw.buf)) {
+				t.Fatalf("r.buf is not a suffix of input: %q and %q", input, sw.buf)
 			}
 		})
 	}
diff --git a/conn_test.go b/conn_test.go
@@ -351,13 +351,12 @@ func (tt *connTest) goDiscardLoop(c *websocket.Conn) {
 	ctx, cancel := context.WithCancel(tt.ctx)
 
 	discardLoopErr := xsync.Go(func() error {
+		defer c.Close(websocket.StatusInternalError, "")
+
 		for {
 			_, _, err := c.Read(ctx)
-			if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
-				return nil
-			}
 			if err != nil {
-				return err
+				return assertCloseStatus(websocket.StatusNormalClosure, err)
 			}
 		}
 	})
diff --git a/read.go b/read.go
@@ -87,15 +87,11 @@ func newMsgReader(c *Conn) *msgReader {
 }
 
 func (mr *msgReader) resetFlate() {
-	if mr.flateContextTakeover() && mr.dict == nil {
-		mr.dict = newSlidingWindow(32768)
-	}
-
 	if mr.flateContextTakeover() {
-		mr.flateReader = getFlateReader(readerFunc(mr.read), mr.dict.buf)
-	} else {
-		mr.flateReader = getFlateReader(readerFunc(mr.read), nil)
+		mr.dict.init(32768)
 	}
+
+	mr.flateReader = getFlateReader(readerFunc(mr.read), mr.dict.buf)
 	mr.limitReader.r = mr.flateReader
 	mr.flateTail.Reset(deflateMessageTail)
 }
@@ -111,9 +107,7 @@ func (mr *msgReader) close() {
 	mr.c.readMu.Lock(context.Background())
 	mr.returnFlateReader()
 
-	if mr.dict != nil {
-		returnSlidingWindow(mr.dict)
-	}
+	mr.dict.close()
 }
 
 func (mr *msgReader) flateContextTakeover() bool {
@@ -325,7 +319,7 @@ type msgReader struct {
 	flateReader io.Reader
 	flateTail   strings.Reader
 	limitReader *limitReader
-	dict        *slidingWindow
+	dict        slidingWindow
 
 	fin           bool
 	payloadLength int64
