mod: test and write compression

client_server_test.go

@@ -142,6 +142,45 @@ func sendRecv(t *testing.T, ws *Conn) {
 	}
 }
 
+func multipleSendRecv(t *testing.T, ws *Conn) {
+	message := "Hello World!"
+	if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("SetWriteDeadline: %v", err)
+	}
+	if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil {
+		t.Fatalf("WriteMessage: %v", err)
+	}
+	if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	_, p, err := ws.ReadMessage()
+	if err != nil {
+		t.Fatalf("ReadMessage: %v", err)
+	}
+	if string(p) != message {
+		t.Fatalf("message=%s, want %s", p, message)
+	}
+
+	message_2 := "Can you read message?"
+	if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("SetWriteDeadline: %v", err)
+	}
+	if err := ws.WriteMessage(TextMessage, []byte(message_2)); err != nil {
+		t.Fatalf("_WriteMessage: %v", err)
+	}
+	if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("_SetReadDeadline: %v", err)
+	}
+
+	_, p, err = ws.ReadMessage()
+	if err != nil {
+		t.Fatalf("_ReadMessage: %v", err) // _ReadMessage: websocket: close 1006 (abnormal closure): unexpected EOF
+	}
+	if string(p) != message {
+		t.Fatalf("_message=%s, want %s", p, message_2)
+	}
+}
+
 func TestProxyDial(t *testing.T) {
 
 	s := newServer(t)
@@ -522,6 +561,23 @@ func TestDialCompression(t *testing.T) {
 	sendRecv(t, ws)
 }
 
+func TestDialCompressionOfContextTakeover(t *testing.T) {
+	s := newServer(t)
+	defer s.Close()
+
+	dialer := cstDialer
+	dialer.EnableCompression = true
+	dialer.EnableContextTakeover = true
+	ws, _, err := dialer.Dial(s.URL, nil)
+	if err != nil {
+		t.Fatalf("Dial: %v", err)
+	}
+	defer ws.Close()
+
+	// Todo multiple send and receive.
+	multipleSendRecv(t, ws)
+}
+
 func TestSocksProxyDial(t *testing.T) {
 	s := newServer(t)
 	defer s.Close()

compression.go

@@ -20,6 +20,7 @@ const (
 
 var (
 	flateWriterPools     [maxCompressionLevel - minCompressionLevel + 1]sync.Pool
+	flateWriterDictPools [maxCompressionLevel - minCompressionLevel + 1]sync.Pool
 	flateReaderPool      = sync.Pool{New: func() interface{} {
 		return flate.NewReader(nil)
 	}}
@@ -66,12 +67,16 @@ func compressNoContextTakeover(w io.WriteCloser, level int, dict []byte) io.Writ
 }
 
 func compressContextTakeover(w io.WriteCloser, level int, dict []byte) io.WriteCloser {
-	p := &flateWriterPools[level-minCompressionLevel]
+	p := &flateWriterDictPools[level-minCompressionLevel]
 	tw := &truncWriter{w: w}
 
-	// WriterDict's Reset just restores the dictionary.
-	// Initialization is done with New. (If possible get struct from sync.Pool)
-	fw, _ := flate.NewWriterDict(tw, level, dict)
+	fw, _ := p.Get().(*flate.Writer)
+	if fw == nil {
+		// use WriterDict
+		fw, _ = flate.NewWriterDict(tw, level, dict)
+	} else {
+		fw.Reset(tw)
+	}
 
 	return &flateWriteWrapper{fw: fw, tw: tw, p: p}
 }

compression_test.go

@@ -65,6 +65,20 @@ func BenchmarkWriteWithCompression(b *testing.B) {
 	b.ReportAllocs()
 }
 
+func BenchmarkWriteWithCompressionOfContextTakeover(b *testing.B) {
+	w := ioutil.Discard
+	c := newConn(fakeNetConn{Reader: nil, Writer: w}, false, 1024, 1024)
+	messages := textMessages(100)
+	c.enableWriteCompression = true
+	c.contextTakeover = true
+	c.newCompressionWriter = compressContextTakeover
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.WriteMessage(TextMessage, messages[i%len(messages)])
+	}
+	b.ReportAllocs()
+}
+
 func TestValidCompressionLevel(t *testing.T) {
 	c := newConn(fakeNetConn{}, false, 1024, 1024)
 	for _, level := range []int{minCompressionLevel - 1, maxCompressionLevel + 1} {

conn_test.go

@@ -494,3 +494,14 @@ func TestBufioReuse(t *testing.T) {
 	}
 
 }
+
+func BenchmarkAddDict(b *testing.B) {
+	w := ioutil.Discard
+	c := newConn(fakeNetConn{Reader: nil, Writer: w}, false, 1024, 1024)
+	messages := textMessages(100)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.AddRxDict(messages[i%len(messages)])
+	}
+	b.ReportAllocs()
+}