diff --git a/lazyClient.go b/lazyClient.go index 311e2cb..f8d9001 100644 --- a/lazyClient.go +++ b/lazyClient.go @@ -1,7 +1,6 @@ package multistream import ( - "bufio" "fmt" "io" "sync" @@ -98,7 +97,9 @@ func (l *lazyClientConn) doWriteHandshake() { // Perform the write handshake but *also* write some extra data. func (l *lazyClientConn) doWriteHandshakeWithData(extra []byte) int { - buf := bufio.NewWriter(l.con) + buf := getWriter(l.con) + defer putWriter(buf) + for _, proto := range l.protos { l.werr = delimWrite(buf, []byte(proto)) if l.werr != nil { diff --git a/multistream.go b/multistream.go index aab3236..85fd23e 100644 --- a/multistream.go +++ b/multistream.go @@ -19,6 +19,12 @@ var ErrTooLarge = errors.New("incoming message was too large") // the multistream muxers on both sides of a channel can work with each other. const ProtocolID = "/multistream/1.0.0" +var writerPool = sync.Pool{ + New: func() interface{} { + return bufio.NewWriter(nil) + }, +} + // HandlerFunc is a user-provided function used by the MultistreamMuxer to // handle a protocol/stream. type HandlerFunc func(protocol string, rwc io.ReadWriteCloser) error @@ -55,7 +61,9 @@ func writeUvarint(w io.Writer, i uint64) error { } func delimWriteBuffered(w io.Writer, mes []byte) error { - bw := bufio.NewWriter(w) + bw := getWriter(w) + defer putWriter(bw) + err := delimWrite(bw, mes) if err != nil { return err @@ -443,3 +451,14 @@ func (br *byteReader) ReadByte() (byte, error) { } return 0, err } + +func getWriter(w io.Writer) *bufio.Writer { + bw := writerPool.Get().(*bufio.Writer) + bw.Reset(w) + return bw +} + +func putWriter(bw *bufio.Writer) { + bw.Reset(nil) + writerPool.Put(bw) +}