1 files changed, 31 insertions, 18 deletions
diff --git a/db.go b/db.go
index f3fe278..64de16e 100644
--- a/db.go
+++ b/db.go
@@ -12,9 +12,6 @@ import (
 	"unsafe"
 )
 
-// The smallest size that the mmap can be.
-const minMmapSize = 1 << 22 // 4MB
-
 // The largest step that can be taken when remapping the mmap.
 const maxMmapStep = 1 << 30 // 1GB
 
@@ -174,11 +171,9 @@ func (db *DB) mmap(minsz int) error {
 	if size < minsz {
 		size = minsz
 	}
-	size = db.mmapSize(size)
-
-	// Verify the map size is not above the maximum allowed.
-	if size > maxMapSize {
-		return fmt.Errorf("mmap too large")
+	size, err = db.mmapSize(size)
+	if err != nil {
+		return err
 	}
 
 	// Dereference all mmap references before unmapping.
@@ -221,21 +216,39 @@ func (db *DB) munmap() error {
 
 // mmapSize determines the appropriate size for the mmap given the current size
 // of the database. The minimum size is 4MB and doubles until it reaches 1GB.
-func (db *DB) mmapSize(size int) int {
-	if size <= minMmapSize {
-		return minMmapSize
-	} else if size < maxMmapStep {
-		size *= 2
-	} else {
-		size += maxMmapStep
+// Returns an error if the new mmap size is greater than the max allowed.
+func (db *DB) mmapSize(size int) (int, error) {
+	// Double the size from 1MB until 1GB.
+	for i := uint(20); i <= 30; i++ {
+		if size <= 1<<i {
+			return 1 << i, nil
+		}
+	}
+
+	// Verify the requested size is not above the maximum allowed.
+	if size > maxMapSize {
+		return 0, fmt.Errorf("mmap too large")
+	}
+
+	// If larger than 1GB then grow by 1GB at a time.
+	sz := int64(size) + int64(maxMmapStep)
+	if remainder := sz % int64(maxMmapStep); remainder > 0 {
+		sz -= remainder
 	}
 
 	// Ensure that the mmap size is a multiple of the page size.
-	if (size % db.pageSize) != 0 {
-		size = ((size / db.pageSize) + 1) * db.pageSize
+	// This should always be true since we're incrementing in MBs.
+	pageSize := int64(db.pageSize)
+	if (sz % pageSize) != 0 {
+		sz = ((sz / pageSize) + 1) * pageSize
+	}
+
+	// If we've exceeded the max size then only grow up to the max size.
+	if sz > maxMapSize {
+		sz = maxMapSize
 	}
 
-	return size
+	return int(sz), nil
 }
 
 // init creates a new database file and initializes its meta pages.