Merge pull request #103 from benbjohnson/fix-count-cache

Fix count() cache retrieval

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.

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
 	}

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(),

handler.go

@@ -59,6 +59,8 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			pprof.Profile(w, r)
 		case "/debug/pprof/symbol":
 			pprof.Symbol(w, r)
+		case "/debug/pprof/trace":
+			pprof.Trace(w, r)
 		default:
 			pprof.Index(w, r)
 		}