basic implementation of remote agent spawning

cmd/pilosactl/main.go

@@ -24,6 +24,8 @@ import (
 	"time"
 	"unsafe"
 
+	"golang.org/x/crypto/ssh"
+
 	"github.com/pilosa/pilosa"
 	"github.com/pilosa/pilosa/bench"
 	"github.com/pilosa/pilosa/creator"
@@ -1309,24 +1311,19 @@ type BspawnCommand struct {
 	// If AgentHosts is specified, Agents is ignored, and the existing
 	// agents specified here are used.
 	AgentHosts []string
-	// Agents is config for creating a fleet of agents from which to run the
-	// benchmark. TODO: mostly unimplemented.
-	Agents AgentConfig
 
 	// Benchmarks is a slice of Spawns which specifies all of the bagent
 	// commands to run. These will all be run in parallel, started on each
 	// of the agents in a round robin fashion.
 	Benchmarks []Spawn
 
+	SSHUser string
+
 	Stdin  io.Reader `json:"-"`
 	Stdout io.Writer `json:"-"`
 	Stderr io.Writer `json:"-"`
 }
 
-type AgentConfig struct {
-	Type string
-}
-
 // Spawn represents a bagent command run in parallel across Num agents. The
 // bagent command can run multiple Benchmarks serially within itself.
 type Spawn struct {
@@ -1381,7 +1378,10 @@ func (cmd *BspawnCommand) Run(ctx context.Context) error {
 		// must create cluster
 		r, w := io.Pipe()
 		createCmd := NewCreateCommand(cmd.Stdin, w, cmd.Stderr)
-		createCmd.ParseFlags(cmd.CreatorArgs)
+		err := createCmd.ParseFlags(cmd.CreatorArgs)
+		if err != nil {
+			return err
+		}
 		go func() {
 			err := createCmd.Run(ctx)
 			if err != nil {
@@ -1390,22 +1390,51 @@ func (cmd *BspawnCommand) Run(ctx context.Context) error {
 		}()
 		clus := &CreateOutput{}
 		dec := json.NewDecoder(r)
-		err := dec.Decode(clus)
+		err = dec.Decode(clus)
 		if err != nil {
 			return err
 		}
 		cmd.PilosaHosts = clus.Hosts
 	}
-	switch cmd.Agents.Type {
-	case "local":
+	if len(cmd.AgentHosts) > 0 {
+		return cmd.spawnRemote(ctx)
+	} else {
 		return cmd.spawnLocal(ctx)
-	case "remote":
-		return fmt.Errorf("remote type spawning is unimplemented")
-	default:
-		return fmt.Errorf("'%v' is not a supported type of spawn command", cmd.Agents.Type)
 	}
 }
 
+func (cmd *BspawnCommand) spawnRemote(ctx context.Context) error {
+	agentIndex := 0
+	agentConnections, err := pilosactl.SSHClients(cmd.AgentHosts, cmd.SSHUser, "")
+	if err != nil {
+		return err
+	}
+	sessions := make([]*ssh.Session, 0)
+	for _, sp := range cmd.Benchmarks {
+		for i := 0; i < sp.Num; i++ {
+			sess, err := agentConnections[agentIndex].NewSession()
+			if err != nil {
+				return err
+			}
+			sessions = append(sessions, sess)
+			sess.Stdout = cmd.Stdout
+			sess.Stderr = cmd.Stderr
+			err = sess.Start("pilosactl bagent -agentNum=" + strconv.Itoa(i) + " -hosts=" + strings.Join(cmd.PilosaHosts, ",") + " " + strings.Join(sp.Args, " "))
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	for _, sess := range sessions {
+		err = sess.Wait()
+		if err != nil {
+			return fmt.Errorf("error waiting for remote bagent: %v", err)
+		}
+	}
+	return nil
+}
+
 func (cmd *BspawnCommand) spawnLocal(ctx context.Context) error {
 	agents := []*BagentCommand{}
 	for _, sp := range cmd.Benchmarks {

cmd/pilosactl/multidbspawn.json

@@ -1,6 +1,6 @@
 {
     "CreatorArgs": ["-type", "local", "-serverN", "3", "-replicaN", "1", "-log-file-prefix", "multidblog"],
-    "Agents": { "Type": "local" },
+    "AgentHosts": ["localhost"],
     "Benchmarks": [
         {
             "Num": 3,

pilosactl/ssh.go

@@ -55,6 +55,18 @@ func NewSSH(host, username, keyfile string) (*SSH, error) {
 	return &SSH{client: client}, nil
 }
 
+func SSHClients(hosts []string, username, keyfile string) ([]*SSH, error) {
+	clients := make([]*SSH, len(hosts))
+	for i, host := range hosts {
+		client, err := NewSSH(host, username, keyfile)
+		if err != nil {
+			return nil, err
+		}
+		clients[i] = client
+	}
+	return clients, nil
+}
+
 func (s *SSH) NewSession() (*ssh.Session, error) {
 	return s.client.NewSession()
 }