Repository: tetratelabs/wazero
Pull request #1156 (2023-02-23 23:51:12)

Fix issue with a delayed check of the context cancelation
If a context is canceled after a guest func call has returned it could cause the module to close in some cases. This change ensures that a delayed check of the context cancelation is ignored.

I've also reduced the cost of context cancelation a bit.


Patch:
diff --git a/internal/sys/file_table.go b/internal/descriptor/table.go
similarity index 51%
rename from internal/sys/file_table.go
rename to internal/descriptor/table.go
index 0c85cf19..0ed376f2 100644
--- a/internal/sys/file_table.go
+++ b/internal/descriptor/table.go
@@ -1,26 +1,26 @@
-package sys
+package descriptor
 
 import "math/bits"
 
-// FileTable is a data structure mapping 32 bit file descriptor to file objects.
+// Table is a data structure mapping 32 bit descriptor to items.
 //
 // The data structure optimizes for memory density and lookup performance,
 // trading off compute at insertion time. This is a useful compromise for the
-// use cases we employ it with: files are usually read or written a lot more
-// often than they are opened, each operation requires a table lookup so we are
-// better off spending extra compute to insert files in the table in order to
+// use cases we employ it with: items are usually accessed a lot more often
+// than they are inserted, each operation requires a table lookup so we are
+// better off spending extra compute to insert items in the table in order to
 // get cheaper lookups. Memory efficiency is also crucial to support scaling
-// with programs that open thousands of files: having a high or non-linear
+// with programs that maintain thousands of items: having a high or non-linear
 // memory-to-item ratio could otherwise be used as an attack vector by malicous
 // applications attempting to damage performance of the host.
-type FileTable struct {
+type Table[Key ~uint32, Item any] struct {
 	masks []uint64
-	files []*FileEntry
+	items []Item
 }
 
-// Len returns the number of files stored in the table.
-func (t *FileTable) Len() (n int) {
-	// We could make this a O(1) operation if we cached the number of files in
+// Len returns the number of items stored in the table.
+func (t *Table[Key, Item]) Len() (n int) {
+	// We could make this a O(1) operation if we cached the number of items in
 	// the table. More state usually means more problems, so until we have a
 	// clear need for this, the simple implementation may be a better trade off.
 	for _, mask := range t.masks {
@@ -29,9 +29,9 @@ func (t *FileTable) Len() (n int) {
 	return n
 }
 
-// Grow ensures that t has enough room for n files, potentially reallocating the
-// internal buffers if their capacity was too small to hold this many files.
-func (t *FileTable) Grow(n int) {
+// Grow ensures that t has enough room for n items, potentially reallocating the
+// internal buffers if their capacity was too small to hold this many items.
+func (t *Table[Key, Item]) Grow(n int) {
 	// Round up to a multiple of 64 since this is the smallest increment due to
 	// using 64 bits masks.
 	n = (n*64 + 63) / 64
@@ -40,20 +40,20 @@ func (t *FileTable) Grow(n int) {
 		masks := make([]uint64, n)
 		copy(masks, t.masks)
 
-		files := make([]*FileEntry, n*64)
-		copy(files, t.files)
+		items := make([]Item, n*64)
+		copy(items, t.items)
 
 		t.masks = masks
-		t.files = files
+		t.items = items
 	}
 }
 
-// Insert inserts the given file to the table, returning the fd that it is
+// Insert inserts the given item to the table, returning the key that it is
 // mapped to.
 //
-// The method does not perform deduplication, it is possible for the same file
-// to be inserted multiple times, each insertion will return a different fd.
-func (t *FileTable) Insert(file *FileEntry) (fd uint32) {
+// The method does not perform deduplication, it is possible for the same item
+// to be inserted multiple times, each insertion will return a different key.
+func (t *Table[Key, Item]) Insert(item Item) (key Key) {
 	offset := 0
 insert:
 	// Note: this loop could be made a lot more efficient using vectorized
@@ -63,10 +63,10 @@ insert:
 		if ^mask != 0 { // not full?
 			shift := bits.TrailingZeros64(^mask)
 			index += offset
-			fd = uint32(index)*64 + uint32(shift)
-			t.files[fd] = file
+			key = Key(index)*64 + Key(shift)
+			t.items[key] = item
 			t.masks[index] = mask | uint64(1<<shift)
-			return fd
+			return key
 		}
 	}
 
@@ -80,52 +80,53 @@ insert:
 	goto insert
 }
 
-// Lookup returns the file associated with the given fd (may be nil).
-func (t *FileTable) Lookup(fd uint32) (file *FileEntry, found bool) {
-	if i := int(fd); i >= 0 && i < len(t.files) {
-		index := uint(fd) / 64
-		shift := uint(fd) % 64
+// Lookup returns the item associated with the given key (may be nil).
+func (t *Table[Key, Item]) Lookup(key Key) (item Item, found bool) {
+	if i := int(key); i >= 0 && i < len(t.items) {
+		index := uint(key) / 64
+		shift := uint(key) % 64
 		if (t.masks[index] & (1 << shift)) != 0 {
-			file, found = t.files[i], true
+			item, found = t.items[i], true
 		}
 	}
 	return
 }
 
-// InsertAt inserts the given `file` at the file descriptor `fd`.
-func (t *FileTable) InsertAt(file *FileEntry, fd uint32) {
-	if diff := int(fd) - t.Len(); diff > 0 {
+// InsertAt inserts the given `item` at the item descriptor `key`.
+func (t *Table[Key, Item]) InsertAt(item Item, key Key) {
+	if diff := int(key) - t.Len(); diff > 0 {
 		t.Grow(diff)
 	}
-	index := uint(fd) / 64
-	shift := uint(fd) % 64
+	index := uint(key) / 64
+	shift := uint(key) % 64
 	t.masks[index] |= 1 << shift
-	t.files[fd] = file
+	t.items[key] = item
 }
 
-// Delete deletes the file stored at the given fd from the table.
-func (t *FileTable) Delete(fd uint32) {
-	if index, shift := fd/64, fd%64; int(index) < len(t.masks) {
+// Delete deletes the item stored at the given key from the table.
+func (t *Table[Key, Item]) Delete(key Key) {
+	if index, shift := key/64, key%64; int(index) < len(t.masks) {
 		mask := t.masks[index]
 		if (mask & (1 << shift)) != 0 {
-			t.files[fd] = nil
+			var zero Item
+			t.items[key] = zero
 			t.masks[index] = mask & ^uint64(1<<shift)
 		}
 	}
 }
 
-// Range calls f for each file and its associated fd in the table. The function
+// Range calls f for each item and its associated key in the table. The function
 // f might return false to interupt the iteration.
-func (t *FileTable) Range(f func(uint32, *FileEntry) bool) {
+func (t *Table[Key, Item]) Range(f func(Key, Item) bool) {
 	for i, mask := range t.masks {
 		if mask == 0 {
 			continue
 		}
-		for j := uint32(0); j < 64; j++ {
+		for j := Key(0); j < 64; j++ {
 			if (mask & (1 << j)) == 0 {
 				continue
 			}
-			if fd := uint32(i)*64 + j; !f(fd, t.files[fd]) {
+			if key := Key(i)*64 + j; !f(key, t.items[key]) {
 				return
 			}
 		}
@@ -133,11 +134,12 @@ func (t *FileTable) Range(f func(uint32, *FileEntry) bool) {
 }
 
 // Reset clears the content of the table.
-func (t *FileTable) Reset() {
+func (t *Table[Key, Item]) Reset() {
 	for i := range t.masks {
 		t.masks[i] = 0
 	}
-	for i := range t.files {
-		t.files[i] = nil
+	var zero Item
+	for i := range t.items {
+		t.items[i] = zero
 	}
 }
diff --git a/internal/sys/fs.go b/internal/sys/fs.go
index 71360e39..692578f9 100644
--- a/internal/sys/fs.go
+++ b/internal/sys/fs.go
@@ -9,6 +9,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/tetratelabs/wazero/internal/descriptor"
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/sysfs"
 )
@@ -221,6 +222,10 @@ type FSContext struct {
 	openedFiles FileTable
 }
 
+// FileTable is an specialization of the descriptor.Table type used to map file
+// descriptors to file entries.
+type FileTable = descriptor.Table[uint32, *FileEntry]
+
 // NewFSContext creates a FSContext with stdio streams and an optional
 // pre-opened filesystem.
 //
diff --git a/internal/wasm/call_context.go b/internal/wasm/call_context.go
index c70c2e88..5443caf8 100644
--- a/internal/wasm/call_context.go
+++ b/internal/wasm/call_context.go
@@ -73,29 +73,36 @@ func (m *CallContext) FailIfClosed() (err error) {
 //
 // Callers of this function must invoke the returned context.CancelFunc to release the spawned Goroutine.
 func (m *CallContext) CloseModuleOnCanceledOrTimeout(ctx context.Context) context.CancelFunc {
-	goroutineDone, cancelFn := context.WithCancel(context.Background())
-	go m.closeModuleOnCanceledOrTimeoutClosure(ctx, goroutineDone)()
-	return cancelFn
+	// Creating an empty channel in this case is a bit more efficient than
+	// creating a context.Context and canceling it with the same effect. We
+	// really just need to be notified when to stop listening to the users
+	// context. Closing the channel will unblock the select in the goroutine
+	// causing it to return an stop listening to ctx.Done().
+	cancelChan := make(chan struct{})
+	go m.closeModuleOnCanceledOrTimeout(ctx, cancelChan)
+	return func() { close(cancelChan) }
 }
 
-// closeModuleOnCanceledOrTimeoutClosure is extracted from CloseModuleOnCanceledOrTimeout for testing.
-func (m *CallContext) closeModuleOnCanceledOrTimeoutClosure(ctx, goroutineDone context.Context) func() {
-	return func() {
-		for {
-			select {
-			case <-ctx.Done():
-				if errors.Is(ctx.Err(), context.Canceled) {
-					// TODO: figure out how to report error here.
-					_ = m.CloseWithExitCode(ctx, sys.ExitCodeContextCanceled)
-				} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-					// TODO: figure out how to report error here.
-					_ = m.CloseWithExitCode(ctx, sys.ExitCodeDeadlineExceeded)
-				}
-				return
-			case <-goroutineDone.Done():
-				return
+// closeModuleOnCanceledOrTimeout is extracted from CloseModuleOnCanceledOrTimeout for testing.
+func (m *CallContext) closeModuleOnCanceledOrTimeout(ctx context.Context, cancelChan <-chan struct{}) {
+	select {
+	case <-ctx.Done():
+		select {
+		case <-cancelChan:
+			// In some cases by the time this goroutine is scheduled, the caller
+			// has already closed both the context and the cancelChan. In this
+			// case go will randomize which branch of the outer select to enter
+			// and we don't want to close the module.
+		default:
+			if errors.Is(ctx.Err(), context.Canceled) {
+				// TODO: figure out how to report error here.
+				_ = m.CloseWithExitCode(ctx, sys.ExitCodeContextCanceled)
+			} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				// TODO: figure out how to report error here.
+				_ = m.CloseWithExitCode(ctx, sys.ExitCodeDeadlineExceeded)
 			}
 		}
+	case <-cancelChan:
 	}
 }
 


Test patch:
diff --git a/internal/sys/file_table_test.go b/internal/descriptor/table_test.go
similarity index 99%
rename from internal/sys/file_table_test.go
rename to internal/descriptor/table_test.go
index 050ceda1..753b583b 100644
--- a/internal/sys/file_table_test.go
+++ b/internal/descriptor/table_test.go
@@ -1,4 +1,4 @@
-package sys_test
+package descriptor_test
 
 import (
 	"testing"
diff --git a/internal/sys/fs_test.go b/internal/sys/fs_test.go
index 97a516d2..82fccd8f 100644
--- a/internal/sys/fs_test.go
+++ b/internal/sys/fs_test.go
@@ -331,13 +331,19 @@ func TestFSContext_ChangeOpenFlag(t *testing.T) {
 	// Without APPEND.
 	fd, err := c.OpenFile(dirFs, fileName, os.O_RDWR, 0o600)
 	require.NoError(t, err)
-	require.Equal(t, c.openedFiles.files[fd].openFlag&syscall.O_APPEND, 0)
+	f0, ok := c.openedFiles.Lookup(fd)
+	require.True(t, ok)
+	require.Equal(t, f0.openFlag&syscall.O_APPEND, 0)
 
 	// Set the APPEND flag.
 	require.NoError(t, c.ChangeOpenFlag(fd, syscall.O_APPEND))
-	require.Equal(t, c.openedFiles.files[fd].openFlag&syscall.O_APPEND, syscall.O_APPEND)
+	f1, ok := c.openedFiles.Lookup(fd)
+	require.True(t, ok)
+	require.Equal(t, f1.openFlag&syscall.O_APPEND, syscall.O_APPEND)
 
 	// Remove the APPEND flag.
 	require.NoError(t, c.ChangeOpenFlag(fd, 0))
-	require.Equal(t, c.openedFiles.files[fd].openFlag&syscall.O_APPEND, 0)
+	f2, ok := c.openedFiles.Lookup(fd)
+	require.True(t, ok)
+	require.Equal(t, f2.openFlag&syscall.O_APPEND, 0)
 }
diff --git a/internal/wasm/call_context_test.go b/internal/wasm/call_context_test.go
index ab758a86..2f69c535 100644
--- a/internal/wasm/call_context_test.go
+++ b/internal/wasm/call_context_test.go
@@ -361,19 +361,31 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 
 	t.Run("cancel works", func(t *testing.T) {
 		cc := &CallContext{Closed: new(uint64), module: &ModuleInstance{Name: "test"}, s: s}
-		goroutineDone, cancelFn := context.WithCancel(context.Background())
-		fn := cc.closeModuleOnCanceledOrTimeoutClosure(context.Background(), goroutineDone)
+		cancelChan := make(chan struct{})
 		var wg sync.WaitGroup
 		wg.Add(1)
 
-		// Ensure that fn returned by closeModuleOnCanceledOrTimeoutClosure exists after cancelFn is called.
+		// Ensure that fn returned by closeModuleOnCanceledOrTimeout exists after cancelFn is called.
 		go func() {
 			defer wg.Done()
-			fn()
+			cc.closeModuleOnCanceledOrTimeout(context.Background(), cancelChan)
 		}()
-		cancelFn()
+		close(cancelChan)
 		wg.Wait()
 	})
+
+	t.Run("no close on all resources canceled", func(t *testing.T) {
+		cc := &CallContext{Closed: new(uint64), module: &ModuleInstance{Name: "test"}, s: s}
+		cancelChan := make(chan struct{})
+		close(cancelChan)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		cc.closeModuleOnCanceledOrTimeout(ctx, cancelChan)
+
+		err := cc.FailIfClosed()
+		require.Nil(t, err)
+	})
 }
 
 type mockCloser struct{ called int }