ssh: add support for lazy spawning of pure SSH connections

Right now, we spawn potentially several connections without necessarily
needing to.  For example, if we're transferring two objects, we can
practically use at most three connections: one for the batch and one for
each object.

To make this more efficient and avoid needless overhead, let's not
actually create the connection until we attempt to acquire it.  At that
point, if it doesn't exist, we'll spawn a new one (using the control
path socket if possible) and then start sending data.  Otherwise, unless
it's the initial connection, let's just stub it out until we actually
need it.

Note that we still create the connection before it's needed because we
create all workers up front, but that will change in a future commit now
that we have this change in place.

Co-authored-by: Chris Darroch <[email protected]>

Author: bk2204

locking/ssh.go

@@ -17,7 +17,7 @@ type sshLockClient struct {
 	*lfsapi.Client
 }
 
-func (c *sshLockClient) connection() *ssh.PktlineConnection {
+func (c *sshLockClient) connection() (*ssh.PktlineConnection, error) {
 	return c.transfer.Connection(0)
 }
 
@@ -146,10 +146,13 @@ func (c *sshLockClient) Lock(remote string, lockReq *lockRequest) (*lockResponse
 	if lockReq.Ref != nil {
 		args = append(args, fmt.Sprintf("refname=%s", lockReq.Ref.Name))
 	}
-	conn := c.connection()
+	conn, err := c.connection()
+	if err != nil {
+		return nil, 0, err
+	}
 	conn.Lock()
 	defer conn.Unlock()
-	err := conn.SendMessage("lock", args)
+	err = conn.SendMessage("lock", args)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -168,10 +171,13 @@ func (c *sshLockClient) Unlock(ref *git.Ref, remote, id string, force bool) (*un
 	if ref != nil {
 		args = append(args, fmt.Sprintf("refname=%s", ref.Name))
 	}
-	conn := c.connection()
+	conn, err := c.connection()
+	if err != nil {
+		return nil, 0, err
+	}
 	conn.Lock()
 	defer conn.Unlock()
-	err := conn.SendMessage(fmt.Sprintf("unlock %s", id), args)
+	err = conn.SendMessage(fmt.Sprintf("unlock %s", id), args)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -191,10 +197,13 @@ func (c *sshLockClient) Search(remote string, searchReq *lockSearchRequest) (*lo
 	for key, value := range values {
 		args = append(args, fmt.Sprintf("%s=%s", key, value))
 	}
-	conn := c.connection()
+	conn, err := c.connection()
+	if err != nil {
+		return nil, 0, err
+	}
 	conn.Lock()
 	defer conn.Unlock()
-	err := conn.SendMessage("list-lock", args)
+	err = conn.SendMessage("list-lock", args)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -225,10 +234,13 @@ func (c *sshLockClient) SearchVerifiable(remote string, vreq *lockVerifiableRequ
 	if vreq.Limit > 0 {
 		args = append(args, fmt.Sprintf("limit=%d", vreq.Limit))
 	}
-	conn := c.connection()
+	conn, err := c.connection()
+	if err != nil {
+		return nil, 0, err
+	}
 	conn.Lock()
 	defer conn.Unlock()
-	err := conn.SendMessage("list-lock", args)
+	err = conn.SendMessage("list-lock", args)
 	if err != nil {
 		return nil, 0, err
 	}

ssh/connection.go

@@ -86,14 +86,28 @@ func (tr *SSHTransfer) IsMultiplexingEnabled() bool {
 }
 
 // Connection returns the nth connection (starting from 0) in this transfer
-// instance or nil if there is no such item.
-func (tr *SSHTransfer) Connection(n int) *PktlineConnection {
+// instance if it is initialized and otherwise initializes a new connection and
+// saves it in the nth position.  In all cases, nil is returned if n is greater
+// than the maximum number of connections.
+func (tr *SSHTransfer) Connection(n int) (*PktlineConnection, error) {
 	tr.lock.RLock()
-	defer tr.lock.RUnlock()
 	if n >= len(tr.conn) {
-		return nil
+		tr.lock.RUnlock()
+		return nil, nil
+	}
+	if tr.conn[n] != nil {
+		defer tr.lock.RUnlock()
+		return tr.conn[n], nil
+	}
+	tr.lock.RUnlock()
+
+	tr.lock.Lock()
+	defer tr.lock.Unlock()
+	if tr.conn[n] != nil {
+		return tr.conn[n], nil
 	}
-	return tr.conn[n]
+	conn, _, err := tr.spawnConnection(n)
+	return conn, err
 }
 
 // ConnectionCount returns the number of connections this object has.
@@ -122,6 +136,15 @@ func (tr *SSHTransfer) SetConnectionCountAtLeast(n int) error {
 	return tr.setConnectionCount(n)
 }
 
+func (tr *SSHTransfer) spawnConnection(n int) (*PktlineConnection, string, error) {
+	conn, _, controlPath, err := startConnection(n, tr.osEnv, tr.gitEnv, tr.meta, tr.operation, tr.controlPath)
+	if err != nil {
+		tracerx.Printf("failed to spawn pure SSH connection: %s", err)
+		return nil, "", err
+	}
+	return conn, controlPath, err
+}
+
 func (tr *SSHTransfer) setConnectionCount(n int) error {
 	count := len(tr.conn)
 	if n < count {
@@ -130,6 +153,10 @@ func (tr *SSHTransfer) setConnectionCount(n int) error {
 			tn = 1
 		}
 		for _, item := range tr.conn[tn:count] {
+			if item == nil {
+				tracerx.Printf("skipping uninitialized lazy pure SSH connection (%d -> %d)", count, n)
+				continue
+			}
 			tracerx.Printf("terminating pure SSH connection (%d -> %d)", count, n)
 			if err := item.End(); err != nil {
 				return err
@@ -138,14 +165,15 @@ func (tr *SSHTransfer) setConnectionCount(n int) error {
 		tr.conn = tr.conn[0:tn]
 	} else if n > count {
 		for i := count; i < n; i++ {
-			conn, _, controlPath, err := startConnection(i, tr.osEnv, tr.gitEnv, tr.meta, tr.operation, tr.controlPath)
-			if err != nil {
-				tracerx.Printf("failed to spawn pure SSH connection: %s", err)
-				return err
-			}
-			tr.conn = append(tr.conn, conn)
 			if i == 0 {
+				conn, controlPath, err := tr.spawnConnection(i)
+				if err != nil {
+					return err
+				}
+				tr.conn = append(tr.conn, conn)
 				tr.controlPath = controlPath
+			} else {
+				tr.conn = append(tr.conn, nil)
 			}
 		}
 	}

tq/ssh.go

@@ -25,13 +25,13 @@ type SSHBatchClient struct {
 }
 
 func (a *SSHBatchClient) batchInternal(args []string, batchLines []string) (int, []string, []string, error) {
-	conn := a.transfer.Connection(0)
-	if conn == nil {
-		return 0, nil, nil, errors.Errorf(tr.Tr.Get("could not get connection for batch request"))
+	conn, err := a.transfer.Connection(0)
+	if err != nil {
+		return 0, nil, nil, errors.Wrap(err, tr.Tr.Get("could not get connection for batch request"))
 	}
 	conn.Lock()
 	defer conn.Unlock()
-	err := conn.SendMessageWithLines("batch", args, batchLines)
+	err = conn.SendMessageWithLines("batch", args, batchLines)
 	if err != nil {
 		return 0, nil, nil, errors.Wrap(err, tr.Tr.Get("batch request"))
 	}
@@ -158,7 +158,7 @@ type SSHAdapter struct {
 // Implementations can run some startup logic here & return some context if needed
 func (a *SSHAdapter) WorkerStarting(workerNum int) (interface{}, error) {
 	a.transfer.SetConnectionCountAtLeast(workerNum + 1)
-	return a.transfer.Connection(workerNum), nil
+	return a.transfer.Connection(workerNum)
 }
 
 // WorkerEnding is called when a worker goroutine is shutting down