diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index aea1aee0a..92d75128b 100644 --- a/cmd/pilosactl/main.go +++ b/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 { diff --git a/cmd/pilosactl/multidbspawn.json b/cmd/pilosactl/multidbspawn.json index 08c00f776..83d6d8a4c 100644 --- a/cmd/pilosactl/multidbspawn.json +++ b/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, diff --git a/pilosactl/ssh.go b/pilosactl/ssh.go index 95795c381..38898dea3 100644 --- a/pilosactl/ssh.go +++ b/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() }