// Sleep delays the execution of the remainder of the chain of commands. // // Params: // -duration (time.Duration): Time to sleep. // -message (string): The message to log when entering sleep. func Sleep(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { dur := p.Get("duration", 10*time.Millisecond).(time.Duration) msg := p.Get("messages", "Sleeping").(string) log.Info(c, msg) time.Sleep(dur) log.Info(c, "Woke up.") return true, nil }
// handleConn handles an individual client connection. // // It manages the connection, but passes channels on to `answer()`. func (s *server) handleConn(conn net.Conn, conf *ssh.ServerConfig) { defer conn.Close() log.Info(s.c, "Accepted connection.") _, chans, reqs, err := ssh.NewServerConn(conn, conf) if err != nil { // Handshake failure. log.Errf(s.c, "Failed handshake: %s (%v)", err, conn) return } // Discard global requests. We're only concerned with channels. safely.GoDo(s.c, func() { ssh.DiscardRequests(reqs) }) condata := sshConnection(conn) // Now we handle the channels. for incoming := range chans { log.Infof(s.c, "Channel type: %s\n", incoming.ChannelType()) if incoming.ChannelType() != "session" { incoming.Reject(ssh.UnknownChannelType, "Unknown channel type") } channel, req, err := incoming.Accept() if err != nil { // Should close request and move on. panic(err) } safely.GoDo(s.c, func() { s.answer(channel, req, condata) }) } conn.Close() }
// RunOnce runs the equivalent of `confd --onetime`. // // This may run the process repeatedly until either we time out (~20 minutes) or // the templates are successfully built. // // Importantly, this blocks until the run is complete. // // Params: // - node (string): The etcd node to use. (Only etcd is currently supported) // // Returns: // - The []bytes from stdout and stderr when running the program. // func RunOnce(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { node := p.Get("node", defaultEtcd).(string) dargs := []string{"-onetime", "-node", node, "-log-level", "error"} log.Info(c, "Building confd templates. This may take a moment.") limit := 1200 timeout := time.Second * 3 var lasterr error start := time.Now() for i := 0; i < limit; i++ { out, err := exec.Command("confd", dargs...).CombinedOutput() if err == nil { log.Infof(c, "Templates generated for %s on run %d", node, i) return out, nil } log.Debugf(c, "Recoverable error: %s", err) log.Debugf(c, "Output: %q", out) lasterr = err time.Sleep(timeout) log.Infof(c, "Re-trying template build. (Elapsed time: %d)", time.Now().Sub(start)/time.Second) } return nil, fmt.Errorf("Could not build confd templates before timeout. Last error: %s", lasterr) }
// Ping handles a simple test SSH exec. // // Returns the string PONG and exit status 0. // // Params: // - channel (ssh.Channel): The channel to respond on. // - request (*ssh.Request): The request. // func Ping(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { channel := p.Get("channel", nil).(ssh.Channel) req := p.Get("request", nil).(*ssh.Request) log.Info(c, "PING\n") if _, err := channel.Write([]byte("pong")); err != nil { log.Errf(c, "Failed to write to channel: %s", err) } sendExitStatus(0, channel) req.Reply(true, nil) return nil, nil }
// setJoinMode determines what mode to start the etcd server in. // // In discovery mode, this will use the discovery URL to join a new cluster. // In "existing" mode, this will join to an existing cluster directly. // // Params: // - client (etcd.Getter): initialized etcd client // - path (string): path to get. This will go through os.ExpandEnv(). // - desiredLen (string): The number of nodes to expect in etcd. This is // usually stored as a string. // // Returns: // string "existing" or "new" func setJoinMode(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { cli := p.Get("client", nil).(client.Client) dlen := p.Get("desiredLen", "3").(string) path := p.Get("path", "").(string) path = os.ExpandEnv(path) state := "existing" dint, err := strconv.Atoi(dlen) if err != nil { log.Warnf(c, "Expected integer length, got '%s'. Defaulting to 3", dlen) dint = 3 } // Ideally, this should look in the /deis/status directory in the discovery // service. That will indicate how many members have been online in the // last few hours. This is a good indicator of whether a cluster exists, // even if not all the hosts are healthy. res, err := etcd.SimpleGet(cli, path, true) if err != nil { // This means that the discovery server is brand new, and nothing // has written to it yet. So we're new. if strings.Contains(err.Error(), "Key not found") { return "new", nil } log.Errf(c, "Failed simple get of %s: %s", path, err) return state, err } if !res.Node.Dir { //return state, errors.New("Expected a directory node in discovery service") log.Info(c, "No status information found in discovery service. Assuming new.") state = "new" } else if len(res.Node.Nodes) < dint { log.Info(c, "Cluster has not reached consensus number. Assuming new.") state = "new" } os.Setenv("ETCD_INITIAL_CLUSTER_STATE", state) return state, nil }
// listen handles accepting and managing connections. However, since closer // is len(1), it will not block the sender. func (s *server) listen(l net.Listener, conf *ssh.ServerConfig, closer chan interface{}) error { cxt := s.c log.Info(cxt, "Accepting new connections.") defer l.Close() // FIXME: Since Accept blocks, closer may not be checked often enough. for { log.Info(cxt, "Checking closer.") if len(closer) > 0 { <-closer log.Info(cxt, "Shutting down SSHD listener.") return nil } conn, err := l.Accept() if err != nil { log.Warnf(cxt, "Error during Accept: %s", err) // We shouldn't kill the listener because of an error. return err } safely.GoDo(cxt, func() { s.handleConn(conn, conf) }) } }
// Wait waits for a sync.WaitGroup to finish. // // Params: // - wg (Waiter): The thing to wait for. // - msg (string): The message to print when done. If this is empty, nothing is sent. // - waiting (string): String to tell what we're waiting for. If empty, nothing is displayed. // - failures (int): The number of failures that occurred while waiting. // // Returns: // Nothing. func Wait(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { ok, missing := p.RequiresValue("wg") if !ok { return nil, &cookoo.FatalError{"Missing required fields: " + strings.Join(missing, ", ")} } wg := p.Get("wg", nil).(Waiter) msg := p.Get("msg", "").(string) fails := p.Get("failures", 0).(int) waitmsg := p.Get("waiting", "").(string) if len(waitmsg) > 0 { log.Info(c, waitmsg) } wg.Wait() if len(msg) > 0 { log.Info(c, msg) } if fails > 0 { return nil, fmt.Errorf("There were %d failures while waiting.", fails) } return nil, nil }
// KillOnExit kills PIDs when the program exits. // // Otherwise, this blocks until an os.Interrupt or os.Kill is received. // // Params: // This treats Params as a map of process names (unimportant) to PIDs. It then // attempts to kill all of the pids that it receives. func KillOnExit(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) { sigs := make(chan os.Signal, 1) signal.Notify(sigs, os.Interrupt, os.Kill) safely.GoDo(c, func() { log.Info(c, "Builder is running.") <-sigs c.Log("info", "Builder received signal to stop.") pids := p.AsMap() killed := 0 for name, pid := range pids { if pid, ok := pid.(int); ok { if proc, err := os.FindProcess(pid); err == nil { log.Infof(c, "Killing %s (pid=%d)", name, pid) proc.Kill() killed++ } } } }) return nil, nil }