diff --git a/bitmap.go b/bitmap.go index 9e14195ca..6d8c7f634 100644 --- a/bitmap.go +++ b/bitmap.go @@ -28,10 +28,12 @@ type Bitmap struct { } // NewBitmap returns a new instance of Bitmap. -func NewBitmap() *Bitmap { - return &Bitmap{ - tree: rbtree.NewTree(rbtreeItemCompare), +func NewBitmap(bits ...uint64) *Bitmap { + bm := &Bitmap{tree: rbtree.NewTree(rbtreeItemCompare)} + for _, i := range bits { + bm.setBit(i) } + return bm } // Chunk returns the chunk within the bitmap. @@ -357,8 +359,8 @@ func (b *Bitmap) Bits() []uint64 { return result } -// SetBit sets the i-th bit of the bitmap. -func (b *Bitmap) SetBit(i uint64) (changed bool) { +// setBit sets the i-th bit of the bitmap. +func (b *Bitmap) setBit(i uint64) (changed bool) { address := deref(i) chunk := b.Chunk(&Chunk{address.ChunkKey, make(Blocks, 32)}) @@ -375,8 +377,8 @@ func (b *Bitmap) SetBit(i uint64) (changed bool) { return changed } -// ClearBit clears the i-th bit of the bitmap. -func (b *Bitmap) ClearBit(i uint64) (changed bool) { +// clearBit clears the i-th bit of the bitmap. +func (b *Bitmap) clearBit(i uint64) (changed bool) { address := deref(i) chunk := b.Chunk(&Chunk{address.ChunkKey, make(Blocks, 32)}) diff --git a/cache.go b/cache.go index 7040b124c..5a5316469 100644 --- a/cache.go +++ b/cache.go @@ -17,7 +17,7 @@ type Cache interface { io.ReaderFrom Add(bitmapID, category uint64, bm *Bitmap) - Get(bitmapID uint64) (bm *Bitmap, ok bool) + Get(bitmapID uint64) *Bitmap Len() int // Updates the cache, if necessary. @@ -50,12 +50,12 @@ func (c *LRUCache) Add(bitmapID, category uint64, bm *Bitmap) { } // Get returns a bitmap with a given id. -func (c *LRUCache) Get(bitmapID uint64) (bm *Bitmap, ok bool) { - value, ok := c.cache.Get(bitmapID) +func (c *LRUCache) Get(bitmapID uint64) *Bitmap { + bm, ok := c.cache.Get(bitmapID) if !ok { - return nil, false + return nil } - return value.(*Bitmap), true + return bm.(*Bitmap) } // Len returns the number of items in the cache. @@ -68,11 +68,9 @@ func (c *LRUCache) Invalidate() {} func (c *LRUCache) Pairs() []Pair { a := make([]Pair, 0, len(c.keys)) for k := range c.keys { - bm, _ := c.Get(k) - a = append(a, Pair{ Key: k, - Count: bm.Count(), + Count: c.Get(k).Count(), }) } return a @@ -156,12 +154,12 @@ func (c *RankCache) Add(bitmapID, category uint64, bm *Bitmap) { } // Get returns a bitmap with a given id. -func (c *RankCache) Get(bitmapID uint64) (bm *Bitmap, ok bool) { +func (c *RankCache) Get(bitmapID uint64) *Bitmap { entry, ok := c.entries[bitmapID] if !ok { - return nil, false + return nil } - return entry.bitmap, true + return entry.bitmap } // Len returns the number of items in the cache. diff --git a/executor.go b/executor.go index f80e50549..4ab485f0b 100644 --- a/executor.go +++ b/executor.go @@ -160,7 +160,10 @@ func (e *Executor) executeGetSlice(db string, c *pql.Get, slice uint64) (*Bitmap frame = DefaultFrame } - f := e.Index().Fragment(db, frame, slice) + f, err := e.Index().Fragment(db, frame, slice) + if err != nil { + return nil, fmt.Errorf("fragment: %s", err) + } return f.Bitmap(c.ID), nil } @@ -240,8 +243,11 @@ func (e *Executor) executeSet(db string, c *pql.Set) error { for _, node := range e.Cluster.SliceNodes(slice) { // Update locally if host matches. if node.Host == e.Host { - f := e.Index().Fragment(db, c.Frame, slice) - f.Bitmap(c.ID).SetBit(c.ProfileID) + f, err := e.Index().Fragment(db, c.Frame, slice) + if err != nil { + return fmt.Errorf("fragment: %s", err) + } + f.SetBit(c.ID, c.ProfileID) continue } diff --git a/executor_test.go b/executor_test.go index 33d993e18..817434795 100644 --- a/executor_test.go +++ b/executor_test.go @@ -12,10 +12,12 @@ import ( // Ensure a get query can be executed. func TestExecutor_Execute_Get(t *testing.T) { - e := NewExecutor(NewCluster(1)) - e.Index().Fragment("d", "f", 0).Bitmap(10).SetBit(3) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit(SliceWidth + 1) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "f", 0).MustSetBit(10, 3) + idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+1) + e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`get(id=10, frame=f)`), nil); err != nil { t.Fatal(err) } else if chunks := res.(*pilosa.Bitmap).Chunks(); len(chunks) != 2 { @@ -29,13 +31,14 @@ func TestExecutor_Execute_Get(t *testing.T) { // Ensure a difference query can be executed. func TestExecutor_Execute_Difference(t *testing.T) { - e := NewExecutor(NewCluster(1)) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(1) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(2) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(3) - - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(2) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "general", 0).MustSetBit(10, 1) + idx.MustFragment("d", "general", 0).MustSetBit(10, 2) + idx.MustFragment("d", "general", 0).MustSetBit(10, 3) + idx.MustFragment("d", "general", 0).MustSetBit(11, 2) + e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`difference(get(id=10), get(id=11))`), nil); err != nil { t.Fatal(err) } else if chunks := res.(*pilosa.Bitmap).Chunks(); len(chunks) != 1 { @@ -47,15 +50,17 @@ func TestExecutor_Execute_Difference(t *testing.T) { // Ensure an intersect query can be executed. func TestExecutor_Execute_Intersect(t *testing.T) { - e := NewExecutor(NewCluster(1)) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(1) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(SliceWidth + 1) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(SliceWidth + 2) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "general", 0).MustSetBit(10, 1) + idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+1) + idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+2) - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(1) - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(2) - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(SliceWidth + 2) + idx.MustFragment("d", "general", 0).MustSetBit(11, 1) + idx.MustFragment("d", "general", 0).MustSetBit(11, 2) + idx.MustFragment("d", "general", 1).MustSetBit(11, SliceWidth+2) + e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`intersect(get(id=10), get(id=11))`), nil); err != nil { t.Fatal(err) } else if chunks := res.(*pilosa.Bitmap).Chunks(); len(chunks) != 2 { @@ -69,14 +74,16 @@ func TestExecutor_Execute_Intersect(t *testing.T) { // Ensure a union query can be executed. func TestExecutor_Execute_Union(t *testing.T) { - e := NewExecutor(NewCluster(1)) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(0) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(SliceWidth + 1) - e.Index().Fragment("d", "general", 0).Bitmap(10).SetBit(SliceWidth + 2) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "general", 0).MustSetBit(10, 0) + idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+1) + idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+2) - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(2) - e.Index().Fragment("d", "general", 0).Bitmap(11).SetBit(SliceWidth + 2) + idx.MustFragment("d", "general", 0).MustSetBit(11, 2) + idx.MustFragment("d", "general", 1).MustSetBit(11, SliceWidth+2) + e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`union(get(id=10), get(id=11))`), nil); err != nil { t.Fatal(err) } else if chunks := res.(*pilosa.Bitmap).Chunks(); len(chunks) != 2 { @@ -90,11 +97,13 @@ func TestExecutor_Execute_Union(t *testing.T) { // Ensure a count query can be executed. func TestExecutor_Execute_Count(t *testing.T) { - e := NewExecutor(NewCluster(1)) - e.Index().Fragment("d", "f", 0).Bitmap(10).SetBit(3) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit(SliceWidth + 1) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit(SliceWidth + 2) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "f", 0).MustSetBit(10, 3) + idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+1) + idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+2) + e := NewExecutor(idx.Index, NewCluster(1)) if n, err := e.Execute("d", MustParse(`count(get(id=10, frame=f))`), nil); err != nil { t.Fatal(err) } else if n != uint64(3) { @@ -104,13 +113,15 @@ func TestExecutor_Execute_Count(t *testing.T) { // Ensure a set query can be executed. func TestExecutor_Execute_Set(t *testing.T) { - e := NewExecutor(NewCluster(1)) + idx := MustOpenIndex() + defer idx.Close() + e := NewExecutor(idx.Index, NewCluster(1)) if _, err := e.Execute("d", MustParse(`set(id=10, frame=f, profile_id=1)`), nil); err != nil { t.Fatal(err) } - f := e.Index().Fragment("d", "f", 0) + f := idx.MustFragment("d", "f", 0) if n := f.Bitmap(10).Count(); n != 1 { t.Fatalf("unexpected bitmap count: %d", n) } @@ -136,18 +147,21 @@ func TestExecutor_Execute_Remote_Bitmap(t *testing.T) { } // Set bits in slice 0 & 2. - bm := pilosa.NewBitmap() - bm.SetBit((0 * SliceWidth) + 1) - bm.SetBit((0 * SliceWidth) + 2) - bm.SetBit((2 * SliceWidth) + 4) + bm := pilosa.NewBitmap( + (0*SliceWidth)+1, + (0*SliceWidth)+2, + (2*SliceWidth)+4, + ) return bm, nil } // Create local executor data. // The local node owns slice 1. - e := NewExecutor(c) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit((1 * SliceWidth) + 1) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) + e := NewExecutor(idx.Index, c) if res, err := e.Execute("d", MustParse(`get(id=10, frame=f)`), nil); err != nil { t.Fatal(err) } else if chunks := res.(*pilosa.Bitmap).Chunks(); len(chunks) != 3 { @@ -174,10 +188,12 @@ func TestExecutor_Execute_Remote_Count(t *testing.T) { } // Create local executor data. The local node owns slice 1. - e := NewExecutor(c) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit((1 * SliceWidth) + 1) - e.Index().Fragment("d", "f", 1).Bitmap(10).SetBit((1 * SliceWidth) + 2) + idx := MustOpenIndex() + defer idx.Close() + idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) + idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+2) + e := NewExecutor(idx.Index, c) if n, err := e.Execute("d", MustParse(`count(get(id=10, frame=f))`), nil); err != nil { t.Fatal(err) } else if n != uint64(12) { @@ -208,13 +224,16 @@ func TestExecutor_Execute_Remote_Set(t *testing.T) { } // Create local executor data. - e := NewExecutor(c) + idx := MustOpenIndex() + defer idx.Close() + + e := NewExecutor(idx.Index, c) if _, err := e.Execute("d", MustParse(`set(id=10, frame=f, profile_id=2)`), nil); err != nil { t.Fatal(err) } // Verify that one bit is set on both node's index. - if n := e.Index().Fragment("d", "f", 0).Bitmap(10).Count(); n != 1 { + if n := idx.MustFragment("d", "f", 0).Bitmap(10).Count(); n != 1 { t.Fatalf("unexpected local count: %d", n) } if !remoteCalled { @@ -229,13 +248,10 @@ type Executor struct { // NewExecutor returns a new instance of Executor. // The executor always matches the hostname of the first cluster node. -func NewExecutor(cluster *pilosa.Cluster) *Executor { - e := &Executor{ - Executor: pilosa.NewExecutor(pilosa.NewIndex()), - } +func NewExecutor(index *pilosa.Index, cluster *pilosa.Cluster) *Executor { + e := &Executor{Executor: pilosa.NewExecutor(index)} e.Cluster = cluster e.Host = cluster.Nodes[0].Host - return e } diff --git a/fragment.go b/fragment.go index 7404f8a19..18c5d2215 100644 --- a/fragment.go +++ b/fragment.go @@ -1,10 +1,17 @@ package pilosa import ( + "errors" + "fmt" + "os" "sort" "strings" "sync" + "syscall" "time" + "unsafe" + + "github.com/umbel/pilosa/roaring" ) // SliceWidth is the number of profile IDs in a slice. @@ -19,13 +26,20 @@ type Fragment struct { frame string slice uint64 + // File-backed storage + path string + file *os.File + storage *roaring.Bitmap + storageData []byte + // Bitmap cache. cache Cache } // NewFragment returns a new instance of Fragment. -func NewFragment(db, frame string, slice uint64) *Fragment { +func NewFragment(path, db, frame string, slice uint64) *Fragment { f := &Fragment{ + path: path, db: db, frame: frame, slice: slice, @@ -44,6 +58,115 @@ func NewFragment(db, frame string, slice uint64) *Fragment { return f } +// Open opens the underlying storage. +func (f *Fragment) Open() error { + f.mu.Lock() + defer f.mu.Unlock() + + // Initialize storage in a function so we can close if anything goes wrong. + if err := func() error { + // Create a roaring bitmap to serve as storage for the slice. + f.storage = roaring.NewBitmap() + + // Open the data file to be mmap'd and used as an ops log. + file, err := os.OpenFile(f.path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) + if err != nil { + return err + } + f.file = file + + // Lock the underlying file. + if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err == nil { + return nil + } + + // If the file is empty then initialize it with an empty bitmap. + fi, err := f.file.Stat() + if err != nil { + return err + } else if fi.Size() == 0 { + if _, err := f.storage.WriteTo(f.file); err != nil { + return err + } + } + + // Mmap the underlying file so it can be zero copied. + storageData, err := syscall.Mmap(int(f.file.Fd()), 0, int(fi.Size()), syscall.PROT_READ, syscall.MAP_SHARED) + if err != nil { + return fmt.Errorf("mmap: %s", err) + } + f.storageData = storageData + + // Advise the kernel that the mmap is accessed randomly. + if err := madvise(f.storageData, syscall.MADV_RANDOM); err != nil { + return fmt.Errorf("madvise: %s", err) + } + + // Attach the mmap file to the bitmap. + if err := f.storage.UnmarshalBinary((*[0x7FFFFFFF]byte)(unsafe.Pointer(&f.storageData[0]))[:]); err != nil { + return fmt.Errorf("unmarshal storage: %s", err) + } + + // Attach the file to the bitmap to act as a write-ahead log. + f.storage.OpWriter = f.file + + return nil + + }(); err != nil { + f.close() + return err + } + + return nil +} + +// Close flushes the underlying storage, closes the file and unlocks it. +func (f *Fragment) Close() error { + f.mu.Lock() + defer f.mu.Unlock() + return f.close() +} + +func (f *Fragment) close() error { + // Clear the storage bitmap so it doesn't access the closed mmap. + f.storage = roaring.NewBitmap() + + // Unmap the file. + if f.storageData != nil { + if err := syscall.Munmap(f.storageData); err != nil { + return fmt.Errorf("munmap: %s", err) + } + f.storageData = nil + } + + // Flush file, unlock & close. + if f.file != nil { + if err := f.file.Sync(); err != nil { + return fmt.Errorf("sync: %s", err) + } + if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_UN); err != nil { + return fmt.Errorf("unlock: %s", err) + } + if err := f.file.Close(); err != nil { + return fmt.Errorf("close file: %s", err) + } + } + + return nil +} + +// Path returns the path the fragment was initialized with. +func (f *Fragment) Path() string { return f.path } + +// DB returns the database the fragment was initialized with. +func (f *Fragment) DB() string { return f.db } + +// Frame returns the frame the fragment was initialized with. +func (f *Fragment) Frame() string { return f.frame } + +// Slice returns the slice the fragment was initialized with. +func (f *Fragment) Slice() uint64 { return f.slice } + // Bitmap returns a bitmap by ID. func (f *Fragment) Bitmap(bitmapID uint64) *Bitmap { f.mu.Lock() @@ -53,15 +176,20 @@ func (f *Fragment) Bitmap(bitmapID uint64) *Bitmap { func (f *Fragment) bitmap(bitmapID uint64) *Bitmap { // Read from cache. - if bm, ok := f.cache.Get(bitmapID); ok { + if bm := f.cache.Get(bitmapID); bm != nil { return bm } - // Read from storage engine. - // bm, filter := f.storage.Fetch(bitmapID, f.db, f.frame, f.slice) - + // Read bitmap from storage. bm := NewBitmap() + f.storage.ForEachRange(uint32(bitmapID)*SliceWidth, uint32(bitmapID+1)*SliceWidth, func(i uint32) { + profileID := (f.slice * SliceWidth) + (uint64(i) % SliceWidth) + bm.setBit(profileID) + }) + + // Add to the cache. f.cache.Add(bitmapID, 0 /*filter*/, bm) + return bm } @@ -95,6 +223,63 @@ func (f *Fragment) TopNAll(n int, categories []uint64) []Pair { return results } +// SetBit sets a bit for a given profile & bitmap within the fragment. +// This updates both the on-disk storage and the in-cache bitmap. +func (f *Fragment) SetBit(bitmapID, profileID uint64) error { + f.mu.Lock() + defer f.mu.Unlock() + + // Determine the position of the bit in the storage. + pos, err := f.pos(bitmapID, profileID) + if err != nil { + return err + } + + // Write to storage. + if err := f.storage.Add(pos); err != nil { + return err + } + + // Update the cache. + f.bitmap(bitmapID).setBit(profileID) + + return nil +} + +// ClearBit clears a bit for a given profile & bitmap within the fragment. +// This updates both the on-disk storage and the in-cache bitmap. +func (f *Fragment) ClearBit(bitmapID, profileID uint64) error { + f.mu.Lock() + defer f.mu.Unlock() + + // Determine the position of the bit in the storage. + pos, err := f.pos(bitmapID, profileID) + if err != nil { + return err + } + + // Write to storage. + if err := f.storage.Remove(pos); err != nil { + return err + } + + // Update the cache. + f.bitmap(bitmapID).clearBit(profileID) + + return nil +} + +// pos translates the bitmap ID and profile ID into a position in the storage bitmap. +func (f *Fragment) pos(bitmapID, profileID uint64) (uint32, error) { + // Return an error if the profile ID is out of the range of the fragment's slice. + minProfileID := f.slice * SliceWidth + if profileID < minProfileID || profileID >= minProfileID+SliceWidth { + return 0, errors.New("profile out of bounds") + } + + return uint32((bitmapID * SliceWidth) + (profileID % SliceWidth)), nil +} + func (f *Fragment) TopN(src *Bitmap, n int, categories []uint64) []Pair { f.mu.Lock() defer f.mu.Unlock() @@ -228,3 +413,11 @@ func (f *Fragment) Range(bitmapID uint64, start, end time.Time) *Bitmap { } return result } + +func madvise(b []byte, advice int) (err error) { + _, _, e1 := syscall.Syscall(syscall.SYS_MADVISE, uintptr(unsafe.Pointer(&b[0])), uintptr(len(b)), uintptr(advice)) + if e1 != 0 { + err = e1 + } + return +} diff --git a/fragment_test.go b/fragment_test.go index e7d3e593d..306914f82 100644 --- a/fragment_test.go +++ b/fragment_test.go @@ -1,8 +1,128 @@ package pilosa_test import ( + "io/ioutil" + "os" + "testing" + "github.com/umbel/pilosa" ) // SliceWidth is a helper reference to use when testing. const SliceWidth = pilosa.SliceWidth + +// Ensure a fragment can set a bit and retrieve it. +func TestFragment_SetBit(t *testing.T) { + f := MustOpenFragment("d", "f", 0) + defer f.Close() + + // Set bits on the fragment. + if err := f.SetBit(120, 1); err != nil { + t.Fatal(err) + } else if err := f.SetBit(120, 6); err != nil { + t.Fatal(err) + } else if err := f.SetBit(121, 0); err != nil { + t.Fatal(err) + } + + // Verify counts on bitmaps. + if n := f.Bitmap(120).Count(); n != 2 { + t.Fatalf("unexpected count: %d", n) + } else if n := f.Bitmap(121).Count(); n != 1 { + t.Fatalf("unexpected count: %d", n) + } + + // Close and reopen the fragment & verify the data. + if err := f.Reopen(); err != nil { + t.Fatal(err) + } else if n := f.Bitmap(120).Count(); n != 2 { + t.Fatalf("unexpected count (reopen): %d", n) + } else if n := f.Bitmap(121).Count(); n != 1 { + t.Fatalf("unexpected count (reopen): %d", n) + } +} + +// Ensure a fragment can clear a set bit. +func TestFragment_ClearBit(t *testing.T) { + f := MustOpenFragment("d", "f", 0) + defer f.Close() + + // Set and then clear bits on the fragment. + if err := f.SetBit(1000, 1); err != nil { + t.Fatal(err) + } else if err := f.SetBit(1000, 2); err != nil { + t.Fatal(err) + } else if err := f.ClearBit(1000, 1); err != nil { + t.Fatal(err) + } + + // Verify count on bitmap. + if n := f.Bitmap(1000).Count(); n != 1 { + t.Fatalf("unexpected count: %d", n) + } + + // Close and reopen the fragment & verify the data. + if err := f.Reopen(); err != nil { + t.Fatal(err) + } else if n := f.Bitmap(1000).Count(); n != 1 { + t.Fatalf("unexpected count (reopen): %d", n) + } +} + +// Fragment is a test wrapper for pilosa.Fragment. +type Fragment struct { + *pilosa.Fragment +} + +// NewFragment returns a new instance of Fragment with a temporary path. +func NewFragment(db, frame string, slice uint64) *Fragment { + file, err := ioutil.TempFile("", "pilosa-fragment-") + if err != nil { + panic(err) + } + file.Close() + return &Fragment{Fragment: pilosa.NewFragment(file.Name(), db, frame, slice)} +} + +// MustOpenFragment creates and opens an fragment at a temporary path. Panic on error. +func MustOpenFragment(db, frame string, slice uint64) *Fragment { + f := NewFragment(db, frame, slice) + if err := f.Open(); err != nil { + panic(err) + } + return f +} + +// Close closes the fragment and removes all underlying data. +func (f *Fragment) Close() error { + defer os.Remove(f.Path()) + return f.Fragment.Close() +} + +// Reopen closes the fragment and reopens it as a new instance. +func (f *Fragment) Reopen() error { + path := f.Path() + if err := f.Close(); err != nil { + return err + } + + f = &Fragment{Fragment: pilosa.NewFragment(path, f.DB(), f.Frame(), f.Slice())} + if err := f.Open(); err != nil { + return err + } + return nil +} + +// MustSetBit sets a bit on a bitmap. Panic on error. +func (f *Fragment) MustSetBit(bitmapID, profileID uint64) { + if err := f.SetBit(bitmapID, profileID); err != nil { + panic(err) + } +} + +// MustClearBit clears a bit on a bitmap. Panic on error. +func (f *Fragment) MustClearBit(bitmapID, profileID uint64) { + if err := f.ClearBit(bitmapID, profileID); err != nil { + panic(err) + } +} diff --git a/handler_test.go b/handler_test.go index 22e784f6d..f2fb0c2d8 100644 --- a/handler_test.go +++ b/handler_test.go @@ -138,12 +138,7 @@ func TestHandler_Query_Uint64_Protobuf(t *testing.T) { func TestHandler_Query_Bitmap_JSON(t *testing.T) { h := NewHandler() h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64) (interface{}, error) { - bm := pilosa.NewBitmap() - bm.SetBit(1) - bm.SetBit(3) - bm.SetBit(66) - bm.SetBit(pilosa.SliceWidth + 1) - return bm, nil + return pilosa.NewBitmap(1, 3, 66, pilosa.SliceWidth+1), nil } w := httptest.NewRecorder() @@ -159,10 +154,7 @@ func TestHandler_Query_Bitmap_JSON(t *testing.T) { func TestHandler_Query_Bitmap_Protobuf(t *testing.T) { h := NewHandler() h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64) (interface{}, error) { - bm := pilosa.NewBitmap() - bm.SetBit(1) - bm.SetBit(pilosa.SliceWidth + 1) - return bm, nil + return pilosa.NewBitmap(1, pilosa.SliceWidth+1), nil } w := httptest.NewRecorder() diff --git a/index.go b/index.go index 8ee0906f4..c2025fc84 100644 --- a/index.go +++ b/index.go @@ -1,23 +1,51 @@ package pilosa import ( + "fmt" + "log" + "os" + "path/filepath" + "strconv" "sync" ) // Index represents a container for fragments. type Index struct { mu sync.Mutex + path string sliceN uint64 fragments map[fragmentKey]*Fragment } // NewIndex returns a new instance of Index. -func NewIndex() *Index { +func NewIndex(path string) *Index { return &Index{ + path: path, fragments: make(map[fragmentKey]*Fragment), } } +// Open initializes the root data directory for the index. +func (i *Index) Open() error { + if err := os.MkdirAll(i.path, 0777); err != nil { + return err + } + return nil +} + +// Close closes all open fragments. +func (i *Index) Close() error { + for key, f := range i.fragments { + if err := f.Close(); err != nil { + log.Println("error closing fragment(%v): %s", key, err) + } + } + return nil +} + +// Path returns the path the index was initialized with. +func (i *Index) Path() string { return i.path } + // SliceN returs the total number of slices managed by the index. func (i *Index) SliceN() uint64 { i.mu.Lock() @@ -25,9 +53,14 @@ func (i *Index) SliceN() uint64 { return i.sliceN } +// FragmentPath returns the path where a given fragment is stored. +func (i *Index) FragmentPath(db, frame string, slice uint64) string { + return filepath.Join(i.path, db, frame, strconv.FormatUint(slice, 10)) +} + // Fragment returns the fragment for a database, frame & slice. // The fragment is created if it doesn't already exist. -func (i *Index) Fragment(db, frame string, slice uint64) *Fragment { +func (i *Index) Fragment(db, frame string, slice uint64) (*Fragment, error) { i.mu.Lock() defer i.mu.Unlock() @@ -39,10 +72,22 @@ func (i *Index) Fragment(db, frame string, slice uint64) *Fragment { // Create fragment, if not exists. key := fragmentKey{db, frame, slice} if i.fragments[key] == nil { - i.fragments[key] = NewFragment(db, frame, slice) + path := i.FragmentPath(db, frame, slice) + + // Create parent directory, if necessary. + if err := os.MkdirAll(filepath.Dir(path), 0777); err != nil { + return nil, fmt.Errorf("parent fragment dir: %s", err) + } + + // Initialize and open fragment. + f := NewFragment(path, db, frame, slice) + if err := f.Open(); err != nil { + return nil, err + } + i.fragments[key] = f } - return i.fragments[key] + return i.fragments[key], nil } // fragmentKey is the map key for fragment look ups. diff --git a/index_test.go b/index_test.go new file mode 100644 index 000000000..4f4d66597 --- /dev/null +++ b/index_test.go @@ -0,0 +1,46 @@ +package pilosa_test + +import ( + "io/ioutil" + "os" + + "github.com/umbel/pilosa" +) + +// Index is a test wrapper for pilosa.Index. +type Index struct { + *pilosa.Index +} + +// NewIndex returns a new instance of Index with a temporary path. +func NewIndex() *Index { + path, err := ioutil.TempDir("", "pilosa-") + if err != nil { + panic(err) + } + return &Index{Index: pilosa.NewIndex(path)} +} + +// MustOpenIndex creates and opens an index at a temporary path. Panic on error. +func MustOpenIndex() *Index { + i := NewIndex() + if err := i.Open(); err != nil { + panic(err) + } + return i +} + +// Close closes the index and removes all underlying data. +func (i *Index) Close() error { + defer os.RemoveAll(i.Path()) + return i.Index.Close() +} + +// MustFragment returns a given fragment. Panic on error. +func (i *Index) MustFragment(db, frame string, slice uint64) *Fragment { + f, err := i.Index.Fragment(db, frame, slice) + if err != nil { + panic(err) + } + return &Fragment{Fragment: f} +}