diff --git a/cache.go b/cache.go index 03e02d15f..2ff5f60ce 100644 --- a/cache.go +++ b/cache.go @@ -51,7 +51,8 @@ func (c *LRUCache) Add(bitmapID, n uint64) { // Get returns a bitmap with a given id. func (c *LRUCache) Get(bitmapID uint64) uint64 { n, _ := c.cache.Get(bitmapID) - return n.(uint64) + nn, _ := n.(uint64) + return nn } // Len returns the number of items in the cache. diff --git a/cmd/pilosa/main.go b/cmd/pilosa/main.go index fd26a5685..7107745fe 100644 --- a/cmd/pilosa/main.go +++ b/cmd/pilosa/main.go @@ -9,6 +9,7 @@ import ( "os" "os/signal" "path/filepath" + "runtime/pprof" "strings" "time" @@ -45,18 +46,38 @@ func main() { os.Exit(2) } + // Start CPU profiling. + if m.CPUProfile != "" { + f, err := os.Create(m.CPUProfile) + if err != nil { + fmt.Fprintf(m.Stderr, "create cpu profile: %v", err) + os.Exit(1) + } + defer f.Close() + + fmt.Fprintln(m.Stderr, "Starting cpu profile") + pprof.StartCPUProfile(f) + time.AfterFunc(m.CPUTime, func() { + fmt.Fprintln(m.Stderr, "Stopping cpu profile") + pprof.StopCPUProfile() + f.Close() + }) + } + // Execute the program. if err := m.Run(); err != nil { fmt.Fprintln(m.Stderr, err) + fmt.Fprintln(m.Stderr, "stopping profile") os.Exit(1) } // First SIGKILL causes server to shut down gracefully. - // Second signal causes a hard shutdown. c := make(chan os.Signal, 2) signal.Notify(c, os.Interrupt) sig := <-c fmt.Fprintf(m.Stderr, "Received %s; gracefully shutting down...\n", sig.String()) + + // Second signal causes a hard shutdown. go func() { <-c; os.Exit(1) }() if err := m.Close(); err != nil { @@ -73,6 +94,10 @@ type Main struct { ConfigPath string Config *Config + // Profiling options. + CPUProfile string + CPUTime time.Duration + // Standard input/output Stdin io.Reader Stdout io.Writer @@ -127,8 +152,10 @@ func (m *Main) Close() error { // ParseFlags parses command line flags from args. func (m *Main) ParseFlags(args []string) error { fs := flag.NewFlagSet("pilosa", flag.ContinueOnError) - fs.SetOutput(m.Stderr) + fs.StringVar(&m.CPUProfile, "cpuprofile", "", "cpu profile") + fs.DurationVar(&m.CPUTime, "cputime", 30*time.Second, "cpu profile duration") fs.StringVar(&m.ConfigPath, "config", "", "config path") + fs.SetOutput(m.Stderr) if err := fs.Parse(args); err != nil { return err } diff --git a/fragment.go b/fragment.go index 6268ab729..919f1d0a5 100644 --- a/fragment.go +++ b/fragment.go @@ -527,6 +527,16 @@ func (f *Fragment) topBitmapPairs(bitmapIDs []uint64) []BitmapPair { // Otherwise retrieve specific bitmaps. pairs := make([]BitmapPair, len(bitmapIDs)) for i, bitmapID := range bitmapIDs { + // Look up cache first, if available. + if n := f.cache.Get(bitmapID); n > 0 { + pairs[i] = BitmapPair{ + ID: bitmapID, + Count: n, + } + continue + } + + // Otherwise load from storage. pairs[i] = BitmapPair{ ID: bitmapID, Count: f.Bitmap(bitmapID).Count(),