Fix sandbox allocator tests

sandbox/allocator/allocator.go

@@ -2,11 +2,10 @@ package allocator
 
 import (
 	"sync"
-
-	"github.com/joshjms/castletown/config"
 )
 
 const DEFAULT_SIZE uint32 = 65536
+const START_UID_GID uint32 = 1000000
 
 type Range struct {
 	UidStart uint32
@@ -16,61 +15,51 @@ type Range struct {
 }
 
 type Allocator struct {
-	ranges []Range
-	used   []bool
+	used map[int]bool
+	mex  int
 
 	mu sync.Mutex
 }
 
-func NewAllocator() (*Allocator, error) {
-	startUid := uint32(100000)
-	startGid := uint32(100000)
-	maxContainers := config.MaxConcurrency
-
-	ranges := make([]Range, maxContainers)
-	used := make([]bool, maxContainers)
-	for i := range maxContainers {
-		ranges[i] = Range{
-			UidStart: startUid + uint32(i)*DEFAULT_SIZE,
-			UidSize:  DEFAULT_SIZE,
-			GidStart: startGid + uint32(i)*DEFAULT_SIZE,
-			GidSize:  DEFAULT_SIZE,
-		}
-		used[i] = false
-	}
-
+func NewAllocator() *Allocator {
 	return &Allocator{
-		ranges: ranges,
-		used:   used,
-	}, nil
+		used: make(map[int]bool),
+		mex:  0,
+	}
 }
 
 func (a *Allocator) Allocate() (int, Range) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	for i, used := range a.used {
-		if !used {
-			a.used[i] = true
-			return i, a.ranges[i]
-		}
+	a.used[a.mex] = true
+	r := Range{
+		UidStart: START_UID_GID + uint32(a.mex)*DEFAULT_SIZE,
+		UidSize:  DEFAULT_SIZE,
+		GidStart: START_UID_GID + uint32(a.mex)*DEFAULT_SIZE,
+		GidSize:  DEFAULT_SIZE,
+	}
+	use := a.mex
+
+	for a.used[a.mex] {
+		a.mex++
 	}
 
-	return -1, Range{}
+	return use, r
 }
 
 func (a *Allocator) Free(i int) int {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	if i < 0 || i >= len(a.used) {
+	if i < 0 || !a.used[i] {
 		return -1
 	}
 
-	if !a.used[i] {
-		return -1
+	delete(a.used, i)
+	if i < a.mex {
+		a.mex = i
 	}
 
-	a.used[i] = false
-	return i
+	return 0
 }

sandbox/allocator/allocator_test.go

@@ -3,22 +3,24 @@ package allocator_test
 import (
 	"testing"
 
+	"github.com/joshjms/castletown/config"
 	"github.com/joshjms/castletown/sandbox/allocator"
 	"github.com/stretchr/testify/require"
 )
 
 func TestAllocator(t *testing.T) {
-	a, err := allocator.NewAllocator()
-	require.NoError(t, err, "Failed to create allocator: %v", err)
+	config.UseDefaults()
+
+	a := allocator.NewAllocator()
 
 	i1, _ := a.Allocate()
 	i2, _ := a.Allocate()
 	i3, _ := a.Allocate()
 	a.Free(i1)
 	i4, _ := a.Allocate()
 
-	require.Equal(t, i1, 0, "incorrect index for first allocation, expected 0")
-	require.Equal(t, i2, 1, "incorrect index for second allocation, expected 1")
-	require.Equal(t, i3, -1, "incorrect index for third allocation, expected -1")
-	require.Equal(t, i4, 0, "incorrect index for fourth allocation, expected 0")
+	require.Equal(t, 0, i1, "incorrect index for first allocation, expected 0")
+	require.Equal(t, 1, i2, "incorrect index for second allocation, expected 1")
+	require.Equal(t, 2, i3, "incorrect index for third allocation, expected 2")
+	require.Equal(t, 0, i4, "incorrect index for fourth allocation, expected 0")
 }