Reimplement Wait method to make it can be called concurrently from multiple goroutines

wutz · wutz · commit 3a53acc96bf9 · 2018-12-05T14:47:03.000+08:00
diff --git a/client.go b/client.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"os"
 	"path"
+	"sync"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -122,7 +123,7 @@ func NewClientPipe(rd io.Reader, wr io.WriteCloser, opts ...ClientOption) (*Clie
 				WriteCloser: wr,
 			},
 			inflight: make(map[uint32]chan<- result),
-			errCh:    make(chan error, 1),
+			errCond:  sync.NewCond(new(sync.Mutex)),
 		},
 		maxPacket:             1 << 15,
 		maxConcurrentRequests: 64,
diff --git a/conn.go b/conn.go
@@ -36,13 +36,21 @@ type clientConn struct {
 	wg         sync.WaitGroup
 	sync.Mutex                          // protects inflight
 	inflight   map[uint32]chan<- result // outstanding requests
-	errCh      chan error
+
+	errCond *sync.Cond
+	err     error
 }
 
 // Wait blocks until the conn has shut down, and return the error
-// causing the shutdown.
+// causing the shutdown. It can be called concurrently from multiple
+// goroutines.
 func (c *clientConn) Wait() error {
-	return <-c.errCh
+	c.errCond.L.Lock()
+	defer c.errCond.L.Unlock()
+	for c.err == nil {
+		c.errCond.Wait()
+	}
+	return c.err
 }
 
 // Close closes the SFTP session.
@@ -56,7 +64,10 @@ func (c *clientConn) loop() {
 	err := c.recv()
 	if err != nil {
 		c.broadcastErr(err)
-		c.errCh <- err
+		c.errCond.L.Lock()
+		c.err = err
+		c.errCond.Broadcast()
+		c.errCond.L.Unlock()
 	}
 }
 
