Don't create an extra fd bitmap to allocate a new fd.

4 files changed, 88 insertions, 74 deletions

diff --git a/pkg/bitmap/BUILD b/pkg/bitmap/BUILD
index 1af60002a..0f1e7006d 100644
--- a/pkg/bitmap/BUILD
+++ b/pkg/bitmap/BUILD
@@ -6,7 +6,6 @@ go_library(
     name = "bitmap",
     srcs = ["bitmap.go"],
     visibility = ["//:sandbox"],
-    deps = ["@org_golang_x_sys//unix:go_default_library"],
 )
 
 go_test(
@@ -14,5 +13,4 @@ go_test(
     size = "small",
     srcs = ["bitmap_test.go"],
     library = ":bitmap",
-    deps = ["@org_golang_x_sys//unix:go_default_library"],
 )
diff --git a/pkg/bitmap/bitmap.go b/pkg/bitmap/bitmap.go
index 40cf88103..803b7b3c7 100644
--- a/pkg/bitmap/bitmap.go
+++ b/pkg/bitmap/bitmap.go
@@ -25,7 +25,7 @@ import (
 // +stateify savable
 type Bitmap struct {
 	// numOnes is the number of ones in the bitmap.
-	numOnes  uint32
+	numOnes uint32
 
 	// bitBlock holds the bits. The type of bitBlock is uint64 which means
 	// each number in bitBlock contains 64 entries.
@@ -56,6 +56,28 @@ func (b *Bitmap) Minimum() uint32 {
 	return math.MaxInt32
 }
 
+// FirstZero returns the first unset bit from the range [start, ).
+func (b *Bitmap) FirstZero(start uint32) uint32 {
+	i, nbit := int(start/64), start%64
+	n := len(b.bitBlock)
+	if i >= n {
+		return math.MaxInt32
+	}
+	w := b.bitBlock[i] | ((1 << nbit) - 1)
+	for {
+		if w != ^uint64(0) {
+			r := bits.TrailingZeros64(^w)
+			return uint32(r + i*64)
+		}
+		i++
+		if i == n {
+			break
+		}
+		w = b.bitBlock[i]
+	}
+	return math.MaxInt32
+}
+
 // Maximum return the largest value in the Bitmap.
 func (b *Bitmap) Maximum() uint32 {
 	for i := len(b.bitBlock) - 1; i >= 0; i-- {
@@ -127,13 +149,13 @@ func (b *Bitmap) flipRange(begin, end uint32) {
 	beginBlock := begin / 64
 	endBlock := end / 64
 	if beginBlock == endBlock {
-		b.bitBlock[endBlock] ^= ((^uint64(0) << uint(begin%64)) & ((uint64(1) << (uint(end) % 64 + 1))-1))
+		b.bitBlock[endBlock] ^= ((^uint64(0) << uint(begin%64)) & ((uint64(1) << (uint(end)%64 + 1)) - 1))
 	} else {
 		b.bitBlock[beginBlock] ^= ^(^uint64(0) << uint(begin%64))
 		for i := beginBlock; i < endBlock; i++ {
 			b.bitBlock[i] = ^b.bitBlock[i]
 		}
-		b.bitBlock[endBlock] ^= ((uint64(1) << (uint(end) % 64 + 1))-1)
+		b.bitBlock[endBlock] ^= ((uint64(1) << (uint(end)%64 + 1)) - 1)
 	}
 }
 
@@ -143,13 +165,13 @@ func (b *Bitmap) clearRange(begin, end uint32) {
 	beginBlock := begin / 64
 	endBlock := end / 64
 	if beginBlock == endBlock {
-		b.bitBlock[beginBlock] &= (((uint64(1) << uint(begin%64)) - 1) | ^((uint64(1) << (uint(end) % 64 + 1)) - 1))
+		b.bitBlock[beginBlock] &= (((uint64(1) << uint(begin%64)) - 1) | ^((uint64(1) << (uint(end)%64 + 1)) - 1))
 	} else {
 		b.bitBlock[beginBlock] &= ((uint64(1) << uint(begin%64)) - 1)
 		for i := beginBlock + 1; i < endBlock; i++ {
 			b.bitBlock[i] &= ^b.bitBlock[i]
 		}
-		b.bitBlock[endBlock] &= ^((uint64(1) << (uint(end) % 64 + 1)) - 1)
+		b.bitBlock[endBlock] &= ^((uint64(1) << (uint(end)%64 + 1)) - 1)
 	}
 }
 
@@ -198,7 +220,7 @@ func (b *Bitmap) ToSlice() []uint32 {
 			// Extract the lowest set 1 bit.
 			j := bitBlock & -bitBlock
 			// Interpret the bit as the in32 number it represents and add it to result.
-			bitmapSlice = append(bitmapSlice, uint32((base + int(bits.OnesCount64(j - 1)))))
+			bitmapSlice = append(bitmapSlice, uint32((base + int(bits.OnesCount64(j-1)))))
 			bitBlock ^= j
 		}
 		base += 64
diff --git a/pkg/bitmap/bitmap_test.go b/pkg/bitmap/bitmap_test.go
index a7c337b55..37f068438 100644
--- a/pkg/bitmap/bitmap_test.go
+++ b/pkg/bitmap/bitmap_test.go
@@ -15,6 +15,7 @@
 package bitmap
 
 import (
+	"math"
 	"reflect"
 	"testing"
 )
@@ -29,9 +30,9 @@ func generateFilledSlice(min, max, length int) []uint32 {
 	}
 	randSlice := make([]uint32, length)
 	if length != 0 {
-		rangeNum := uint32((max - min)/length)
+		rangeNum := uint32((max - min) / length)
 		randSlice[0], randSlice[length-1] = uint32(min), uint32(max)
-		for i := 1 ; i < length - 1; i++ {
+		for i := 1; i < length-1; i++ {
 			randSlice[i] = randSlice[i-1] + rangeNum
 		}
 	}
@@ -50,9 +51,9 @@ func generateFilledBitmap(min, max, fillNum int) ([]uint32, Bitmap) {
 }
 
 func TestNewBitmap(t *testing.T) {
-	tests :=  []struct {
-		name string
-		size int
+	tests := []struct {
+		name       string
+		size       int
 		expectSize int
 	}{
 		{"length 1", 1, 1},
@@ -62,7 +63,7 @@ func TestNewBitmap(t *testing.T) {
 
 	for _, tt := range tests {
 		tt := tt
-		t.Run(tt.name, func(t *testing.T){
+		t.Run(tt.name, func(t *testing.T) {
 			if bitmap := BitmapWithSize(uint32(tt.size)); len(bitmap.bitBlock) != tt.expectSize {
 				t.Errorf("BitmapWithSize created bitmap with %v, bitBlock size: %d, wanted: %d", tt.name, len(bitmap.bitBlock), tt.expectSize)
 			}
@@ -72,14 +73,14 @@ func TestNewBitmap(t *testing.T) {
 
 func TestAdd(t *testing.T) {
 	tests := []struct {
-		name string
+		name       string
 		bitmapSize int
-		addNum int
+		addNum     int
 	}{
 		{"Add with null bitmap.bitBlock", 0, 10},
-		{"Add without extending bitBlock",64, 10},
+		{"Add without extending bitBlock", 64, 10},
 		{"Add without extending bitblock with margin number", 63, 64},
-		{"Add with extended one block",1024, 1025},
+		{"Add with extended one block", 1024, 1025},
 		{"Add with extended more then one block", 1024, 2048},
 	}
 
@@ -110,14 +111,14 @@ func TestRemove(t *testing.T) {
 	}
 	bitmapSlice := bitmap.ToSlice()
 	if !reflect.DeepEqual(bitmapSlice, secondSlice) {
-		 t.Errorf("After Remove() firstSlice, remained slice: %v, wanted: %v", bitmapSlice, secondSlice)
+		t.Errorf("After Remove() firstSlice, remained slice: %v, wanted: %v", bitmapSlice, secondSlice)
 	}
 
 	for i := 0; i < 50; i++ {
 		bitmap.Remove(secondSlice[i])
 	}
 	bitmapSlice = bitmap.ToSlice()
-	emptySlice := make([]uint32,0)
+	emptySlice := make([]uint32, 0)
 	if !reflect.DeepEqual(bitmapSlice, emptySlice) {
 		t.Errorf("After Remove secondSlice, remained slice: %v, wanted: %v", bitmapSlice, emptySlice)
 	}
@@ -127,9 +128,9 @@ func TestRemove(t *testing.T) {
 // Verify flip bits within one bitBlock, one bit and bits cross multi bitBlocks.
 func TestFlipRange(t *testing.T) {
 	tests := []struct {
-		name string
-		flipRangeMin int
-		flipRangeMax int
+		name           string
+		flipRangeMin   int
+		flipRangeMax   int
 		filledSliceLen int
 	}{
 		{"Flip one number in bitmap", 77, 77, 1},
@@ -141,7 +142,7 @@ func TestFlipRange(t *testing.T) {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
 			fillSlice, bitmap := generateFilledBitmap(tt.flipRangeMin, tt.flipRangeMax, tt.filledSliceLen)
-			flipFillSlice := make([]uint32,0)
+			flipFillSlice := make([]uint32, 0)
 			for i, j := tt.flipRangeMin, 0; i <= tt.flipRangeMax; i++ {
 				if uint32(i) != fillSlice[j] {
 					flipFillSlice = append(flipFillSlice, uint32(i))
@@ -150,7 +151,7 @@ func TestFlipRange(t *testing.T) {
 				}
 			}
 
-			bitmap.FlipRange(uint32(tt.flipRangeMin), uint32(tt.flipRangeMax + 1))
+			bitmap.FlipRange(uint32(tt.flipRangeMin), uint32(tt.flipRangeMax+1))
 			flipBitmapSlice := bitmap.ToSlice()
 			if !reflect.DeepEqual(flipFillSlice, flipBitmapSlice) {
 				t.Errorf("%v, flipped slice: %v, wanted: %v", tt.name, flipBitmapSlice, flipFillSlice)
@@ -162,10 +163,10 @@ func TestFlipRange(t *testing.T) {
 // Verify clear bits within one bitBlock, one bit and bits cross multi bitBlocks.
 func TestClearRange(t *testing.T) {
 	tests := []struct {
-		name string
+		name          string
 		clearRangeMin int
 		clearRangeMax int
-		bitmapSize int
+		bitmapSize    int
 	}{
 		{"ClearRange clear one number", 5, 5, 64},
 		{"ClearRange clear numbers within one bitBlock", 4, 61, 64},
@@ -177,9 +178,9 @@ func TestClearRange(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			bitmap := BitmapWithSize(uint32(tt.bitmapSize))
 			bitmap.FlipRange(uint32(0), uint32(tt.bitmapSize))
-			bitmap.ClearRange(uint32(tt.clearRangeMin), uint32(tt.clearRangeMax + 1))
+			bitmap.ClearRange(uint32(tt.clearRangeMin), uint32(tt.clearRangeMax+1))
 			clearedBitmapSlice := bitmap.ToSlice()
-			clearedSlice := make([]uint32,0)
+			clearedSlice := make([]uint32, 0)
 			for i := 0; i < tt.bitmapSize; i++ {
 				if i < tt.clearRangeMin || i > tt.clearRangeMax {
 					clearedSlice = append(clearedSlice, uint32(i))
@@ -218,22 +219,22 @@ func TestMinimum(t *testing.T) {
 func TestMaximum(t *testing.T) {
 	randSlice, bitmap := generateFilledBitmap(0, 1024, 200)
 	max := bitmap.Maximum()
-	if max != randSlice[len(randSlice) - 1] {
-		t.Errorf("Maximum() returns: %v, wanted: %v", max, randSlice[len(randSlice) - 1])
+	if max != randSlice[len(randSlice)-1] {
+		t.Errorf("Maximum() returns: %v, wanted: %v", max, randSlice[len(randSlice)-1])
 	}
 
 	bitmap.ClearRange(uint32(1000), uint32(1025))
 	max = bitmap.Maximum()
 	bitmapSlice := bitmap.ToSlice()
-	if max != bitmapSlice[len(bitmapSlice) - 1] {
-		t.Errorf("After ClearRange, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice) - 1])
+	if max != bitmapSlice[len(bitmapSlice)-1] {
+		t.Errorf("After ClearRange, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice)-1])
 	}
 
 	bitmap.FlipRange(uint32(1001), uint32(1021))
 	max = bitmap.Maximum()
 	bitmapSlice = bitmap.ToSlice()
-	if max != bitmapSlice[len(bitmapSlice) - 1] {
-		t.Errorf("After Flip, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice) - 1])
+	if max != bitmapSlice[len(bitmapSlice)-1] {
+		t.Errorf("After Flip, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice)-1])
 	}
 }
 
@@ -261,7 +262,7 @@ func TestBitmapNumOnes(t *testing.T) {
 	}
 
 	// Add 10 number.
-	for i := 1080; i < 1090; i++{
+	for i := 1080; i < 1090; i++ {
 		bitmap.Add(uint32(i))
 	}
 	bitmapOnes = bitmap.GetNumOnes()
@@ -270,7 +271,7 @@ func TestBitmapNumOnes(t *testing.T) {
 	}
 
 	// Add the 10 number again, the length supposed not change.
-	for i := 1080; i < 1090; i++{
+	for i := 1080; i < 1090; i++ {
 		bitmap.Add(uint32(i))
 	}
 	bitmapOnes = bitmap.GetNumOnes()
@@ -292,3 +293,14 @@ func TestBitmapNumOnes(t *testing.T) {
 		t.Errorf("After ClearRange, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 20)
 	}
 }
+
+func TestFirstZero(t *testing.T) {
+	bitmap := BitmapWithSize(uint32(1000))
+	bitmap.FlipRange(200, 400)
+	for i, j := range map[uint32]uint32{0: 0, 201: 400, 200: 400, 199: 199, 400: 400, 10000: math.MaxInt32} {
+		v := bitmap.FirstZero(i)
+		if v != j {
+			t.Errorf("Minimum() returns: %v, wanted: %v", v, j)
+		}
+	}
+}
diff --git a/pkg/sentry/kernel/fd_table.go b/pkg/sentry/kernel/fd_table.go
index 5648be11b..9f7702fcc 100644
--- a/pkg/sentry/kernel/fd_table.go
+++ b/pkg/sentry/kernel/fd_table.go
@@ -271,7 +271,7 @@ func (f *FDTable) NewFDs(ctx context.Context, minFD int32, files []*fs.File, fla
 		if lim.Cur != limits.Infinity {
 			end = int32(lim.Cur)
 		}
-		if minFD >= end {
+		if minFD+int32(len(files)) > end {
 			return nil, unix.EMFILE
 		}
 	}
@@ -280,17 +280,10 @@ func (f *FDTable) NewFDs(ctx context.Context, minFD int32, files []*fs.File, fla
 
 	// max is used as the largest number in fdBitmap + 1.
 	max := int32(0)
-	// flipFdBitmap is flip of fdBitmap which is used to find free fd in fdBitmap.
-	flipFdBitmap := f.fdBitmap.Clone()
 
 	if !f.fdBitmap.IsEmpty() {
 		max = int32(f.fdBitmap.Maximum())
 		max++
-		flipFdBitmap.FlipRange(uint32(0), uint32(max))
-		// Clear 0 to (minFD-1) bit in flipFdBitmap, thus it will not use fd that is less than minFD.
-		if minFD > 0 {
-			flipFdBitmap.ClearRange(uint32(0), uint32(minFD))
-		}
 	}
 
 	// Adjust max in case it is less than minFD.
@@ -301,20 +294,18 @@ func (f *FDTable) NewFDs(ctx context.Context, minFD int32, files []*fs.File, fla
 	for len(fds) < len(files) {
 		// Try to use free bit in fdBitmap.
 		// If all bits in fdBitmap are used, expand fd to the max.
-		if !flipFdBitmap.IsEmpty() {
-			fd := flipFdBitmap.Minimum()
-			f.fdBitmap.Add(fd)
-			flipFdBitmap.Remove(fd)
-			f.set(ctx, int32(fd), files[len(fds)], flags)
-			fds = append(fds, int32(fd))
-		} else if max < end {
-			f.fdBitmap.Add(uint32(max))
-			f.set(ctx, max, files[len(fds)], flags)
-			fds = append(fds, max)
+		fd := f.fdBitmap.FirstZero(uint32(minFD))
+		if fd == math.MaxInt32 {
+			fd = uint32(max)
 			max++
-		} else {
+		}
+		if fd >= uint32(end) {
 			break
 		}
+		f.fdBitmap.Add(fd)
+		f.set(ctx, int32(fd), files[len(fds)], flags)
+		fds = append(fds, int32(fd))
+		minFD = int32(fd)
 	}
 
 	// Failure? Unwind existing FDs.
@@ -366,17 +357,10 @@ func (f *FDTable) NewFDsVFS2(ctx context.Context, minFD int32, files []*vfs.File
 
 	// max is used as the largest number in fdBitmap + 1.
 	max := int32(0)
-	// flipFdBitmap is flip of fdBitmap which is used to find free fd in fdBitmap.
-	flipFdBitmap := f.fdBitmap.Clone()
 
 	if !f.fdBitmap.IsEmpty() {
 		max = int32(f.fdBitmap.Maximum())
 		max++
-		flipFdBitmap.FlipRange(uint32(0), uint32(max))
-		// Clear 0 to (minFD-1) bit in flipFdBitmap, thus it will not use fd that is less than minFD.
-		if minFD > 0 {
-			flipFdBitmap.ClearRange(uint32(0), uint32(minFD))
-		}
 	}
 
 	// Adjust max in case it is less than minFD.
@@ -387,20 +371,18 @@ func (f *FDTable) NewFDsVFS2(ctx context.Context, minFD int32, files []*vfs.File
 	for len(fds) < len(files) {
 		// Try to use free bit in fdBitmap.
 		// If all bits in fdBitmap are used, expand fd to the max.
-		if !flipFdBitmap.IsEmpty() {
-			fd := flipFdBitmap.Minimum()
-			f.fdBitmap.Add(fd)
-			flipFdBitmap.Remove(fd)
-			f.setVFS2(ctx, int32(fd), files[len(fds)], flags)
-			fds = append(fds, int32(fd))
-		} else if max < end {
-			f.fdBitmap.Add(uint32(max))
-			f.setVFS2(ctx, max, files[len(fds)], flags)
-			fds = append(fds, max)
+		fd := f.fdBitmap.FirstZero(uint32(minFD))
+		if fd == math.MaxInt32 {
+			fd = uint32(max)
 			max++
-		} else {
+		}
+		if fd >= uint32(end) {
 			break
 		}
+		f.fdBitmap.Add(fd)
+		f.setVFS2(ctx, int32(fd), files[len(fds)], flags)
+		fds = append(fds, int32(fd))
+		minFD = int32(fd)
 	}
 	// Failure? Unwind existing FDs.
 	if len(fds) < len(files) {