diff --git a/cmd/pilosa-bench/main.go b/cmd/pilosa-bench/main.go new file mode 100644 index 000000000..c059576bd --- /dev/null +++ b/cmd/pilosa-bench/main.go @@ -0,0 +1,213 @@ +package main + +import ( + "errors" + "flag" + "fmt" + "io" + "io/ioutil" + "log" + "os" + "path/filepath" + "strconv" + + "github.com/umbel/pilosa" +) + +// BasePort is the initial port used when generating a cluster. +const BasePort = 16000 + +// Default settings for command line arguments. +const ( + DefaultServerN = 1 + DefaultReplicaN = 1 +) + +func main() { + m := NewMain() + + // Parse command line arguments. + if err := m.ParseFlags(os.Args[1:]); err == flag.ErrHelp { + fmt.Fprintln(m.Stderr, m.Usage()) + os.Exit(1) + } else if err != nil { + fmt.Fprintln(m.Stderr, err) + os.Exit(1) + } + + // Execute the program. + if err := m.Run(); err != nil { + fmt.Fprintln(m.Stderr, err) + os.Exit(1) + } +} + +// Main represents the main program execution. +type Main struct { + ServerN int + ReplicaN int + + Verbose bool + Work bool + + // Standard input/output + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer +} + +// NewMain returns a new instance of Main. +func NewMain() *Main { + return &Main{ + ServerN: DefaultServerN, + ReplicaN: DefaultReplicaN, + + Stdin: os.Stdin, + Stdout: os.Stdout, + Stderr: os.Stderr, + } +} + +// Usage returns usage documentation. +func (m *Main) Usage() string { + return ` +pilosa-bench is a tool for testing and benchmarking pilosa clusters. + +Usage: + + pilosa-bench [arguments] + +The following arguments are available: + + -servers N + The number of servers to generate for the cluster. + Defaults to 1 server. + + -replicas N + Replication factor for data in the cluster. + Defaults to 1 replica. + + -v + Logs all server output to stderr in addition to server log. + + -work + Prints the temporary work directory and does not delete it + after the benchmark has completed. + +`[1:] +} + +// ParseFlags parses command line flags from args. +func (m *Main) ParseFlags(args []string) error { + fs := flag.NewFlagSet("pilosa-bench", flag.ContinueOnError) + fs.SetOutput(ioutil.Discard) + fs.IntVar(&m.ServerN, "servers", DefaultServerN, "server count") + fs.IntVar(&m.ReplicaN, "replicas", DefaultReplicaN, "replication factor") + fs.BoolVar(&m.Verbose, "v", false, "verbose") + fs.BoolVar(&m.Work, "work", false, "verbose") + if err := fs.Parse(args); err != nil { + return err + } + return nil +} + +// Run executes the main program execution. +func (m *Main) Run(args ...string) error { + // Validate arguments. + if m.ServerN <= 0 { + return errors.New("server count must be at least 1") + } else if m.ReplicaN <= 0 { + return errors.New("replication factor must be at least 1") + } else if m.ReplicaN > m.ServerN { + return errors.New("replication factor must be less than server count") + } + + logger := m.logger() + var logFiles []*os.File + + // Create a base temp path. + path, err := ioutil.TempDir("", "pilosa-bench-") + if err != nil { + return err + } + if m.Work { + fmt.Printf("WORK=%s\n\n", path) + } + + // Build cluster configuration. + cluster := pilosa.Cluster{ + ReplicaN: m.ReplicaN, + } + for i := 0; i < m.ServerN; i++ { + cluster.Nodes = append(cluster.Nodes, &pilosa.Node{ + Host: fmt.Sprintf("localhost:%d", BasePort+i), + }) + } + + // Build servers. + servers := make([]*pilosa.Server, m.ServerN) + for i := range 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 file. + f, err := os.Create(filepath.Join(path, strconv.Itoa(i), "log")) + if err != nil { + return err + } + logFiles = append(logFiles, f) + + // Set log and optionally write out to stderr as well. + s.LogOutput = f + if m.Verbose { + s.LogOutput = io.MultiWriter(m.Stderr, s.LogOutput) + } + + servers[i] = s + } + + // Open all servers. + for i, s := range servers { + logger.Printf("starting server #%d: %s", i, s.Host) + if err := s.Open(); err != nil { + return err + } + } + + // FIXME(benbjohnson): Execute benchmark testing. + + // Close all servers. + for i, s := range servers { + logger.Printf("closing server #%d", i) + if err := s.Close(); err != nil { + logger.Printf("error closing server: %s", err) + } + } + + // Close all logs. + for _, f := range logFiles { + f.Close() + } + + // If work flag is not set then delete all data & logs. + if !m.Work { + if err := os.RemoveAll(path); err != nil { + return err + } + } + + return nil +} + +// Close gracefully closes the program. +func (m *Main) Close() error { return nil } + +func (m *Main) logger() *log.Logger { return log.New(m.Stderr, "", log.LstdFlags) } diff --git a/cmd/pilosa-bench/main_test.go b/cmd/pilosa-bench/main_test.go new file mode 100644 index 000000000..4cc7d67c1 --- /dev/null +++ b/cmd/pilosa-bench/main_test.go @@ -0,0 +1,74 @@ +package main_test + +import ( + "bytes" + "io" + "os" + "testing" + + main "github.com/umbel/pilosa/cmd/pilosa-bench" +) + +// Ensure server count flag can be parsed. +func TestMain_ParseFlags_ServerN(t *testing.T) { + m := NewMain() + if err := m.ParseFlags([]string{"-servers", "2"}); err != nil { + t.Fatal(err) + } else if m.ServerN != 2 { + t.Fatalf("unexpected server count: %d", m.ServerN) + } +} + +// Ensure replication factor flag can be parsed. +func TestMain_ParseFlags_ReplicaN(t *testing.T) { + m := NewMain() + if err := m.ParseFlags([]string{"-servers=4", "-replicas", "3"}); err != nil { + t.Fatal(err) + } else if m.ReplicaN != 3 { + t.Fatalf("unexpected replica count: %d", m.ReplicaN) + } +} + +// Ensure verbose flag can be parsed. +func TestMain_ParseFlags_Verbose(t *testing.T) { + m := NewMain() + if err := m.ParseFlags([]string{"-v"}); err != nil { + t.Fatal(err) + } else if !m.Verbose { + t.Fatalf("expected verbose flag") + } +} + +// Ensure work flag can be parsed. +func TestMain_ParseFlags_Work(t *testing.T) { + m := NewMain() + if err := m.ParseFlags([]string{"-work"}); err != nil { + t.Fatal(err) + } else if !m.Work { + t.Fatalf("expected work flag") + } +} + +// Main represents a test wrapper for main.Main. +type Main struct { + *main.Main + + Stdin bytes.Buffer + Stdout bytes.Buffer + Stderr bytes.Buffer +} + +// NewMain returns a new instance of Main. +func NewMain() *Main { + m := &Main{Main: main.NewMain()} + m.Main.Stdin = &m.Stdin + m.Main.Stdout = &m.Stdout + m.Main.Stderr = &m.Stderr + + if testing.Verbose() { + m.Main.Stdout = io.MultiWriter(os.Stdout, m.Main.Stdout) + m.Main.Stderr = io.MultiWriter(os.Stderr, m.Main.Stderr) + } + + return m +}