diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 312853e24..6537c6dd8 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -192,6 +192,25 @@ func SplitSize(mb MultiBuffer, size int32) (MultiBuffer, MultiBuffer) { return mb, r } +// WriteMultiBuffer writes all buffers from the MultiBuffer to the Writer one by one, and return error if any, with leftover MultiBuffer. +func WriteMultiBuffer(writer io.Writer, mb MultiBuffer) (MultiBuffer, error) { + for { + mb2, b := SplitFirst(mb) + mb = mb2 + if b == nil { + break + } + + _, err := writer.Write(b.Bytes()) + b.Release() + if err != nil { + return mb, nil + } + } + + return nil, nil +} + // Len returns the total number of bytes in the MultiBuffer. func (mb MultiBuffer) Len() int32 { if mb == nil { diff --git a/common/buf/writer.go b/common/buf/writer.go index 275ccf694..8e22a16d6 100644 --- a/common/buf/writer.go +++ b/common/buf/writer.go @@ -219,19 +219,9 @@ type SequentialWriter struct { // WriteMultiBuffer implements Writer. func (w *SequentialWriter) WriteMultiBuffer(mb MultiBuffer) error { - defer ReleaseMulti(mb) - - for _, b := range mb { - if b.IsEmpty() { - continue - } - - if err := WriteAllBytes(w.Writer, b.Bytes()); err != nil { - return err - } - } - - return nil + mb, err := WriteMultiBuffer(w.Writer, mb) + ReleaseMulti(mb) + return err } type noOpWriter byte diff --git a/transport/internet/quic/conn.go b/transport/internet/quic/conn.go index afefbc09f..fc4f10f14 100644 --- a/transport/internet/quic/conn.go +++ b/transport/internet/quic/conn.go @@ -172,21 +172,9 @@ func (c *interConn) ReadMultiBuffer() (buf.MultiBuffer, error) { func (c *interConn) WriteMultiBuffer(mb buf.MultiBuffer) error { mb = buf.Compact(mb) - - for { - mb2, b := buf.SplitFirst(mb) - mb = mb2 - if b == nil { - break - } - - if err := buf.WriteAllBytes(c, b.Bytes()); err != nil { - buf.ReleaseMulti(mb) - return err - } - } - - return nil + mb, err := buf.WriteMultiBuffer(c, mb) + buf.ReleaseMulti(mb) + return err } func (c *interConn) Write(b []byte) (int, error) { diff --git a/transport/internet/websocket/connection.go b/transport/internet/websocket/connection.go index c15e7f7df..f9a2f3a17 100644 --- a/transport/internet/websocket/connection.go +++ b/transport/internet/websocket/connection.go @@ -70,21 +70,9 @@ func (c *connection) Write(b []byte) (int, error) { func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) error { mb = buf.Compact(mb) - - for { - mb2, b := buf.SplitFirst(mb) - mb = mb2 - if b == nil { - break - } - - if err := buf.WriteAllBytes(c, b.Bytes()); err != nil { - buf.ReleaseMulti(mb) - return err - } - } - - return nil + mb, err := buf.WriteMultiBuffer(c, mb) + buf.ReleaseMulti(mb) + return err } func (c *connection) Close() error {