diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 119357d95..0853b8172 100644 --- a/cmd/pilosactl/main.go +++ b/cmd/pilosactl/main.go @@ -12,8 +12,8 @@ import ( "io/ioutil" "log" "math/rand" - "net" "os" + "os/signal" "path/filepath" "sort" "strconv" @@ -24,7 +24,6 @@ import ( "time" "unsafe" - "github.com/BurntSushi/toml" "github.com/pilosa/pilosa" "github.com/pilosa/pilosa/bench" "github.com/pilosa/pilosa/creator" @@ -1064,120 +1063,78 @@ type CreateOutput struct { // Run executes cluster creation. func (cmd *CreateCommand) Run(ctx context.Context) error { var clus creator.Cluster - output := &CreateOutput{} switch cmd.Type { case "local": - var err error - clus, err = creator.NewLocalCluster(cmd.ReplicaN, cmd.ServerN) - if err != nil { - return fmt.Errorf("running create command: %v", err) + clus = &creator.LocalCluster{ + ReplicaN: cmd.ReplicaN, + ServerN: cmd.ServerN, } - defer clus.Shutdown() - output.Hosts = clus.Hosts() - - logReaders := clus.Logs() - if cmd.LogFilePrefix != "" { - output.LogFiles = make([]string, len(clus.Hosts())) - } - for i, _ := range clus.Hosts() { - var f io.Writer = cmd.Stderr - var err error - if cmd.LogFilePrefix != "" { - f, err = os.Create(cmd.LogFilePrefix + strconv.Itoa(i)) - if err != nil { - return err - } - output.LogFiles[i] = f.(*os.File).Name() - } - - go func(i int, f io.Writer) { - _, err := io.Copy(f, logReaders[i]) - if err != nil { - fmt.Fprintf(cmd.Stderr, "Error copying cluster logs: '%v'", err) - } - }(i, f) - } - - enc := json.NewEncoder(cmd.Stdout) - err = enc.Encode(output) - if err != nil { - return err - } - select {} case "AWS": return fmt.Errorf("AWS cluster type is not yet implemented") case "": - if len(cmd.Hosts) == 0 { - return fmt.Errorf("no type or hosts specified - cannot continue") + clus = &creator.RemoteCluster{ + ClusterHosts: cmd.Hosts, + ReplicaN: cmd.ReplicaN, + SSHUser: cmd.SSHUser, + Stderr: cmd.Stderr, } - // TODO: build pilosa - // TODO: copy binary to hosts - // build config - conf := pilosa.NewConfigForHosts(cmd.Hosts) - conf.Cluster.ReplicaN = cmd.ReplicaN - - // copy config to remote hosts and start pilosa - waitall := &sync.WaitGroup{} - for _, hostport := range cmd.Hosts { - host, port, err := net.SplitHostPort(hostport) - if err != nil { - return err - } - conf.Host = hostport - conf.DataDir = "~/.pilosa" + port - - client, err := pilosactl.NewSSH(host, cmd.SSHUser, "") - if err != nil { - return err - } - sess, err := client.NewSession() - if err != nil { - return err - } - configname := "pilosa" + port + ".conf" - w, err := sess.StdinPipe() - err = sess.Start("cat > " + configname) - if err != nil { - return err - } - enc := toml.NewEncoder(w) - err = enc.Encode(conf) - if err != nil { - return fmt.Errorf("encoding config: %v", err) - } - err = w.Close() - if err != nil { - return err - } - err = sess.Wait() - if err != nil { - return err - } - - sess, err = client.NewSession() - if err != nil { - return err - } - sess.Stdout = cmd.Stdout - sess.Stderr = cmd.Stderr - err = sess.Start("pilosa -config " + configname) - if err != nil { - return err - } - waitall.Add(1) - go func() { - defer waitall.Done() - err = sess.Wait() - if err != nil { - fmt.Fprintf(cmd.Stderr, "problem with remote pilosa process: %v", err) - } - }() - } - waitall.Wait() - return nil default: return fmt.Errorf("Unknown cluster type %v", cmd.Type) } + + err := clus.Start() + if err != nil { + return fmt.Errorf("starting cluster: %v", err) + } + + c := make(chan os.Signal, 1) + signal.Notify(c, os.Interrupt) + go func() { + for range c { + fmt.Fprintf(cmd.Stderr, "caught signal\n") + err := clus.Shutdown() + code := 0 + if err != nil { + code = 1 + } + os.Exit(code) + } + }() + + defer clus.Shutdown() + output := &CreateOutput{} + output.Hosts = clus.Hosts() + + logReaders := clus.Logs() + if cmd.LogFilePrefix != "" { + output.LogFiles = make([]string, len(clus.Hosts())) + } + for i, _ := range clus.Hosts() { + var f io.Writer = cmd.Stderr + var err error + if cmd.LogFilePrefix != "" { + f, err = os.Create(cmd.LogFilePrefix + strconv.Itoa(i)) + if err != nil { + return err + } + output.LogFiles[i] = f.(*os.File).Name() + } + + go func(i int, f io.Writer) { + _, err := io.Copy(f, logReaders[i]) + if err != nil { + fmt.Fprintf(cmd.Stderr, "Error copying cluster logs: '%v'", err) + } + }(i, f) + } + + enc := json.NewEncoder(cmd.Stdout) + err = enc.Encode(output) + if err != nil { + return err + } + select {} + } // BagentCommand represents a command for running a benchmark agent. A benchmark diff --git a/creator/creator.go b/creator/creator.go index a6b0f54f4..a7840dad2 100644 --- a/creator/creator.go +++ b/creator/creator.go @@ -1,100 +1,11 @@ // creator contains code for standing up pilosa clusters package creator -import ( - "fmt" - "io" - "io/ioutil" - "os" - "path/filepath" - "strconv" - - "github.com/pilosa/pilosa" -) +import "io" type Cluster interface { + Start() error Hosts() []string Shutdown() error Logs() []io.Reader } - -type localcluster struct { - hosts []string - logs []io.Reader - servers []*pilosa.Server - cluster *pilosa.Cluster - path string -} - -func NewLocalCluster(replicaN, serverN int) (Cluster, error) { - BasePort := 19327 - - localCluster := &localcluster{ - hosts: make([]string, serverN), - servers: make([]*pilosa.Server, serverN), - logs: make([]io.Reader, serverN), - } - path, err := ioutil.TempDir("", "pilosa-bench-") - if err != nil { - return localCluster, err - } - localCluster.path = path - - // Build cluster configuration. - cluster := pilosa.NewCluster() - cluster.ReplicaN = replicaN - - for i := 0; i < serverN; i++ { - cluster.Nodes = append(cluster.Nodes, &pilosa.Node{ - Host: fmt.Sprintf("localhost:%d", BasePort+i), - }) - } - localCluster.cluster = cluster - - // Build servers. - for i := range localCluster.servers { - // Make server work directory. - if err := os.MkdirAll(filepath.Join(path, strconv.Itoa(i)), 0777); err != nil { - return localCluster, err - } - - // Build server. - s := pilosa.NewServer() - s.Host = fmt.Sprintf("localhost:%d", BasePort+i) - s.Cluster = cluster - s.Index.Path = filepath.Join(path, strconv.Itoa(i), "data") - - // Create log stream - localCluster.logs[i], s.LogOutput = io.Pipe() - - localCluster.servers[i] = s - } - - // Open all servers. - for i, s := range localCluster.servers { - if err := s.Open(); err != nil { - return localCluster, err - } - localCluster.hosts[i] = s.Host - } - - return localCluster, nil -} - -func (c *localcluster) Hosts() []string { return c.hosts } -func (c *localcluster) Logs() []io.Reader { return c.logs } -func (c *localcluster) Shutdown() error { - errs := "" - for _, s := range c.servers { - if err := s.Close(); err != nil { - errs = errs + err.Error() + "; " - } - } - if err := os.RemoveAll(c.path); err != nil { - errs = errs + err.Error() + ";" - } - if errs != "" { - return fmt.Errorf(errs) - } - return nil -} diff --git a/creator/local.go b/creator/local.go new file mode 100644 index 000000000..4d4b6b8ef --- /dev/null +++ b/creator/local.go @@ -0,0 +1,94 @@ +package creator + +import ( + "fmt" + "io" + "io/ioutil" + "os" + "path/filepath" + "strconv" + + "github.com/pilosa/pilosa" +) + +type LocalCluster struct { + ReplicaN int + ServerN int + hosts []string + logs []io.Reader + servers []*pilosa.Server + cluster *pilosa.Cluster + path string +} + +func (localCluster *LocalCluster) Start() error { + BasePort := 19327 + + localCluster.hosts = make([]string, localCluster.ServerN) + localCluster.servers = make([]*pilosa.Server, localCluster.ServerN) + localCluster.logs = make([]io.Reader, localCluster.ServerN) + + path, err := ioutil.TempDir("", "pilosa-bench-") + if err != nil { + return err + } + localCluster.path = path + + // Build cluster configuration. + cluster := pilosa.NewCluster() + cluster.ReplicaN = localCluster.ReplicaN + + for i := 0; i < localCluster.ServerN; i++ { + cluster.Nodes = append(cluster.Nodes, &pilosa.Node{ + Host: fmt.Sprintf("localhost:%d", BasePort+i), + }) + } + localCluster.cluster = cluster + + // Build servers. + for i := range localCluster.servers { + // Make server work directory. + if err := os.MkdirAll(filepath.Join(path, strconv.Itoa(i)), 0777); err != nil { + return err + } + + // Build server. + s := pilosa.NewServer() + s.Host = fmt.Sprintf("localhost:%d", BasePort+i) + s.Cluster = cluster + s.Index.Path = filepath.Join(path, strconv.Itoa(i), "data") + + // Create log stream + localCluster.logs[i], s.LogOutput = io.Pipe() + + localCluster.servers[i] = s + } + + // Open all servers. + for i, s := range localCluster.servers { + if err := s.Open(); err != nil { + return err + } + localCluster.hosts[i] = s.Host + } + + return nil +} + +func (c *LocalCluster) Hosts() []string { return c.hosts } +func (c *LocalCluster) Logs() []io.Reader { return c.logs } +func (c *LocalCluster) Shutdown() error { + errs := "" + for _, s := range c.servers { + if err := s.Close(); err != nil { + errs = errs + err.Error() + "; " + } + } + if err := os.RemoveAll(c.path); err != nil { + errs = errs + err.Error() + ";" + } + if errs != "" { + return fmt.Errorf(errs) + } + return nil +} diff --git a/creator/remote.go b/creator/remote.go new file mode 100644 index 000000000..84cf9fa89 --- /dev/null +++ b/creator/remote.go @@ -0,0 +1,162 @@ +package creator + +import ( + "fmt" + "io" + "net" + "sync" + + "time" + + "github.com/BurntSushi/toml" + "github.com/pilosa/pilosa" + "github.com/pilosa/pilosa/pilosactl" + "golang.org/x/crypto/ssh" +) + +type RemoteCluster struct { + ClusterHosts []string + ReplicaN int + SSHUser string + Keyfile string + Key []byte + Stderr io.Writer + wg *sync.WaitGroup + logs []io.Reader + sessions []*ssh.Session + pipeRs []*io.PipeReader + pipeWs []*io.PipeWriter + stdins []io.WriteCloser +} + +func (c *RemoteCluster) Start() error { + c.logs = make([]io.Reader, 0) + if len(c.ClusterHosts) == 0 { + return fmt.Errorf("no type or hosts specified - cannot continue") + } + // TODO: build pilosa + // TODO: copy binary to hosts + // build config + conf := pilosa.NewConfigForHosts(c.ClusterHosts) + conf.Cluster.ReplicaN = c.ReplicaN + + // copy config to remote hosts and start pilosa + c.wg = &sync.WaitGroup{} + for _, hostport := range c.ClusterHosts { + + // Set up config for this host + host, port, err := net.SplitHostPort(hostport) + if err != nil { + return err + } + conf.Host = hostport + 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() + if err != nil { + return err + } + configname := "pilosa" + port + ".conf" + w, err := sess.StdinPipe() + err = sess.Start("cat > " + configname) + if err != nil { + return err + } + enc := toml.NewEncoder(w) + err = enc.Encode(conf) + if err != nil { + return fmt.Errorf("encoding config: %v", err) + } + err = w.Close() + if err != nil { + return err + } + err = sess.Wait() + if err != nil { + return err + } + + // Start pilosa on remote host + sess, err = client.NewSession() + if err != nil { + return err + } + // Have to request pty in order to be able to kill remote process + // reliably. + modes := ssh.TerminalModes{ + ssh.ISIG: 1, + ssh.ECHO: 0, + } + err = sess.RequestPty("vt100", 40, 80, modes) + if err != nil { + return fmt.Errorf("request pty error: %v", err) + } + pipeR, pipeW := io.Pipe() + sess.Stdout = pipeW + sess.Stderr = pipeW + inpipe, err := sess.StdinPipe() + if err != nil { + return err + } + c.logs = append(c.logs, pipeR) + c.sessions = append(c.sessions, sess) + c.pipeRs = append(c.pipeRs, pipeR) + c.pipeWs = append(c.pipeWs, pipeW) + c.stdins = append(c.stdins, inpipe) + + err = sess.Start("pilosa -config " + configname) + if err != nil { + return err + } + c.wg.Add(1) + go func() { + defer c.wg.Done() + fmt.Fprintf(c.Stderr, "start waiting on %v\n", sess) + err = sess.Wait() + fmt.Fprintf(c.Stderr, "done waiting on session\n") + if err != nil { + fmt.Fprintf(c.Stderr, "problem with remote pilosa process: %v", err) + } + }() + } + return nil +} + +func (c *RemoteCluster) Hosts() []string { return c.ClusterHosts } +func (c *RemoteCluster) Logs() []io.Reader { return c.logs } +func (c *RemoteCluster) Shutdown() error { + for i, sess := range c.sessions { + var err error + _, err = c.stdins[i].Write([]byte{3}) // Send Control C + if err != nil { + fmt.Fprintf(c.Stderr, "Error write-signaling remote process: %v\n", err) + } + err = sess.Signal(ssh.SIGINT) + if err != nil { + fmt.Fprintf(c.Stderr, "Error signaling remote process: %v\n", err) + } + err = sess.Close() + if err != nil { + fmt.Fprintf(c.Stderr, "Error closing remote session: %v\n", err) + } + } + done := make(chan struct{}, 1) + go func() { + fmt.Fprintf(c.Stderr, "Waiting\n") + c.wg.Wait() + done <- struct{}{} + }() + select { + case <-done: + return nil + case <-time.After(time.Second * 5): + return fmt.Errorf("timed out waiting for remote processes to exit") + } +}