From 73c0ea2f0cc69e3dca01ecb6fbbbdc0c960132f0 Mon Sep 17 00:00:00 2001
From: jaffee <jaffee@umbel.com>
Date: Thu, 22 Dec 2016 09:23:55 -0600
Subject: [PATCH] refactor/cleanup remote file creation

---
 cmd/pilosactl/multidbspawn.json |  2 +-
 creator/remote.go               | 19 +++-----------
 pilosactl/ssh.go                | 45 ++++++++++++++++++++++++++++++---
 3 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/cmd/pilosactl/multidbspawn.json b/cmd/pilosactl/multidbspawn.json
index 83d6d8a4c..6be8b40f3 100644
--- a/cmd/pilosactl/multidbspawn.json
+++ b/cmd/pilosactl/multidbspawn.json
@@ -1,5 +1,5 @@
 {
-    "CreatorArgs": ["-type", "local", "-serverN", "3", "-replicaN", "1", "-log-file-prefix", "multidblog"],
+    "CreatorArgs": ["-hosts", "localhost:19444,localhost:19445,localhost:19446", "-log-file-prefix", "multidblog"],
     "AgentHosts": ["localhost"],
     "Benchmarks": [
         {
diff --git a/creator/remote.go b/creator/remote.go
index dd8baa564..41c1c9ca7 100644
--- a/creator/remote.go
+++ b/creator/remote.go
@@ -57,21 +57,14 @@ func (c *RemoteCluster) Start() error {
 		conf.DataDir = "~/.pilosa" + port
 
 		// Connect to remote host
-		client, err := pilosactl.NewSSH(host, c.SSHUser, "")
-		if err != nil {
-			return err
-		}
-
-		// Create config file on remote host
-		sess, err := client.NewSession()
+		client, err := pilosactl.NewSSH(host, c.SSHUser, "", c.Stderr)
 		if err != nil {
 			return err
 		}
 		configname := "pilosa" + port + ".conf"
-		w, err := sess.StdinPipe()
-		err = sess.Start("cat > " + configname)
+		w, err := client.OpenFile(configname)
 		if err != nil {
-			return err
+			return fmt.Errorf("opening remote config file: %v", err)
 		}
 		enc := toml.NewEncoder(w)
 		err = enc.Encode(conf)
@@ -82,13 +75,9 @@ func (c *RemoteCluster) Start() error {
 		if err != nil {
 			return err
 		}
-		err = sess.Wait()
-		if err != nil {
-			return err
-		}
 
 		// Start pilosa on remote host
-		sess, err = client.NewSession()
+		sess, err := client.NewSession()
 		if err != nil {
 			return err
 		}
diff --git a/pilosactl/ssh.go b/pilosactl/ssh.go
index 38898dea3..98ab6d18a 100644
--- a/pilosactl/ssh.go
+++ b/pilosactl/ssh.go
@@ -2,6 +2,7 @@ package pilosactl
 
 import (
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"os/user"
@@ -13,12 +14,13 @@ import (
 
 type SSH struct {
 	client *ssh.Client
+	Stderr io.Writer
 }
 
 // NewSSH wraps up some of the complexity of using the crypto/ssh pacakge
 // directly assuming you want to connect using public key auth and you can pass
 // a keyfile or your key is accessible through ssh agent.
-func NewSSH(host, username, keyfile string) (*SSH, error) {
+func NewSSH(host, username, keyfile string, stderr io.Writer) (*SSH, error) {
 	if username == "" {
 		user, err := user.Current()
 		if err != nil {
@@ -52,13 +54,13 @@ func NewSSH(host, username, keyfile string) (*SSH, error) {
 		return nil, fmt.Errorf("NewSHH failed Dial: %v ", err)
 	}
 
-	return &SSH{client: client}, nil
+	return &SSH{client: client, Stderr: stderr}, nil
 }
 
-func SSHClients(hosts []string, username, keyfile string) ([]*SSH, error) {
+func SSHClients(hosts []string, username, keyfile string, stderr io.Writer) ([]*SSH, error) {
 	clients := make([]*SSH, len(hosts))
 	for i, host := range hosts {
-		client, err := NewSSH(host, username, keyfile)
+		client, err := NewSSH(host, username, keyfile, stderr)
 		if err != nil {
 			return nil, err
 		}
@@ -70,3 +72,38 @@ func SSHClients(hosts []string, username, keyfile string) ([]*SSH, error) {
 func (s *SSH) NewSession() (*ssh.Session, error) {
 	return s.client.NewSession()
 }
+
+type remoteFile struct {
+	w    io.WriteCloser
+	sess *ssh.Session
+}
+
+func (r *remoteFile) Write(p []byte) (n int, err error) {
+	return r.w.Write(p)
+}
+
+func (r *remoteFile) Close() error {
+	errc := r.w.Close()
+	errw := r.sess.Wait()
+	if errc != nil || errw != nil {
+		return fmt.Errorf("error closing remote file - close: '%v', wait: '%v'", errc, errw)
+	}
+	return nil
+}
+
+func (s *SSH) OpenFile(name string) (io.WriteCloser, error) {
+	sess, err := s.NewSession()
+	if err != nil {
+		return nil, err
+	}
+	w, err := sess.StdinPipe()
+	if err != nil {
+		return nil, err
+	}
+	err = sess.Start("cat > " + name)
+	if err != nil {
+		return nil, err
+	}
+
+	return &remoteFile{w: w, sess: sess}, nil
+}