Use connect pool

Author: wzshiming

client.go

@@ -15,11 +15,11 @@ import (
 // NewDialer returns a new Dialer that dials through the provided
 // proxy server's network and address.
 func NewDialer(addr string) (*Dialer, error) {
-	host, config, err := clientConfig(addr)
+	config, err := parseClientConfig(addr)
 	if err != nil {
 		return nil, err
 	}
-	return NewDialerWithConfig(host, config)
+	return NewDialerWithConfig(config.host, config.clientConfig)
 }
 
 func NewDialerWithConfig(host string, config *ssh.ClientConfig) (*Dialer, error) {
@@ -29,10 +29,15 @@ func NewDialerWithConfig(host string, config *ssh.ClientConfig) (*Dialer, error)
 	}, nil
 }
 
-func clientConfig(addr string) (host string, config *ssh.ClientConfig, err error) {
+type clientConfig struct {
+	host         string
+	clientConfig *ssh.ClientConfig
+}
+
+func parseClientConfig(addr string) (*clientConfig, error) {
 	ur, err := url.Parse(addr)
 	if err != nil {
-		return "", nil, err
+		return nil, err
 	}
 
 	user := ""
@@ -43,7 +48,7 @@ func clientConfig(addr string) (host string, config *ssh.ClientConfig, err error
 		pwd, isPwd = ur.User.Password()
 	}
 
-	config = &ssh.ClientConfig{
+	config := &ssh.ClientConfig{
 		User:            user,
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
@@ -54,23 +59,26 @@ func clientConfig(addr string) (host string, config *ssh.ClientConfig, err error
 
 	identityDatas, err := getQuery(ur.Query()["identity_data"], ur.Query()["identity_file"])
 	if err != nil {
-		return "", nil, err
+		return nil, err
 	}
 	for _, data := range identityDatas {
 		signer, err := ssh.ParsePrivateKey(data)
 		if err != nil {
-			return "", nil, err
+			return nil, err
 		}
 		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
 	}
 
-	host = ur.Hostname()
+	host := ur.Hostname()
 	port := ur.Port()
 	if port == "" {
 		port = "22"
 	}
-	host = net.JoinHostPort(host, port)
-	return host, config, nil
+
+	return &clientConfig{
+		clientConfig: config,
+		host:         net.JoinHostPort(host, port),
+	}, nil
 }
 
 type Dialer struct {
@@ -79,20 +87,22 @@ type Dialer struct {
 	// ProxyDial specifies the optional dial function for
 	// establishing the transport connection.
 	ProxyDial func(context.Context, string, string) (net.Conn, error)
-	sshCli    *ssh.Client
-	host      string
-	config    *ssh.ClientConfig
+
+	host   string
+	config *ssh.ClientConfig
+
+	pool sync.Pool
 }
 
 func (d *Dialer) Close() error {
-	d.mut.Lock()
-	defer d.mut.Unlock()
-	if d.sshCli == nil {
-		return nil
+	for {
+		a := d.pool.Get()
+		if a == nil {
+			break
+		}
+		a.(*ssh.Client).Close()
 	}
-	err := d.sshCli.Close()
-	d.sshCli = nil
-	return err
+	return nil
 }
 
 func (d *Dialer) proxyDial(ctx context.Context, network, address string) (net.Conn, error) {
@@ -105,12 +115,36 @@ func (d *Dialer) proxyDial(ctx context.Context, network, address string) (net.Co
 }
 
 func (d *Dialer) SSHClient(ctx context.Context) (*ssh.Client, error) {
+	return d.GetClient(ctx)
+}
+
+func (d *Dialer) GetClient(ctx context.Context) (*ssh.Client, error) {
+	a := d.pool.Get()
+	if a != nil {
+		return a.(*ssh.Client), nil
+	}
+
 	d.mut.Lock()
 	defer d.mut.Unlock()
-	cli := d.sshCli
-	if cli != nil {
-		return cli, nil
+
+	a = d.pool.Get()
+	if a != nil {
+		return a.(*ssh.Client), nil
+	}
+
+	cli, err := d.sshClient(ctx)
+	if err != nil {
+		return nil, err
 	}
+
+	return cli, nil
+}
+
+func (d *Dialer) PutClient(cli *ssh.Client) {
+	d.pool.Put(cli)
+}
+
+func (d *Dialer) sshClient(ctx context.Context) (*ssh.Client, error) {
 	conn, err := d.proxyDial(ctx, "tcp", d.host)
 	if err != nil {
 		return nil, err
@@ -120,39 +154,42 @@ func (d *Dialer) SSHClient(ctx context.Context) (*ssh.Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	cli = ssh.NewClient(con, chans, reqs)
-	d.sshCli = cli
-	return cli, nil
+	return ssh.NewClient(con, chans, reqs), nil
 }
 
-func (d *Dialer) CommandDialContext(ctx context.Context, name string, args ...string) (net.Conn, error) {
-	cmd := make([]string, 0, len(args)+1)
-	cmd = append(cmd, name)
+func buildCmd(name string, args ...string) string {
+	cmds := make([]string, 0, len(args)+1)
+	cmds = append(cmds, name)
 	for _, arg := range args {
-		cmd = append(cmd, strconv.Quote(arg))
+		cmds = append(cmds, strconv.Quote(arg))
 	}
-	return d.commandDialContext(ctx, strings.Join(cmd, " "), 1)
+	return strings.Join(cmds, " ")
 }
 
-func (d *Dialer) commandDialContext(ctx context.Context, cmd string, retry int) (net.Conn, error) {
-	cli, err := d.SSHClient(ctx)
+func (d *Dialer) CommandDialContext(ctx context.Context, name string, args ...string) (net.Conn, error) {
+	cli, err := d.GetClient(ctx)
 	if err != nil {
 		return nil, err
 	}
 	sess, err := cli.NewSession()
 	if err != nil {
+		if isSSHError(err) {
+			d.PutClient(cli)
+		} else {
+			cli.Close()
+		}
 		return nil, err
 	}
+	defer d.PutClient(cli)
+
 	conn1, conn2 := net.Pipe()
 	sess.Stdin = conn1
 	sess.Stdout = conn1
 	sess.Stderr = os.Stderr
+
+	cmd := buildCmd(name, args...)
 	err = sess.Start(cmd)
 	if err != nil {
-		if retry != 0 {
-			d.Close()
-			return d.commandDialContext(ctx, cmd, retry-1)
-		}
 		return nil, err
 	}
 	ctx, cancel := context.WithCancel(ctx)
@@ -180,25 +217,42 @@ func (d *Dialer) commandDialContext(ctx context.Context, cmd string, retry int)
 }
 
 func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
-	return d.dialContext(ctx, network, address, 1)
-}
-
-func (d *Dialer) dialContext(ctx context.Context, network, address string, retry int) (net.Conn, error) {
-	cli, err := d.SSHClient(ctx)
+	cli, err := d.GetClient(ctx)
 	if err != nil {
 		return nil, err
 	}
+
 	conn, err := cli.DialContext(ctx, network, address)
 	if err != nil {
-		if ctx.Err() != nil {
-			return nil, err
+		if isSSHError(err) {
+			d.PutClient(cli)
+		} else {
+			cli.Close()
 		}
-		if retry != 0 {
-			d.Close()
-			return d.dialContext(ctx, network, address, retry-1)
+		return nil, err
+	}
+
+	d.PutClient(cli)
+	return conn, nil
+}
+
+func (d *Dialer) Dial(network, address string) (net.Conn, error) {
+	cli, err := d.GetClient(context.Background())
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := cli.Dial(network, address)
+	if err != nil {
+		if isSSHError(err) {
+			d.PutClient(cli)
+		} else {
+			cli.Close()
 		}
 		return nil, err
 	}
+
+	d.PutClient(cli)
 	return conn, nil
 }
 
@@ -209,21 +263,44 @@ func (d *Dialer) Listen(ctx context.Context, network, address string) (net.Liste
 			address = net.JoinHostPort("0.0.0.0", port)
 		}
 	}
-	return d.listen(ctx, network, address, 1)
-}
 
-func (d *Dialer) listen(ctx context.Context, network, address string, retry int) (net.Listener, error) {
-	cli, err := d.SSHClient(ctx)
+	cli, err := d.GetClient(ctx)
 	if err != nil {
 		return nil, err
 	}
+
 	listener, err := cli.Listen(network, address)
 	if err != nil {
-		if retry != 0 {
-			d.Close()
-			return d.listen(ctx, network, address, retry-1)
+		if isSSHError(err) {
+			d.PutClient(cli)
+		} else {
+			cli.Close()
 		}
 		return nil, err
 	}
+
+	listener = &listenerWithCleanup{
+		Listener: listener,
+		cleanup: func() {
+			d.PutClient(cli)
+		},
+	}
+
 	return listener, nil
 }
+
+type listenerWithCleanup struct {
+	net.Listener
+	cleanup func()
+}
+
+func (l *listenerWithCleanup) Close() error {
+	err := l.Listener.Close()
+	l.cleanup()
+	return err
+}
+
+func isSSHError(err error) bool {
+	msg := err.Error()
+	return strings.Contains(msg, "ssh: ")
+}