diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 1cfdedd5c..459d3a32f 100644 --- a/cmd/pilosactl/main.go +++ b/cmd/pilosactl/main.go @@ -1534,7 +1534,7 @@ func (cmd *BspawnCommand) Run(ctx context.Context) error { } func (cmd *BspawnCommand) spawnRemote(ctx context.Context) (map[string]interface{}, error) { - agentIndex := 0 + agentIdx := 0 agentFleet, err := pssh.NewFleet(cmd.AgentHosts, cmd.SSHUser, "", cmd.Stderr) if err != nil { return nil, err @@ -1560,10 +1560,12 @@ func (cmd *BspawnCommand) spawnRemote(ctx context.Context) (map[string]interface for _, sp := range cmd.Benchmarks { results[sp.Name] = make(map[int]interface{}) for i := 0; i < sp.Num; i++ { - sess, err := agentFleet[agentIndex].NewSession() + agentIdx %= len(cmd.AgentHosts) + sess, err := agentFleet[cmd.AgentHosts[agentIdx]].NewSession() if err != nil { return nil, err } + agentIdx += 1 sessions = append(sessions, sess) stdout, err := sess.StdoutPipe() if err != nil { diff --git a/creator/remote.go b/creator/remote.go index 00a9015b9..b3fa56eba 100644 --- a/creator/remote.go +++ b/creator/remote.go @@ -44,11 +44,11 @@ func (c *RemoteCluster) Start() error { return fmt.Errorf("no type or hosts specified - cannot continue") } + fleet, err := pssh.NewFleet(c.ClusterHosts, c.SSHUser, c.Keyfile, c.Stderr) + if err != nil { + return fmt.Errorf("connecting to cluster hosts: %v", err) + } if c.CopyBinary { - fleet, err := pssh.NewFleet(c.ClusterHosts, c.SSHUser, c.Keyfile, c.Stderr) - if err != nil { - return fmt.Errorf("copying binary: %v", err) - } pkg := "github.com/pilosa/pilosa/cmd/pilosa" bin, err := build.Binary(pkg, c.GOOS, c.GOARCH) @@ -79,8 +79,8 @@ func (c *RemoteCluster) Start() error { conf.Host = hostport conf.DataDir = "~/.pilosa" + port - // Connect to remote host - client, err := pssh.NewClient(host, c.SSHUser, "", c.Stderr) + // Get client for host + client, err := fleet.Get(host) if err != nil { return fmt.Errorf("connecting to host: %v", err) } diff --git a/ssh/ssh.go b/ssh/ssh.go index d310e66b1..20a1bea41 100644 --- a/ssh/ssh.go +++ b/ssh/ssh.go @@ -68,17 +68,17 @@ func NewClient(host, username, keyfile string, stderr io.Writer) (*Client, error return &Client{client: client, Stderr: stderr}, nil } -type Fleet []*Client +type Fleet map[string]*Client func NewFleet(hosts []string, username, keyfile string, stderr io.Writer) (Fleet, error) { hosts = DedupHosts(hosts) - clients := make([]*Client, len(hosts)) - for i, host := range hosts { + clients := make(map[string]*Client, len(hosts)) + for _, host := range hosts { client, err := NewClient(host, username, keyfile, stderr) if err != nil { return nil, err } - clients[i] = client + clients[host] = client } return clients, nil } @@ -212,3 +212,14 @@ func (sf Fleet) WriteFile(name, perm string, data io.Reader) error { } return wc.Close() } + +func (sf Fleet) Get(host string) (*Client, error) { + colonIdx := strings.Index(host, ":") + if colonIdx != -1 { + host = host[:colonIdx] + } + if client, ok := sf[host]; ok { + return client, nil + } + return nil, fmt.Errorf("No client found in fleet: %v for host: %v", sf, host) +}