place advisory lock on db

Author: dsa0x

bloom.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/dsa0x/sprout/pkg/murmur"
 	"github.com/edsrzf/mmap-go"
+	"github.com/juju/fslock"
 )
 
 type BloomFilter struct {
@@ -33,6 +34,7 @@ type BloomFilter struct {
 	mem        mmap.MMap
 	pageOffset int
 	lock       sync.Mutex
+	flock      *fslock.Lock
 	byteSize   int
 
 	// m is the number bits per slice(hashFn)
@@ -140,7 +142,17 @@ func NewBloom(opts *BloomOptions) *BloomFilter {
 		opts:       opts,
 	}
 
-	err := bf.mmap()
+	// initialize advisory lock
+	bf.flock = fslock.New(bf.path)
+
+	// open the file
+	err := bf.openFile()
+	if err != nil {
+		log.Fatalf("Error opening file: %v", err)
+	}
+
+	// open mmap the file
+	err = bf.mmap()
 	if err != nil {
 		log.Fatalf("Mmap error: %v", err)
 	}
@@ -251,21 +263,6 @@ func (bf *BloomFilter) hasStore() bool {
 	return bf.db != nil && bf.db.isReady()
 }
 
-func (bf *BloomFilter) unmap() error {
-	var err error
-	if bf.mem != nil {
-		err = bf.mem.Unmap()
-		if err != nil {
-			_ = bf.memFile.Close()
-			return err
-		}
-	}
-	if bf.memFile != nil {
-		return bf.memFile.Close()
-	}
-	return nil
-}
-
 // getBitIndex returns the index and mask for the bit. (unused)
 //
 // The first half of the bits are set at the beginning of the byte,
@@ -325,13 +322,27 @@ func (bf *BloomFilter) Capacity() int {
 	return bf.capacity
 }
 
-// Close closes the file handle to the filter and the persistent store (if any)
+// Close flushes the file to disk and closes the file handle to the filter
 func (bf *BloomFilter) Close() error {
 	if err := bf.mem.Flush(); err != nil {
 		_ = bf.memFile.Close()
 		return err
 	}
-	return bf.unmap()
+	if err := bf.unmap(); err != nil {
+		_ = bf.memFile.Close()
+		return err
+	}
+	if bf.flock != nil {
+		err := bf.flock.Unlock()
+		if err != nil {
+			return err
+		}
+	}
+
+	if bf.memFile != nil {
+		return bf.memFile.Close()
+	}
+	return nil
 }
 
 // Count returns the number of items added to the bloom filter
@@ -351,9 +362,14 @@ func (bf *BloomFilter) DB() interface{} {
 
 // Clear resets all bits in the bloom filter
 func (bf *BloomFilter) Clear() {
-	bf.lock.Lock()
-	defer bf.lock.Unlock()
-	bf.mem = make([]byte, bf.bit_width)
+	mem := make([]byte, bf.bit_width)
+	copy(bf.mem, mem)
+	err := bf.mem.Flush()
+	if err != nil {
+		fmt.Printf("Error flushing filter to disk: %s\n", err)
+		os.Exit(1)
+	}
+	bf.count = 0
 }
 
 type BloomFilterStats struct {
@@ -379,21 +395,47 @@ func (bf *BloomFilter) Stats() BloomFilterStats {
 	}
 }
 
-// mmap opens a the filter file and maps it into memory
-func (bf *BloomFilter) mmap() error {
+func (bf *BloomFilter) unmap() error {
 	var err error
-	bf.memFile, err = os.OpenFile(bf.path, os.O_RDWR|os.O_CREATE, 0666)
-	if err != nil {
-		return fmt.Errorf("Unable to open bloom filter file: %s", err)
+	if bf.mem != nil {
+		err = bf.mem.Unmap()
+		if err != nil {
+			return err
+		}
 	}
+	return nil
+}
+
+// mmap opens the filter file and maps it into memory
+func (bf *BloomFilter) mmap() error {
+
+	var err error
 
 	if err := bf.memFile.Truncate(int64(bf.opts.dataSize)); err != nil {
 		log.Fatalf("Error truncating file: %s", err)
 	}
 
 	bf.mem, err = mmap.MapRegion(bf.memFile, bf.opts.dataSize, mmap.RDWR, 0, 0)
 	if err != nil {
-		return fmt.Errorf("Unable to mmap bloom filter file: %s", err)
+		return fmt.Errorf("unable to mmap bloom filter file: %s", err)
+	}
+
+	return nil
+}
+
+// openFile opens the filter file and locks it
+func (bf *BloomFilter) openFile() error {
+	var err error
+	bf.memFile, err = os.OpenFile(bf.path, os.O_RDWR|os.O_CREATE, 0666)
+	if err != nil {
+		return fmt.Errorf("unable to open bloom filter file: %s", err)
+	}
+
+	if err := bf.flock.TryLock(); err != nil {
+		if err == fslock.ErrLocked {
+			return fmt.Errorf("file is locked by another process")
+		}
+		return fmt.Errorf("unable to lock bloom filter file: %s", err)
 	}
 
 	return nil

bloom_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"os"
+	"os/exec"
 	"testing"
 
 	"github.com/dgraph-io/badger/v3"
@@ -52,7 +53,7 @@ func TestBloomFilter_Add(t *testing.T) {
 		opts := &BloomOptions{
 			Err_rate: 0.01,
 			Capacity: 110000,
-			Path:     "./test.db",
+			Path:     "./test1.db",
 		}
 		bf := NewBloom(opts)
 
@@ -84,7 +85,7 @@ func TestBloomFilter_Add(t *testing.T) {
 		opts := &BloomOptions{
 			Err_rate: 0.01,
 			Capacity: 1000,
-			Path:     "./test.db",
+			Path:     "./test2.db",
 		}
 		bf := NewBloom(opts)
 
@@ -112,7 +113,7 @@ func TestBloomFilter_Add(t *testing.T) {
 		opts := &BloomOptions{
 			Err_rate: 0.1,
 			Capacity: 100,
-			Path:     "./test.db",
+			Path:     "./test3.db",
 		}
 		bf := NewBloom(opts)
 
@@ -134,12 +135,15 @@ func TestBloomFilter_Merge(t *testing.T) {
 		Path:     "./test.db",
 	}
 	bf := NewBloom(opts)
-	bf2 := NewBloom(opts)
+	opts2 := *opts
+	opts2.Path = "./test2.db"
+	bf2 := NewBloom(&opts2)
 
 	defer func() {
 		bf.Close()
 		bf2.Close()
 		os.Remove(opts.Path)
+		os.Remove(opts2.Path)
 	}()
 
 	t.Run("merge success", func(t *testing.T) {
@@ -153,11 +157,12 @@ func TestBloomFilter_Merge(t *testing.T) {
 		opts := &BloomOptions{
 			Err_rate: 0.01,
 			Capacity: 10000,
-			Path:     "./test.db",
+			Path:     "./test3.db",
 		}
 		bf2 := NewBloom(opts)
 		bf.Merge(bf2)
 		defer func() {
+			bf.Close()
 			bf2.Close()
 			os.Remove(opts.Path)
 		}()
@@ -170,8 +175,15 @@ func TestBloomFilter_Merge(t *testing.T) {
 
 	t.Run("object added to the single filters should be found in the resulting merge", func(t *testing.T) {
 		key := []byte("foo")
+		opts := &BloomOptions{
+			Err_rate: 0.01,
+			Capacity: 1000,
+			Path:     "./test4.db",
+		}
 		bf := NewBloom(opts)
-		bf2 := NewBloom(opts)
+		opts2 := *opts
+		opts2.Path = "./test5.db"
+		bf2 := NewBloom(&opts2)
 		bf2.Add(key)
 		err := bf.Merge(bf2)
 		if err != nil {
@@ -180,6 +192,12 @@ func TestBloomFilter_Merge(t *testing.T) {
 		if !bf.Contains(key) {
 			t.Errorf("Expected key %s to be found in the merged filter", string(key))
 		}
+		defer func() {
+			bf.Close()
+			bf2.Close()
+			os.Remove(opts.Path)
+			os.Remove(opts2.Path)
+		}()
 	})
 
 }
@@ -309,3 +327,81 @@ func assertPanic(t *testing.T, fn func()) {
 	fn()
 	t.Errorf("The code did not panic")
 }
+
+func TestBloomFilter_Clear(t *testing.T) {
+	store, cleanupFunc := DBSetupTest(t)
+	defer cleanupFunc()
+	opts := &BloomOptions{
+		Err_rate: 0.01,
+		Capacity: 1000,
+		Database: store,
+		Path:     "./test.db",
+	}
+	bf := NewBloom(opts)
+
+	defer func() {
+		bf.Close()
+		os.Remove(opts.Path)
+	}()
+
+	for i := 0; i < opts.Capacity/2; i++ {
+		bf.Add([]byte(fmt.Sprintf("foo%d", i)))
+	}
+
+	keys := []string{"foo", "baz", "bar"}
+	for _, key := range keys {
+		bf.Add([]byte(key))
+	}
+
+	t.Run("should clear the bloom filter", func(t *testing.T) {
+		bf.Clear()
+		for _, key := range keys {
+			if bf.Contains([]byte(key)) {
+				t.Errorf("Expected key to not be found in bloom filter after clear")
+			}
+		}
+	})
+
+	for i := 0; i < opts.Capacity/2; i++ {
+		bf.Add([]byte(fmt.Sprintf("foo%d", i)))
+	}
+
+	for _, key := range keys {
+		bf.Add([]byte(key))
+		t.Run("Should find the newly added keys", func(t *testing.T) {
+			if !bf.Contains([]byte(key)) {
+				t.Errorf("Expected key to be found in bloom filter after clear")
+			}
+		})
+	}
+}
+func TestBloomFilter_FileLock(t *testing.T) {
+	t.Run("should clear the bloom filter", func(t *testing.T) {
+		store, cleanupFunc := DBSetupTest(t)
+		defer cleanupFunc()
+		if os.Getenv("SPROUT_LOCK") == "1" {
+			opts := &BloomOptions{
+				Err_rate: 0.01,
+				Capacity: 1000,
+				Database: store,
+				Path:     "./test.db",
+			}
+			bf := NewBloom(opts)
+			defer func() {
+				bf.Close()
+				os.Remove(opts.Path)
+			}()
+			NewBloom(opts)
+			return
+		}
+
+		cmd := exec.Command(os.Args[0], "-test.run=TestBloomFilter_FileLock")
+		cmd.Env = append(os.Environ(), "SPROUT_LOCK=1")
+		err := cmd.Run()
+		if e, ok := err.(*exec.ExitError); ok && !e.Success() {
+			return
+		}
+		t.Fatalf("expected file lock error, got none")
+	})
+
+}