change writer without buffer

使用了sync.Pool,从池中取出来的自带一定容量，所以无需使用Buffer

binary/writer.go

@@ -1,53 +1,44 @@
 package binary
 
 import (
-	"bytes"
 	"encoding/binary"
 	"encoding/hex"
 	"sync"
 )
 
-type Writer bytes.Buffer
+type Writer struct {
+	b []byte
+}
 
 var bufferPool = sync.Pool{
 	New: func() interface{} {
-		return new(bytes.Buffer)
+		return new(Writer)
 	},
 }
 
-func NewBuffer() *bytes.Buffer {
+func NewWriter() *Writer {
-	return bufferPool.Get().(*bytes.Buffer)
+	return bufferPool.Get().(*Writer)
 }
 
-func PutBuffer(buf *bytes.Buffer) {
+func PutBuffer(w *Writer) {
 	// See https://golang.org/issue/23199
 	const maxSize = 1 << 16
-	if buf.Cap() < maxSize { // 对于大Buffer直接丢弃
+	if cap(w.b) < maxSize { // 对于大Buffer直接丢弃
-		buf.Reset()
+		w.b = w.b[:0]
-		bufferPool.Put(buf)
+		bufferPool.Put(w)
 	}
 }
 
-func NewWriter() *Writer {
-	return new(Writer)
-}
-
-/*
-func PutWriter(w *Writer) {
-	PutBuffer((*bytes.Buffer)(w))
-}
-*/
-
 func NewWriterF(f func(writer *Writer)) []byte {
-	w := NewBuffer()
+	w := NewWriter()
-	f((*Writer)(w))
+	f(w)
 	b := append([]byte(nil), w.Bytes()...)
 	PutBuffer(w)
 	return b
 }
 
 func (w *Writer) Write(b []byte) {
-	(*bytes.Buffer)(w).Write(b)
+	w.b = append(w.b, b...)
 }
 
 func (w *Writer) WriteHex(h string) {
@@ -56,7 +47,7 @@ func (w *Writer) WriteHex(h string) {
 }
 
 func (w *Writer) WriteByte(b byte) {
-	(*bytes.Buffer)(w).WriteByte(b)
+	w.b = append(w.b, b)
 }
 
 func (w *Writer) WriteUInt16(v uint16) {
@@ -138,5 +129,5 @@ func (w *Writer) WriteTlvLimitedSize(data []byte, limit int) {
 }
 
 func (w *Writer) Bytes() []byte {
-	return (*bytes.Buffer)(w).Bytes()
+	return w.b
 }

binary/writer_test.go

@@ -1,71 +1,10 @@
 package binary
 
 import (
-	"bytes"
-	"encoding/hex"
-	"fmt"
 	"math/rand"
-	"sync"
 	"testing"
 )
 
-func NewWriterFOld(f func(writer *Writer)) []byte {
-	w := new(bytes.Buffer)
-	f((*Writer)(w))
-	return w.Bytes()
-}
-
-func TestNewWriterF(t *testing.T) {
-	wg := sync.WaitGroup{}
-	for i := 0; i < 10000; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			test := make([]byte, 1024)
-			rand.Read(test)
-			b1 := NewWriterFOld(func(writer *Writer) {
-				writer.Write(test)
-				writer.Write(NewWriterFOld(func(writer *Writer) {
-					writer.Write(test)
-					writer.Write(NewWriterFOld(func(writer *Writer) {
-						writer.Write(test)
-					}))
-				}))
-			})
-
-			b2 := NewWriterF(func(writer *Writer) {
-				writer.Write(test)
-				writer.Write(NewWriterF(func(writer *Writer) {
-					writer.Write(test)
-					writer.Write(NewWriterF(func(writer *Writer) {
-						writer.Write(test)
-					}))
-				}))
-			})
-
-			if !bytes.Equal(b1, b2) {
-				fmt.Println(hex.EncodeToString(b1))
-				fmt.Println(hex.EncodeToString(b2))
-				t.Error("Not equal!!!")
-			}
-		}()
-	}
-	wg.Wait()
-}
-
-func BenchmarkNewWriterFOld128(b *testing.B) {
-	test := make([]byte, 128)
-	rand.Read(test)
-	b.StartTimer()
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			NewWriterFOld(func(w *Writer) {
-				w.Write(test)
-			})
-		}
-	})
-}
-
 func BenchmarkNewWriterF128(b *testing.B) {
 	test := make([]byte, 128)
 	rand.Read(test)
@@ -79,21 +18,6 @@ func BenchmarkNewWriterF128(b *testing.B) {
 	})
 }
 
-func BenchmarkNewWriterFOld128_3(b *testing.B) {
-	test := make([]byte, 128)
-	rand.Read(test)
-	b.StartTimer()
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			NewWriterFOld(func(w *Writer) {
-				w.Write(test)
-				w.Write(test)
-				w.Write(test)
-			})
-		}
-	})
-}
-
 func BenchmarkNewWriterF128_3(b *testing.B) {
 	test := make([]byte, 128)
 	rand.Read(test)
@@ -109,23 +33,6 @@ func BenchmarkNewWriterF128_3(b *testing.B) {
 	})
 }
 
-func BenchmarkNewWriterFOld128_5(b *testing.B) {
-	test := make([]byte, 128)
-	rand.Read(test)
-	b.StartTimer()
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			NewWriterFOld(func(w *Writer) {
-				w.Write(test)
-				w.Write(test)
-				w.Write(test)
-				w.Write(test)
-				w.Write(test)
-			})
-		}
-	})
-}
-
 func BenchmarkNewWriterF128_5(b *testing.B) {
 	test := make([]byte, 128)
 	rand.Read(test)

protocol/packets/builders.go

@@ -7,7 +7,7 @@ import (
 )
 
 func BuildLoginPacket(uin int64, bodyType byte, key, body, extraData []byte) []byte {
-	w := binary.NewWriter()
+	return binary.NewWriterF(func(w *binary.Writer) {
 		w.WriteIntLvPacket(4, func(w *binary.Writer) {
 			w.WriteUInt32(0x00_00_00_0A)
 			w.WriteByte(bodyType)
@@ -22,11 +22,11 @@ func BuildLoginPacket(uin int64, bodyType byte, key, body, extraData []byte) []b
 				w.EncryptAndWrite(key, body)
 			}
 		})
-	return w.Bytes()
+	})
 }
 
 func BuildUniPacket(uin int64, seq uint16, commandName string, encryptType byte, sessionID, extraData, key, body []byte) []byte {
-	w := binary.NewWriter()
+	return binary.NewWriterF(func(w *binary.Writer) {
 		w.WriteIntLvPacket(4, func(w *binary.Writer) {
 			w.WriteUInt32(0x0B)
 			w.WriteByte(encryptType)
@@ -37,5 +37,5 @@ func BuildUniPacket(uin int64, seq uint16, commandName string, encryptType byte,
 				w.WriteUniPacket(commandName, sessionID, extraData, body)
 			}))
 		})
-	return w.Bytes()
+	})
 }

protocol/packets/global.go

@@ -33,11 +33,8 @@ type IEncryptMethod interface {
 }
 
 func BuildOicqRequestPacket(uin int64, commandId uint16, encrypt IEncryptMethod, key []byte, bodyFunc func(writer *binary.Writer)) []byte {
-	b := binary.NewWriter()
+	body := encrypt.DoEncrypt(binary.NewWriterF(bodyFunc), key)
-	bodyFunc(b)
+	return binary.NewWriterF(func(p *binary.Writer) {
-
-	body := encrypt.DoEncrypt(b.Bytes(), key)
-	p := binary.NewWriter()
 		p.WriteByte(0x02)
 		p.WriteUInt16(27 + 2 + uint16(len(body)))
 		p.WriteUInt16(8001)
@@ -52,7 +49,7 @@ func BuildOicqRequestPacket(uin int64, commandId uint16, encrypt IEncryptMethod,
 		p.WriteUInt32(0)
 		p.Write(body)
 		p.WriteByte(0x03)
-	return p.Bytes()
+	})
 }
 
 func BuildCode2DRequestPacket(seq uint32, j uint64, cmd uint16, bodyFunc func(writer *binary.Writer)) []byte {
@@ -73,7 +70,7 @@ func BuildCode2DRequestPacket(seq uint32, j uint64, cmd uint16, bodyFunc func(wr
 }
 
 func BuildSsoPacket(seq uint16, appID uint32, commandName, imei string, extData, outPacketSessionId, body, ksid []byte) []byte {
-	p := binary.NewWriter()
+	return binary.NewWriterF(func(p *binary.Writer) {
 		p.WriteIntLvPacket(4, func(writer *binary.Writer) {
 			writer.WriteUInt32(uint32(seq))
 			writer.WriteUInt32(appID)
@@ -101,7 +98,7 @@ func BuildSsoPacket(seq uint16, appID uint32, commandName, imei string, extData,
 		p.WriteIntLvPacket(4, func(writer *binary.Writer) {
 			writer.Write(body)
 		})
-	return p.Bytes()
+	})
 }
 
 func ParseIncomingPacket(payload, d2key []byte) (*IncomingPacket, error) {