From a366cdf6168ee6d9c12d2babca11f8e7b2a4cf0e Mon Sep 17 00:00:00 2001
From: misu <misu@aptpod.co.jp>
Date: Wed, 24 Jan 2018 17:53:38 +0900
Subject: [PATCH] impl: compressContextTakeover

---
 compression.go | 15 +++++++--------
 conn.go        | 14 +++++++++++---
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/compression.go b/compression.go
index 64cbb44..4c5f0aa 100644
--- a/compression.go
+++ b/compression.go
@@ -53,7 +53,7 @@ func isValidCompressionLevel(level int) bool {
 	return minCompressionLevel <= level && level <= maxCompressionLevel
 }
 
-func compressNoContextTakeover(w io.WriteCloser, level int) io.WriteCloser {
+func compressNoContextTakeover(w io.WriteCloser, level int, dict []byte) io.WriteCloser {
 	p := &flateWriterPools[level-minCompressionLevel]
 	tw := &truncWriter{w: w}
 	fw, _ := p.Get().(*flate.Writer)
@@ -65,15 +65,14 @@ func compressNoContextTakeover(w io.WriteCloser, level int) io.WriteCloser {
 	return &flateWriteWrapper{fw: fw, tw: tw, p: p}
 }
 
-func compressContextTakeover(w io.WriteCloser, level int) io.WriteCloser {
+func compressContextTakeover(w io.WriteCloser, level int, dict []byte) io.WriteCloser {
 	p := &flateWriterPools[level-minCompressionLevel]
 	tw := &truncWriter{w: w}
-	fw, _ := p.Get().(*flate.Writer)
-	if fw == nil {
-		fw, _ = flate.NewWriter(tw, level)
-	} else {
-		fw.Reset(tw)
-	}
+
+	// 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)
+
 	return &flateWriteWrapper{fw: fw, tw: tw, p: p}
 }
 
diff --git a/conn.go b/conn.go
index cf4a9af..c740e01 100644
--- a/conn.go
+++ b/conn.go
@@ -243,7 +243,7 @@ type Conn struct {
 
 	enableWriteCompression bool
 	compressionLevel       int
-	newCompressionWriter   func(io.WriteCloser, int) io.WriteCloser
+	newCompressionWriter   func(io.WriteCloser, int, []byte) io.WriteCloser
 
 	// Read fields
 	reader        io.ReadCloser // the current reader returned to the application
@@ -505,9 +505,14 @@ func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
 	}
 	c.writer = mw
 	if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
-		w := c.newCompressionWriter(c.writer, c.compressionLevel)
 		mw.compress = true
-		c.writer = w
+		switch {
+		case c.contextTakeover:
+			c.writer = c.newCompressionWriter(c.writer, c.compressionLevel, c.dict)
+		// no-context-takeover
+		default:
+			c.writer = c.newCompressionWriter(c.writer, c.compressionLevel, nil)
+		}
 	}
 	return c.writer, nil
 }
@@ -758,6 +763,9 @@ func (c *Conn) WriteMessage(messageType int, data []byte) error {
 	if _, err = w.Write(data); err != nil {
 		return err
 	}
+	if c.contextTakeover {
+		c.AddDict(data)
+	}
 	return w.Close()
 }