Browse Source

merge duplicated code

Darien Raymond 7 years ago
parent
commit
88e757e33f

+ 19 - 0
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 {

+ 3 - 13
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

+ 3 - 15
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) {

+ 3 - 15
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 {