Merge pull request #621 from djs55/timeout

go: allow concurrent Write calls

Author: djs55

go/pkg/libproxy/multiplexed.go

@@ -100,13 +100,15 @@ func (c *channel) sendWindowUpdate() error {
 	c.read.advance()
 	seq := c.read.allowed
 	c.m.Unlock()
-	return c.multiplexer.send(NewWindow(c.ID, seq))
+	return c.multiplexer.send(NewWindow(c.ID, seq), nil)
 }
 
 func (c *channel) recvWindowUpdate(seq uint64) {
 	c.m.Lock()
 	c.write.allowed = seq
-	c.c.Signal()
+	// net.Conn says: Multiple goroutines may invoke methods on a Conn simultaneously.
+	// Therefore there can be multiple goroutines blocked in Write, so when the window opens we should wake them all up.
+	c.c.Broadcast()
 	c.m.Unlock()
 }
 
@@ -148,56 +150,40 @@ func (c *channel) Write(p []byte) (int, error) {
 			c.write.current = c.write.current + uint64(toWrite)
 
 			// Don't block holding the metadata mutex.
-			// Note this would allow concurrent calls to Write on the same channel
-			// to conflict, but we regard that as user error.
 			c.m.Unlock()
-
-			// need to write the header and the payload together
-			c.multiplexer.writeMutex.Lock()
-			f := NewData(c.ID, uint32(toWrite))
-			c.multiplexer.appendEvent(&event{eventType: eventSend, frame: f})
-			err1 := f.Write(c.multiplexer.connW)
-			_, err2 := c.multiplexer.connW.Write(p[0:toWrite])
-			err3 := c.multiplexer.connW.Flush()
-			c.multiplexer.writeMutex.Unlock()
-
+			err := c.multiplexer.send(NewData(c.ID, uint32(toWrite)), p[0:toWrite])
 			c.m.Lock()
-			if err1 != nil {
-				return written, err1
-			}
-			if err2 != nil {
-				return written, err2
-			}
-			if err3 != nil {
-				return written, err3
+
+			if err != nil {
+				return written, err
 			}
 			p = p[toWrite:]
 			written = written + toWrite
 			continue
 		}
 
-		// Wait for the write window to be increased (or a timeout)
-		done := make(chan struct{})
-		timeout := make(chan time.Time)
+		// If the client has set a deadline then create a timer:
+		var (
+			timer   *time.Timer
+			timeOut bool
+		)
 		if !c.writeDeadline.IsZero() {
-			go func() {
-				time.Sleep(time.Until(c.writeDeadline))
-				close(timeout)
-			}()
+			timer = time.AfterFunc(time.Until(c.writeDeadline), func() {
+				c.m.Lock()
+				defer c.m.Unlock()
+				timeOut = true
+				c.c.Broadcast()
+			})
+		}
+
+		// Wait for the write window to be increased or a timeout
+		c.c.Wait()
+
+		if timer != nil {
+			timer.Stop()
 		}
-		go func() {
-			c.c.Wait()
-			close(done)
-		}()
-		select {
-		case <-timeout:
-			// clean up the goroutine
-			c.c.Broadcast()
-			<-done
+		if timeOut {
 			return written, &errTimeout{}
-		case <-done:
-			// The timeout will still fire in the background
-			continue
 		}
 	}
 }
@@ -213,7 +199,7 @@ func (c *channel) Close() error {
 	if alreadyClosed {
 		return nil
 	}
-	if err := c.multiplexer.send(NewClose(c.ID)); err != nil {
+	if err := c.multiplexer.send(NewClose(c.ID), nil); err != nil {
 		return err
 	}
 	c.m.Lock()
@@ -240,7 +226,7 @@ func (c *channel) CloseWrite() error {
 	if alreadyShutdown {
 		return nil
 	}
-	if err := c.multiplexer.send(NewShutdown(c.ID)); err != nil {
+	if err := c.multiplexer.send(NewShutdown(c.ID), nil); err != nil {
 		return err
 	}
 	c.m.Lock()
@@ -428,14 +414,22 @@ func (m *multiplexer) appendEvent(e *event) {
 	m.events = m.events.Next()
 }
 
-func (m *multiplexer) send(f *Frame) error {
+// send a frame (header) plus optional payload. If this call fails then the multiplexed connection will be desynchronised.
+func (m *multiplexer) send(f *Frame, payload []byte) error {
 	m.writeMutex.Lock()
 	defer m.writeMutex.Unlock()
+	m.appendEvent(&event{eventType: eventSend, frame: f})
+
 	if err := f.Write(m.connW); err != nil {
-		return err
+		return fmt.Errorf("writing frame %s: %w", f, err)
 	}
-	m.appendEvent(&event{eventType: eventSend, frame: f})
-	return m.connW.Flush()
+	if n, err := m.connW.Write(payload); err != nil || n != len(payload) {
+		return fmt.Errorf("writing frame %s payload length %d: %d, %w", f, len(payload), n, err)
+	}
+	if err := m.connW.Flush(); err != nil {
+		return fmt.Errorf("flushing frame %s: %w", f, err)
+	}
+	return nil
 }
 
 func (m *multiplexer) findFreeChannelID() uint32 {
@@ -485,7 +479,7 @@ func (m *multiplexer) Dial(d Destination) (MultiplexedConn, error) {
 	m.channels[id] = channel
 	m.metadataMutex.Unlock()
 
-	if err := m.send(NewOpen(id, d)); err != nil {
+	if err := m.send(NewOpen(id, d), nil); err != nil {
 		return nil, err
 	}
 	if err := channel.sendWindowUpdate(); err != nil {