fdtable: avoid large arrays

FDTable.descriptorTable is a slice of unsafe.Pointer-s and its maximum length
is MaxInt32. It requires up to 16GB of memory. A process can use just a few
descriptors but sets one or more of them to high numbers. In this case,
FDTable.descriptorTable is extended to the maximum size.

The problem here is that go-runtime zeros memory regions when they are reused.
In the case of fdtable, the memory region is 16GB, so it is a time consuming
operation. Second, it forces the kernel to allocate physical pages to
the entire region.

This change adds another level to descriptorTable, so the first level is
a slice of buckets where each bucket is a slice of descriptors. The bucket
size is fixed to 512 entries to fit one page.

Before:
BenchmarkFDLookupAndDecRef-12              	50834290	        23.70 ns/op
BenchmarkCreateWithMaxFD-12                	       2	7194873988 ns/op
BenchmarkFDLookupAndDecRefConcurrent-12    	23775555	        49.68 ns/op
BenchmarkTableLookup-12                    	412888780	         2.835 ns/op
BenchmarkTableMapLookup-12                 	87944782	        12.84 ns/op

After:
BenchmarkFDLookupAndDecRef-12              	46229940	        25.03 ns/op
BenchmarkCreateWithMaxFD-12                	      13	  82573899 ns/op
BenchmarkFDLookupAndDecRefConcurrent-12    	21889380	        54.13 ns/op
BenchmarkTableLookup-12                    	415851230	         2.821 ns/op
BenchmarkTableMapLookup-12                 	97236267	        11.89 ns/op

Reported-by: syzbot+af17678e3bfb7ca7c65a@syzkaller.appspotmail.com
PiperOrigin-RevId: 539138632

pkg/bitmap/bitmap.go

@@ -250,6 +250,38 @@ func (b *Bitmap) FlipRange(begin, end uint32) {
 	}
 }
 
+// ForEach calls `f` for each set bit in the range [start, end).
+//
+// If f returns false, ForEach stops the iteration.
+func (b *Bitmap) ForEach(start, end uint32, f func(idx uint32) bool) {
+	blockEnd := (end + 63) / 64
+	if blockEnd > uint32(len(b.bitBlock)) {
+		blockEnd = uint32(len(b.bitBlock))
+	}
+	// base is the start number of a bitBlock
+	base := start / 64 * 64
+	blockMask := ^((uint64(1) << (start % 64)) - 1)
+	for i := start / 64; i < blockEnd; i++ {
+		if i == end/64 {
+			blockMask &= (uint64(1) << (end % 64)) - 1
+		}
+		bitBlock := b.bitBlock[i] & blockMask
+		blockMask = ^uint64(0)
+		// Iterate through all the numbers held by this bit block.
+		for bitBlock != 0 {
+			// Extract the lowest set 1 bit.
+			j := bitBlock & -bitBlock
+			// Interpret the bit as the in32 number it represents and add it to result.
+			idx := base + uint32(bits.OnesCount64(j-1))
+			if !f(idx) {
+				return
+			}
+			bitBlock ^= j
+		}
+		base += 64
+	}
+}
+
 // ToSlice transform the Bitmap into slice. For example, a bitmap of [0, 1, 0, 1]
 // will return the slice [1, 3].
 func (b *Bitmap) ToSlice() []uint32 {

pkg/bitmap/bitmap_test.go

@@ -294,6 +294,56 @@ func TestBitmapNumOnes(t *testing.T) {
 	}
 }
 
+type bitmapForEachTestcase struct {
+	start, end uint32
+	expected   map[uint32]bool
+}
+
+func TestForEach(t *testing.T) {
+	const bitmapSize = 1 << 20
+	bitmap := New(bitmapSize)
+	bitmap.FlipRange(200, 400)
+	bitmap.FlipRange(1003, 1004)
+	bitmap.FlipRange(1005, 1006)
+	bitmap.FlipRange(1096, 1098)
+	testcases := []bitmapForEachTestcase{
+		{0, 0, map[uint32]bool{}},
+		{0, 100, map[uint32]bool{}},
+		{1003, 1004, map[uint32]bool{1003: true}},
+		{1003, 1006, map[uint32]bool{1003: true, 1005: true}},
+		{1000, 2000, map[uint32]bool{1003: true, 1005: true, 1096: true, 1097: true}},
+		{1000, 1097, map[uint32]bool{1003: true, 1005: true, 1096: true}},
+		{1060, 1097, map[uint32]bool{1096: true}},
+		{0, bitmapSize, func() map[uint32]bool {
+			m := make(map[uint32]bool)
+			for _, i := range bitmap.ToSlice() {
+				m[i] = true
+			}
+			return m
+		}()},
+		{234, 356, func() map[uint32]bool {
+			m := make(map[uint32]bool)
+			for i := uint32(234); i < 356; i++ {
+				m[i] = true
+			}
+			return m
+		}()},
+	}
+	for _, tc := range testcases {
+		bitmap.ForEach(tc.start, tc.end, func(idx uint32) bool {
+			if _, ok := tc.expected[idx]; !ok {
+				t.Errorf("[%d, %d): unexpeced index: %d", tc.start, tc.end, idx)
+				return false
+			}
+			delete(tc.expected, idx)
+			return true
+		})
+		if len(tc.expected) != 0 {
+			t.Errorf("[%d-%d): leftover: %#v", tc.start, tc.end, tc.expected)
+		}
+	}
+}
+
 type BitmapGetFirstTestcase struct {
 	queryValue    uint32
 	expectedValue uint32

pkg/sentry/kernel/BUILD

@@ -7,6 +7,39 @@ package(
     licenses = ["notice"],
 )
 
+go_template_instance(
+    name = "atomicptr_bucket_slice",
+    out = "atomicptr_bucket_slice_unsafe.go",
+    package = "kernel",
+    prefix = "descriptorBucketSlice",
+    template = "//pkg/sync/atomicptr:generic_atomicptr",
+    types = {
+        "Value": "descriptorBucketSlice",
+    },
+)
+
+go_template_instance(
+    name = "atomicptr_bucket",
+    out = "atomicptr_bucket_unsafe.go",
+    package = "kernel",
+    prefix = "descriptorBucket",
+    template = "//pkg/sync/atomicptr:generic_atomicptr",
+    types = {
+        "Value": "descriptorBucket",
+    },
+)
+
+go_template_instance(
+    name = "atomicptr_descriptor",
+    out = "atomicptr_descriptor_unsafe.go",
+    package = "kernel",
+    prefix = "descriptor",
+    template = "//pkg/sync/atomicptr:generic_atomicptr",
+    types = {
+        "Value": "descriptor",
+    },
+)
+
 declare_mutex(
     name = "cpu_clock_mutex",
     out = "cpu_clock_mutex.go",
@@ -205,6 +238,9 @@ go_library(
     srcs = [
         "abstract_socket_namespace.go",
         "aio.go",
+        "atomicptr_bucket_slice_unsafe.go",
+        "atomicptr_bucket_unsafe.go",
+        "atomicptr_descriptor_unsafe.go",
         "cgroup.go",
         "cgroup_mutex.go",
         "context.go",

pkg/sentry/kernel/fd_table.go

@@ -159,32 +159,22 @@ func (f *FDTable) DecRef(ctx context.Context) {
 //
 // It is the caller's responsibility to acquire an appropriate lock.
 func (f *FDTable) forEachUpTo(ctx context.Context, maxFd int32, fn func(fd int32, file *vfs.FileDescription, flags FDFlags)) {
-	// retries tracks the number of failed TryIncRef attempts for the same FD.
-	retries := 0
-	fds := f.fdBitmap.ToSlice()
 	// Iterate through the fdBitmap.
-	for _, ufd := range fds {
+	f.fdBitmap.ForEach(0, uint32(maxFd), func(ufd uint32) bool {
 		fd := int32(ufd)
-		if fd >= maxFd {
-			break
-		}
 		file, flags, ok := f.get(fd)
 		if !ok {
-			break
+			return true
 		}
 		if file != nil {
 			if !file.TryIncRef() {
-				retries++
-				if retries > 1000 {
-					panic(fmt.Sprintf("File in FD table has been destroyed. FD: %d, File: %+v, Impl: %+v", fd, file, file.Impl()))
-				}
-				continue // Race caught.
+				return true
 			}
 			fn(fd, file, flags)
 			file.DecRef(ctx)
 		}
-		retries = 0
-	}
+		return true
+	})
 }
 
 // forEach iterates over all non-nil files upto maxFd in sorted order.

pkg/sentry/kernel/fd_table_test.go

@@ -219,6 +219,61 @@ func BenchmarkFDLookupAndDecRef(b *testing.B) {
 	})
 }
 
+func BenchmarkNewFDAt(b *testing.B) {
+	const maxLimit = 1 << 31
+	b.StopTimer() // Setup.
+
+	runTest(b, func(ctx context.Context, fdTable *FDTable, fd *vfs.FileDescription, limitSet *limits.LimitSet) {
+		// Remove the previous limit.
+		limitSet.Set(limits.NumberOfFiles, limits.Limit{maxLimit, maxLimit}, true)
+
+		b.StartTimer() // Benchmark.
+		for i := 0; i < b.N; i++ {
+			err := fdTable.NewFDAt(ctx, int32(i%maxLimit), fd, FDFlags{})
+			if err != nil {
+				b.Fatalf("fdTable.NewFDAt: got %v, wanted nil", err)
+			}
+		}
+	})
+}
+
+func BenchmarkFork(b *testing.B) {
+	b.StopTimer() // Setup.
+
+	runTest(b, func(ctx context.Context, fdTable *FDTable, fd *vfs.FileDescription, limitSet *limits.LimitSet) {
+		for i := 0; i < maxFD; i++ {
+			err := fdTable.NewFDAt(ctx, int32(i), fd, FDFlags{})
+			if err != nil {
+				b.Fatalf("fdTable.NewFDs: got %v, wanted nil", err)
+			}
+		}
+
+		b.StartTimer() // Benchmark.
+		for i := 0; i < b.N; i++ {
+			t := fdTable.Fork(ctx, maxFD)
+			t.DecRef(ctx)
+		}
+	})
+}
+
+func BenchmarkCreateWithMaxFD(b *testing.B) {
+	const maxLimit = 1 << 31
+	runTest(b, func(ctx context.Context, _ *FDTable, fd *vfs.FileDescription, limitSet *limits.LimitSet) {
+		// Remove the previous limit.
+		limitSet.Set(limits.NumberOfFiles, limits.Limit{maxLimit, maxLimit}, true)
+
+		for i := 0; i < b.N; i++ {
+			fdTable := new(FDTable)
+			fdTable.init()
+			err := fdTable.NewFDAt(ctx, maxLimit-1, fd, FDFlags{})
+			if err != nil {
+				b.Fatalf("fdTable.NewFDs: got %v, wanted nil", err)
+			}
+			fdTable.DecRef(ctx)
+		}
+	})
+}
+
 func BenchmarkFDLookupAndDecRefConcurrent(b *testing.B) {
 	b.StopTimer() // Setup.
 

pkg/sentry/kernel/fd_table_unsafe.go

@@ -16,19 +16,21 @@ package kernel
 
 import (
 	"math"
-	"sync/atomic"
-	"unsafe"
 
 	"gvisor.dev/gvisor/pkg/bitmap"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
+type descriptorBucket [fdsPerBucket]descriptorAtomicPtr
+type descriptorBucketSlice []descriptorBucketAtomicPtr
+
+// descriptorTable is a two level table. The first level is a slice of
+// *descriptorBucket where each bucket is a slice of *descriptor.
+//
+// All objects are updated atomically.
 type descriptorTable struct {
-	// slice is a *[]unsafe.Pointer, where each element is actually
-	// *descriptor object, updated atomically.
-	//
 	// Changes to the slice itself requiring holding FDTable.mu.
-	slice unsafe.Pointer `state:".(map[int32]*descriptor)"`
+	slice descriptorBucketSliceAtomicPtr `state:".(map[int32]*descriptor)"`
 }
 
 // initNoLeakCheck initializes the table without enabling leak checking.
@@ -36,8 +38,8 @@ type descriptorTable struct {
 // This is used when loading an FDTable after S/R, during which the ref count
 // object itself will enable leak checking if necessary.
 func (f *FDTable) initNoLeakCheck() {
-	var slice []unsafe.Pointer // Empty slice.
-	atomic.StorePointer(&f.slice, unsafe.Pointer(&slice))
+	var slice descriptorBucketSlice // Empty slice.
+	f.slice.Store(&slice)
 }
 
 // init initializes the table with leak checking.
@@ -47,17 +49,30 @@ func (f *FDTable) init() {
 	f.fdBitmap = bitmap.New(uint32(math.MaxUint16))
 }
 
+const (
+	// fdsPerBucketShift is chosen in such a way that the size of bucket is
+	// equal to one page.
+	fdsPerBucketShift = 9
+	fdsPerBucket      = 1 << fdsPerBucketShift
+	fdsPerBucketMask  = fdsPerBucket - 1
+)
+
 // get gets a file entry.
 //
 // The boolean indicates whether this was in range.
 //
 //go:nosplit
 func (f *FDTable) get(fd int32) (*vfs.FileDescription, FDFlags, bool) {
-	slice := *(*[]unsafe.Pointer)(atomic.LoadPointer(&f.slice))
-	if fd >= int32(len(slice)) {
+	slice := *f.slice.Load()
+	bucketN := fd >> fdsPerBucketShift
+	if bucketN >= int32(len(slice)) {
 		return nil, FDFlags{}, false
 	}
-	d := (*descriptor)(atomic.LoadPointer(&slice[fd]))
+	bucket := slice[bucketN].Load()
+	if bucket == nil {
+		return nil, FDFlags{}, false
+	}
+	d := bucket[fd&fdsPerBucketMask].Load()
 	if d == nil {
 		return nil, FDFlags{}, true
 	}
@@ -67,8 +82,8 @@ func (f *FDTable) get(fd int32) (*vfs.FileDescription, FDFlags, bool) {
 // CurrentMaxFDs returns the number of file descriptors that may be stored in f
 // without reallocation.
 func (f *FDTable) CurrentMaxFDs() int {
-	slice := *(*[]unsafe.Pointer)(atomic.LoadPointer(&f.slice))
-	return len(slice)
+	slice := *f.slice.Load()
+	return len(slice) * fdsPerBucket
 }
 
 // set sets the file description referred to by fd to file. If
@@ -79,11 +94,12 @@ func (f *FDTable) CurrentMaxFDs() int {
 //
 // Precondition: mu must be held.
 func (f *FDTable) set(fd int32, file *vfs.FileDescription, flags FDFlags) *vfs.FileDescription {
-	slicePtr := (*[]unsafe.Pointer)(atomic.LoadPointer(&f.slice))
+	slicePtr := f.slice.Load()
 
+	bucketN := fd >> fdsPerBucketShift
 	// Grow the table as required.
-	if length := len(*slicePtr); int(fd) >= length {
-		newLen := int(fd) + 1
+	if length := len(*slicePtr); int(bucketN) >= length {
+		newLen := int(bucketN) + 1
 		if newLen < 2*length {
 			// Ensure the table at least doubles in size without going over the limit.
 			newLen = 2 * length
@@ -91,13 +107,19 @@ func (f *FDTable) set(fd int32, file *vfs.FileDescription, flags FDFlags) *vfs.F
 				newLen = int(MaxFdLimit)
 			}
 		}
-		newSlice := append(*slicePtr, make([]unsafe.Pointer, newLen-length)...)
+		newSlice := append(*slicePtr, make([]descriptorBucketAtomicPtr, newLen-length)...)
 		slicePtr = &newSlice
-		atomic.StorePointer(&f.slice, unsafe.Pointer(slicePtr))
+		f.slice.Store(slicePtr)
 	}
 
 	slice := *slicePtr
 
+	bucket := slice[bucketN].Load()
+	if bucket == nil {
+		bucket = &descriptorBucket{}
+		slice[bucketN].Store(bucket)
+	}
+
 	var desc *descriptor
 	if file != nil {
 		desc = &descriptor{
@@ -107,7 +129,7 @@ func (f *FDTable) set(fd int32, file *vfs.FileDescription, flags FDFlags) *vfs.F
 	}
 
 	// Update the single element.
-	orig := (*descriptor)(atomic.SwapPointer(&slice[fd], unsafe.Pointer(desc)))
+	orig := bucket[fd%fdsPerBucket].Swap(desc)
 
 	// Acquire a table reference.
 	if desc != nil && desc.file != nil {

pkg/sync/atomicptr/generic_atomicptr_unsafe.go

@@ -47,3 +47,8 @@ func (p *AtomicPtr) Load() *Value {
 func (p *AtomicPtr) Store(x *Value) {
 	atomic.StorePointer(&p.ptr, (unsafe.Pointer)(x))
 }
+
+// Swap atomically stores `x` into *p and returns the previous *p value.
+func (p *AtomicPtr) Swap(x *Value) *Value {
+	return (*Value)(atomic.SwapPointer(&p.ptr, (unsafe.Pointer)(x)))
+}