remove fixed 100 second wait and add colored log lines

The fixed wait time was replaced with a loop and retry at the two
milestones in droplet creation. This allows the program to continue
as soon as everything is ready.

go.mod

@@ -5,6 +5,7 @@ go 1.15
 require (
 	github.com/aerissecure/csrange v0.0.0-20190810202023-1552ad60d0d7
 	github.com/digitalocean/godo v1.42.1
+	github.com/fatih/color v1.9.0
 	github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff
 	golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a
 	golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d

go.sum

@@ -5,6 +5,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/digitalocean/godo v1.42.1 h1:SJ/XMVsp5CZmyQal8gLlOl9jSl1i3FaN20LlgtK5ZMs=
 github.com/digitalocean/godo v1.42.1/go.mod h1:p7dOjjtSBqCTUksqtA5Fd3uaKs9kyTq2xcz76ulEJRU=
+github.com/fatih/color v1.9.0 h1:8xPHl4/q1VyqGIPif1F+1V3Y3lSmrq01EabUW3CoW5s=
+github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.5 h1:F768QJ1E9tib+q5Sc8MkdJi1RxLTbRcTf8LJV56aRls=
@@ -13,6 +15,11 @@ github.com/google/go-querystring v1.0.0 h1:Xkwi/a1rcvNg1PPYe5vI8GbeBY/jrVuDX5ASu
 github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO6wN/zVPAxq5ck=
 github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff h1:6NvhExg4omUC9NfA+l4Oq3ibNNeJUdiAF3iBVB0PlDk=
 github.com/jmcvetta/randutil v0.0.0-20150817122601-2bb1b664bcff/go.mod h1:ddfPX8Z28YMjiqoaJhNBzWHapTHXejnB5cDCUWDwriw=
+github.com/mattn/go-colorable v0.1.4 h1:snbPLB8fVfU9iwbbo30TPtbLRzwWu6aJS6Xh4eaaviA=
+github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
+github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
+github.com/mattn/go-isatty v0.0.11 h1:FxPOTFNqGkuDUGi3H/qkUbQO4ZiBa2brKq5r0l8TGeM=
+github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -31,7 +38,9 @@ golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d h1:TzXSXBo42m9gQenoE3b9BG
 golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd h1:xhmwyvizuTgC2qz7ZlMluP20uW+C3Rm0FD/WLDX8884=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

machine.go

@@ -15,6 +15,7 @@ import (
 	"golang.org/x/crypto/ssh"
 
 	"github.com/digitalocean/godo"
+	"github.com/fatih/color"
 )
 
 // Machine is just a wrapper around a created droplet.
@@ -34,15 +35,28 @@ type Machine struct {
 	Ports     string
 }
 
-// Println uses fmt.Println to log to the console, formatted for this machine.
+// Println uses log.Println to log to the console, formatted for this machine.
 func (m *Machine) Println(a ...interface{}) {
-	a = append([]interface{}{fmt.Sprintf("%s (%d):", m.Name, m.Index)}, a...)
+	a = append([]interface{}{fmt.Sprintf("%s (%s):", m.Name, m.Index)}, a...)
 	log.Println(a...)
 }
 
-// Printf uses fmt.Printf to log to the console, formatted for this machine.
+// Printf uses log.Printf to log to the console, formatted for this machine.
 func (m *Machine) Printf(format string, a ...interface{}) {
-	log.Printf(fmt.Sprintf("%s (%d): %s", m.Name, m.Index, format), a...)
+	log.Printf(fmt.Sprintf("%s (%s): %s", m.Name, m.Index, format), a...)
+}
+
+// Colorln uses log.Println to log to the console, formatted for this machine.
+// The log line will be formatted with the specified color.
+func (m *Machine) Colorln(c *color.Color, a ...interface{}) {
+	a = append([]interface{}{fmt.Sprintf("%s (%s):", m.Name, m.Index)}, a...)
+	log.Print(c.Sprintln(a...))
+}
+
+// Colorf uses log.Printf to log to the console, formatted for this machine.
+// The log line will be formatted with the specified color.
+func (m *Machine) Colorf(c *color.Color, format string, a ...interface{}) {
+	log.Print(c.Sprintf(fmt.Sprintf("%s (%s): %s", m.Name, m.Index, format), a...))
 }
 
 // IsReady ensures that the machine has an IP address.
@@ -129,7 +143,7 @@ func (m *Machine) RunCommand(filename string) error {
 	return session.Wait()
 }
 
-// Command generates the command to be run on the remote hosting using the defined Template.
+// Command generates the command to be run on the remote host using the defined Template.
 func (m *Machine) Command() (string, error) {
 	vars := map[string]interface{}{
 		"ports": m.Ports,

main.go

@@ -18,11 +18,12 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/aerissecure/csrange"
-	"github.com/digitalocean/godo"
-
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/terminal"
+
+	"github.com/aerissecure/csrange"
+	"github.com/digitalocean/godo"
+	"github.com/fatih/color"
 )
 
 // TODO: there is a 10 host limit for a single droplet creation call. for larger, split into multiple calls.
@@ -40,6 +41,10 @@ var (
 	force       = flag.Bool("force", false, "Bypass built-in protections that prevent you from deploying more than 50 droplets")
 	showversion = flag.Bool("v", false, "Print version and exit")
 	version     = "1.0.0"
+
+	yellow = color.New(color.FgYellow)
+	green  = color.New(color.FgGreen)
+	red    = color.New(color.FgRed)
 )
 
 func main() {
@@ -80,6 +85,8 @@ func main() {
 	}
 	droplets := []godo.Droplet{}
 
+	log.Printf("creating %d droplets\n", *count)
+
 	for region, c := range regionCountMap {
 		log.Printf("Creating %d droplets in region %s", c, region)
 		drops, _, err := client.Droplets.CreateMultiple(context.Background(), newDropLetMultiCreateRequest(*name, region, keyID, c))
@@ -111,18 +118,13 @@ func main() {
 
 	log.Println("Please CTRL-C to destroy droplets")
 
-	log.Println("Waiting for droplets to finish provisioning")
-	continueTime := time.Now().Add(time.Second * 100).Round(time.Second)
-	for time.Now().Before(continueTime.Add(time.Second * -1)) {
-		log.Printf("Waiting %v seconds...\n", continueTime.Sub(time.Now().Round(time.Second)).Seconds())
-		time.Sleep(time.Second * 10)
-	}
+	log.Println("Provisioning droplets")
 
 	var readyMachines []*Machine
 	for i := range machines {
 		m := &machines[i]
 
-		m.Index = zeroPad(len(readyMachines), i+1)
+		m.Index = zeroPad(*count, i+1)
 		m.Template = *cmd
 
 		sshConfig := &ssh.ClientConfig{
@@ -135,18 +137,22 @@ func main() {
 
 		m.SSHConfig = sshConfig
 
-		if err := m.GetIPs(client); err != nil {
-			m.Printf("Error getting the IPv4 address of droplet: %v", err)
-		}
-
-		m.Printf("IPv4 Address: %s", m.IPv4)
-
-		if !m.IsReady() {
-			m.Println("Droplet not ready yet. Skipping...")
+		// get IPv4 address, waiting until available
+		for {
+			if err := m.GetIPs(client); err != nil {
+				m.Printf("Error getting the IPv4 address of droplet: %v", err)
+			}
+			if !m.IsReady() {
+				m.Colorln(yellow, "Droplet not ready yet, sleeping 5s")
+				time.Sleep(5 * time.Second)
+				continue
+			}
+			break
 		}
 
 		readyMachines = append(readyMachines, m)
-		m.Println("Droplet ready")
+		m.Printf("IPv4 Address: %s", m.IPv4)
+		m.Colorln(green, "Droplet ready")
 	}
 
 	var portBuckets []string
@@ -164,16 +170,22 @@ func main() {
 		wg.Add(1)
 		go func(idx int, m *Machine) {
 			m.Println("Establishing SSH connection...")
-			client, err := ssh.Dial("tcp", m.IPv4+":22", m.SSHConfig)
-			if err != nil {
-				m.Println("Error establishing SSH connection. Results may be incomplete. Please exit and try again")
-				wg.Done()
-				return
-			}
-
-			m.Println("SSH connection established.")
 
-			m.SSHClient = client
+			for {
+				client, err := ssh.Dial("tcp", m.IPv4+":22", m.SSHConfig)
+				if err != nil {
+					fmt.Println("ERROR:", err)
+					m.Colorln(yellow, "Error establishing SSH connection, retrying...")
+					time.Sleep(10 * time.Second)
+					// why is this sleep not working????
+					// wg.Done() // don't call it quits here, keep trying to connect
+					// return
+					continue
+				}
+				m.Colorln(green, "SSH connection established.")
+				m.SSHClient = client
+				break
+			}
 
 			if *pkg != "" {
 				m.Println("Installing packages")
@@ -198,7 +210,7 @@ func main() {
 			}
 
 			m.Printf("Results: %s\n", fname)
-			m.Println("Done. Exiting...")
+			m.Colorln(green, "Done. Exiting...")
 
 			wg.Done()
 		}(i, m)