Parallel vanity

nk/main.go

@@ -20,13 +20,15 @@ import (
 	"encoding/base64"
 	"flag"
 	"fmt"
+	"github.com/nats-io/nkeys"
 	"io"
 	"io/ioutil"
 	"log"
 	"os"
+	"runtime"
 	"strings"
-
-	"github.com/nats-io/nkeys"
+	"sync"
+	"time"
 )
 
 // this will be set during compilation when a release is made on tools
@@ -236,6 +238,7 @@ var b32Enc = base32.StdEncoding.WithPadding(base32.NoPadding)
 
 func createVanityKey(keyType, vanity, entropy string, max int) nkeys.KeyPair {
 	spinners := []rune(`⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏`)
+	steps := len(spinners)
 	pre := preForType(keyType)
 	vanity = strings.ToUpper(vanity)
 	// Check to make sure we can base32 into it by trying to decode it.
@@ -244,19 +247,80 @@ func createVanityKey(keyType, vanity, entropy string, max int) nkeys.KeyPair {
 		log.Fatalf("Can not generate base32 encoded strings to match '%s'", vanity)
 	}
 
-	for i := 0; i < max; i++ {
-		spin := spinners[i%len(spinners)]
-		fmt.Fprintf(os.Stderr, "\r\033[mcomputing\033[m %s ", string(spin))
-		kp := genKeyPair(pre, entropy)
-		pub, _ := kp.PublicKey()
-		if strings.HasPrefix(pub[1:], vanity) {
-			fmt.Fprintf(os.Stderr, "\r")
-			return kp
+	var mu sync.Mutex
+	var found nkeys.KeyPair
+	var done bool
+
+	var spinner sync.WaitGroup
+	go func() {
+		spinner.Add(1)
+		i := 0
+		for {
+			i++
+			if i > steps {
+				i = 0
+			}
+			spin := spinners[i%steps]
+			mu.Lock()
+			if done {
+				mu.Unlock()
+				break
+			}
+			fmt.Fprintf(os.Stderr, "\r\033[mcomputing\033[m %s ", string(spin))
+			mu.Unlock()
+			time.Sleep(50 * time.Millisecond)
+		}
+		fmt.Fprintf(os.Stderr, "\r")
+		spinner.Done()
+	}()
+
+	var workers sync.WaitGroup
+	cpu := runtime.NumCPU()
+	workers.Add(cpu)
+	for i := 0; i < cpu; i++ {
+		tries := max / cpu
+		extra := max % cpu
+		if extra > 0 && extra >= i-1 {
+			tries += 1
 		}
+
+		go func() {
+			for i := 0; i < tries; i++ {
+				// are we done
+				mu.Lock()
+				if found != nil {
+					// unlock fast path
+					mu.Unlock()
+					break
+				}
+				mu.Unlock()
+
+				kp := genKeyPair(pre, entropy)
+				pub, _ := kp.PublicKey()
+				if strings.HasPrefix(pub[1:], vanity) {
+					// lock again to signal the found key
+					mu.Lock()
+					found = kp
+					mu.Unlock()
+					break
+				}
+			}
+			workers.Done()
+		}()
+	}
+	workers.Wait()
+
+	// stop the spinner
+	mu.Lock()
+	done = true
+	mu.Unlock()
+	// wait for any console i/o
+	spinner.Wait()
+
+	if found == nil {
+		log.Fatalf("Failed to generate prefix after %d attempts", max)
 	}
-	fmt.Fprintf(os.Stderr, "\r")
-	log.Fatalf("Failed to generate prefix after %d attempts", max)
-	return nil
+	return found
 }
 
 func readKeyFile(filename string) []byte {