coalesce writes

multiplex.go

@@ -59,7 +59,7 @@ type Multiplex struct {
 	shutdownErr  error
 	shutdownLock sync.Mutex
 
-	wrTkn chan struct{}
+	writeCh chan []byte
 
 	nstreams chan *Stream
 
@@ -76,13 +76,12 @@ func NewMultiplex(con net.Conn, initiator bool) *Multiplex {
 		channels:  make(map[streamID]*Stream),
 		closed:    make(chan struct{}),
 		shutdown:  make(chan struct{}),
-		wrTkn:     make(chan struct{}, 1),
+		writeCh:   make(chan []byte, 16),
 		nstreams:  make(chan *Stream, 16),
 	}
 
 	go mp.handleIncoming()
-
-	mp.wrTkn <- struct{}{}
+	go mp.handleOutgoing()
 
 	return mp
 }
@@ -146,45 +145,101 @@ func (mp *Multiplex) IsClosed() bool {
 
 func (mp *Multiplex) sendMsg(ctx context.Context, header uint64, data []byte) error {
 	buf := pool.Get(len(data) + 20)
-	defer pool.Put(buf)
 
 	n := 0
 	n += binary.PutUvarint(buf[n:], header)
 	n += binary.PutUvarint(buf[n:], uint64(len(data)))
 	n += copy(buf[n:], data)
 
 	select {
-	case tkn := <-mp.wrTkn:
-		defer func() { mp.wrTkn <- tkn }()
+	case mp.writeCh <- buf[:n]:
+		return nil
+	case <-mp.shutdown:
+		return ErrShutdown
 	case <-ctx.Done():
 		return ctx.Err()
 	}
+}
 
-	if mp.isShutdown() {
-		return ErrShutdown
-	}
+func (mp *Multiplex) handleOutgoing() {
+	for {
+		select {
+		case <-mp.shutdown:
+			return
 
-	dl, hasDl := ctx.Deadline()
-	if hasDl {
-		if err := mp.con.SetWriteDeadline(dl); err != nil {
-			return err
+		case data := <-mp.writeCh:
+			err := mp.writeMsg(data)
+			if err != nil {
+				log.Warningf("Error writing data: %s", err.Error())
+			}
 		}
 	}
+}
 
-	written, err := mp.con.Write(buf[:n])
-	if err != nil && (written > 0 || isFatalNetworkError(err)) {
-		// Bail. We've written partial message or it's a fatal error and can't do anything
-		// about this.
-		mp.closeNoWait()
-		return err
+func (mp *Multiplex) writeMsg(data []byte) error {
+	if len(data) >= 512 {
+		return mp.doWriteMsg(data)
 	}
 
-	if hasDl {
-		// only return this error if we don't *already* have an error from the write.
-		if err2 := mp.con.SetWriteDeadline(time.Time{}); err == nil && err2 != nil {
-			return err2
+	buf := pool.Get(4096)
+	defer pool.Put(buf)
+
+	n := copy(buf, data)
+	pool.Put(data)
+
+	timer := time.NewTimer(100 * time.Millisecond)
+	defer timer.Stop()
+
+	for {
+		select {
+		case data = <-mp.writeCh:
+			wr := copy(buf[n:], data)
+			if wr < len(data) {
+				// we filled the buffer, send it
+				err := mp.doWriteMsg(buf)
+				if err != nil {
+					pool.Put(data)
+					return err
+				}
+
+				if len(data)-wr >= 512 {
+					// the remaining data is not a small write, send it
+					err := mp.doWriteMsg(data[wr:])
+					pool.Put(data)
+					return err
+				}
+
+				n = copy(buf, data[wr:])
+
+				// we've written some, reset the timer to coalesce the rest
+				if !timer.Stop() {
+					<-timer.C
+				}
+				timer.Reset(100 * time.Millisecond)
+			} else {
+				n += wr
+			}
+
+			pool.Put(data)
+
+		case <-timer.C:
+			return mp.doWriteMsg(buf[:n])
+
+		case <-mp.shutdown:
+			return ErrShutdown
 		}
 	}
+}
+
+func (mp *Multiplex) doWriteMsg(data []byte) error {
+	if mp.isShutdown() {
+		return ErrShutdown
+	}
+
+	_, err := mp.con.Write(data)
+	if err != nil {
+		mp.closeNoWait()
+	}
 
 	return err
 }

multiplex_test.go

@@ -40,7 +40,7 @@ func TestSlowReader(t *testing.T) {
 
 	// 100 is large enough that the buffer of the underlying connection will
 	// fill up.
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10000; i++ {
 		_, err = sa.Write(mes)
 		if err != nil {
 			break
@@ -346,7 +346,7 @@ func TestReset(t *testing.T) {
 		t.Fatalf("successfully wrote to reset stream")
 	}
 
-	time.Sleep(10 * time.Millisecond)
+	time.Sleep(200 * time.Millisecond)
 
 	n, err = sb.Write([]byte("test"))
 	if n != 0 {