Add support for limited memory use

Author: mangalaman93

.travis.yml

@@ -1,8 +1,8 @@
 language: go
 
 go:
-  - 1.10.3
   - 1.11
+  - 1.12
 
 os:
   - osx

arena.go

@@ -31,6 +31,11 @@ func newArena(file string, size int) (*arena, error) {
 		return nil, err
 	}
 
+	// We can close the file descriptor here
+	if err := fd.Close(); err != nil {
+		return nil, err
+	}
+
 	return &arena{
 		IMmap: m,
 		size:  size,

arena_test.go

@@ -9,26 +9,47 @@ import (
 )
 
 func TestNewArenaNoDir(t *testing.T) {
-	arena, err := newArena(fmt.Sprintf("%d/temp.dat", time.Now().UnixNano()), 100)
-	if arena != nil || err == nil || os.IsExist(err) {
+	aa, err := newArena(fmt.Sprintf("%d/temp.dat", time.Now().UnixNano()), 100)
+	if aa != nil || err == nil || os.IsExist(err) {
 		t.Fatalf("unexpected return for newArena :: %v", err)
 	}
 }
 
+func TestNewArenaNoReadPerm(t *testing.T) {
+	fileName := path.Join(os.TempDir(), "temp.dat")
+	defer os.Remove(fileName)
+	if _, err := os.OpenFile(fileName, os.O_RDWR|os.O_CREATE, 000); err != nil {
+		t.Fatalf("unable to create file :: %v", err)
+	}
+
+	arena, err := newArena(fileName, 100)
+	if arena != nil || err == nil || !os.IsPermission(err) {
+		aa, err := newArena("/temp.dat", 100)
+		if aa != nil || err == nil || !os.IsPermission(err) {
+			t.Fatalf("unexpected return for newArena :: %v", err)
+		}
+	}
+}
+
 func TestNewArenaNoFile(t *testing.T) {
 	arenaSize := 100
 	fileName := path.Join(os.TempDir(), "temp.dat")
 	defer os.Remove(fileName)
 
-	arena, err := newArena(fileName, arenaSize)
+	aa, err := newArena(fileName, arenaSize)
 	if err != nil {
 		t.Fatalf("error in creating new arena: %v", err)
 	}
-	defer arena.Unmap()
+	defer func() {
+		err := aa.Unmap()
+		if err != nil {
+			t.Fatalf("error occurred while unmapping: %v", err)
+		}
+	}()
 
 	// ensure arena struct stores correct size
-	if arena.size != arenaSize {
-		t.Fatalf("arena size do not match, exp: %v, actual: %v", arenaSize, arena.size)
+	if aa.size != arenaSize {
+		t.Fatalf("arena size do not match, exp: %v, actual: %v", arenaSize, aa.size)
 	}
 
 	// ensure underlined file is of correct size
@@ -55,15 +76,20 @@ func TestNewArenaLargerFile(t *testing.T) {
 	}
 
 	// creating new arena
-	arena, err := newArena(fileName, arenaSize)
+	aa, err := newArena(fileName, arenaSize)
 	if err != nil {
 		t.Fatalf("error in creating new arena: %v", err)
 	}
-	defer arena.Unmap()
+	defer func() {
+		err := aa.Unmap()
+		if err != nil {
+			t.Fatalf("error occurred while unmapping: %v", err)
+		}
+	}()
 
 	// ensure arena struct stores correct size
-	if arena.size != arenaSize {
-		t.Fatalf("arena size do not match, exp: %v, actual: %v", arenaSize, arena.size)
+	if aa.size != arenaSize {
+		t.Fatalf("arena size do not match, exp: %v, actual: %v", arenaSize, aa.size)
 	}
 
 	// ensure underlined file is still of original size
@@ -78,8 +104,8 @@ func TestNewArenaLargerFile(t *testing.T) {
 
 func TestNewArenaNoFolder(t *testing.T) {
 	arenaSize := 100
-	arena, err := newArena("1/2/3/4/5/6/arena.dat", arenaSize)
-	if !os.IsNotExist(err) || arena != nil {
+	aa, err := newArena("1/2/3/4/5/6/aa.dat", arenaSize)
+	if !os.IsNotExist(err) || aa != nil {
 		t.Fatalf("expected file not exists error, returned: %v", err)
 	}
 }

arenamanager.go

@@ -3,6 +3,7 @@ package bigqueue
 import (
 	"errors"
 	"fmt"
+	"os"
 	"path"
 )
 
@@ -11,27 +12,33 @@ const (
 )
 
 var (
-	// errShouldNotReach is returned when code reaches unexpected places
+	// errShouldNotReach is returned when an invariant is not true anymore
 	errShouldNotReach = errors.New("SHOULD NOT REACH HERE")
 )
 
 // arenaManager manages all the arenas for a bigqueue
 type arenaManager struct {
-	dir       string
-	conf      *bqConfig
-	baseAid   int
-	arenaList []*arena
+	dir         string
+	conf        *bqConfig
+	index       *queueIndex
+	baseAid     int
+	arenaList   []*arena
+	inMemArenas int
 }
 
 // newArenaManager returns a pointer to new arenaManager
-func newArenaManager(dir string, conf *bqConfig, headAid, tailAid int) (
+func newArenaManager(dir string, conf *bqConfig, index *queueIndex) (
 	*arenaManager, error) {
 
+	headAid, _ := index.getHead()
+	tailAid, _ := index.getTail()
+
 	numArenas := tailAid + 1 - headAid
 	arenaList := make([]*arena, numArenas)
 	am := &arenaManager{
 		dir:       dir,
 		conf:      conf,
+		index:     index,
 		baseAid:   headAid,
 		arenaList: arenaList,
 	}
@@ -40,6 +47,7 @@ func newArenaManager(dir string, conf *bqConfig, headAid, tailAid int) (
 	if err := am.loadArenaIntoMemory(headAid); err != nil {
 		return nil, err
 	}
+
 	if err := am.loadArenaIntoMemory(tailAid); err != nil {
 		return nil, err
 	}
@@ -63,6 +71,12 @@ func (m *arenaManager) getArena(aid int) (*arena, error) {
 		return aa, nil
 	}
 
+	// before we get a new arena into memory, we need to ensure that after fetching
+	// a new arena into memory, we do not cross the provided memory limit
+	if err := m.ensureEnoughMem(); err != nil {
+		return nil, err
+	}
+
 	// otherwise, get arena into memory
 	if err := m.loadArenaIntoMemory(aid); err != nil {
 		return nil, err
@@ -71,18 +85,103 @@ func (m *arenaManager) getArena(aid int) (*arena, error) {
 	return m.arenaList[relAid], nil
 }
 
+// ensureEnoughMem ensures that at least 1 new arena can be brought into memory
+// TODO: shrink arenaList
+func (m *arenaManager) ensureEnoughMem() error {
+	// Check whether head has moved and arenas can be unmpped
+	// Remove all such arenas from memory, irrespectively
+	headAid, _ := m.index.getHead()
+	for aid := m.baseAid; aid < headAid; aid++ {
+		if err := m.unloadArenaFromMemory(aid); err != nil {
+			return err
+		}
+
+		// disk garbage collection
+		if err := m.deleteArenaBackedFile(aid); err != nil {
+			return err
+		}
+	}
+
+	// if no limit on # of arenas, no need for eviction
+	if m.conf.maxInMemArenas == 0 {
+		return nil
+	}
+
+	// Check whether an eviction is needed to begin with
+	if m.inMemArenas < m.conf.maxInMemArenas {
+		return nil
+	}
+
+	// Start evicting from the arena just before the last arena that we have.
+	// If message size > arena size, last arena may not always be the tail arena.
+	// We always ensure that head and tail arenas are not evicted from memory.
+	// Assuming m.conf.maxInMemArenas >= 3.
+	// Simply iterate from the last arena until enough memory is
+	// available for a new arena to be loaded into memory
+	tailAid, _ := m.index.getTail()
+	curAid := m.baseAid + len(m.arenaList)
+	for m.conf.maxInMemArenas-m.inMemArenas <= 0 {
+		curAid--
+
+		if curAid < 0 {
+			return errShouldNotReach
+		}
+
+		if curAid == tailAid || curAid == headAid {
+			continue
+		}
+
+		if err := m.unloadArenaFromMemory(curAid); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // loadArenaIntoMemory will fetch the arena into memory
 func (m *arenaManager) loadArenaIntoMemory(aid int) error {
+	if m.arenaList[aid-m.baseAid] != nil {
+		return nil
+	}
+
 	filePath := path.Join(m.dir, fmt.Sprintf(cArenaFileFmt, aid))
 	aa, err := newArena(filePath, m.conf.arenaSize)
 	if err != nil {
 		return err
 	}
 
+	m.inMemArenas++
 	m.arenaList[aid-m.baseAid] = aa
 	return nil
 }
 
+// unloadArenaFromMemory will remove the arena from memory
+func (m *arenaManager) unloadArenaFromMemory(aid int) error {
+	if m.arenaList[aid-m.baseAid] == nil {
+		return nil
+	}
+
+	if err := m.arenaList[aid-m.baseAid].Unmap(); err != nil {
+		return err
+	}
+
+	m.inMemArenas--
+	m.arenaList[aid-m.baseAid] = nil
+	return nil
+}
+
+// deleteArenaBackedFile deletes the backed file for given arena with
+// arena id: aid. If file doesn't exist, the error is ignored.
+func (m *arenaManager) deleteArenaBackedFile(aid int) error {
+	filePath := path.Join(m.dir, fmt.Sprintf(cArenaFileFmt, aid))
+	if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) {
+		return err
+	}
+
+	return nil
+}
+
 // close unmaps all the arenas managed by arenaManager
 func (m *arenaManager) close() error {
 	var retErr error

bigqueue.go

@@ -27,7 +27,7 @@ type BigQueue struct {
 }
 
 // NewBigQueue constructs an instance of *BigQueue
-func NewBigQueue(dir string, opts ...Option) (IBigQueue, error) {
+func NewBigQueue(dir string, opts ...Option) (*BigQueue, error) {
 	complete := false
 
 	// setup configuration
@@ -50,9 +50,7 @@ func NewBigQueue(dir string, opts ...Option) (IBigQueue, error) {
 	}()
 
 	// create arena manager
-	headAid, _ := index.getHead()
-	tailAid, _ := index.getTail()
-	am, err := newArenaManager(dir, conf, headAid, tailAid)
+	am, err := newArenaManager(dir, conf, index)
 	if err != nil {
 		return nil, err
 	}