wasi: optimize and benchmark fd_read (#909)

Signed-off-by: Adrian Cole <[email protected]>

imports/wasi_snapshot_preview1/fs.go

@@ -613,6 +613,7 @@ func fdReadFn(ctx context.Context, mod api.Module, stack []uint64) (nread uint32
 func fdReadOrPread(ctx context.Context, mod api.Module, stack []uint64, isPread bool) (uint32, Errno) {
 	sysCtx := mod.(*wasm.CallContext).Sys
 	mem := mod.Memory()
+
 	fd := uint32(stack[0])
 	iovs := uint32(stack[1])
 	iovsCount := uint32(stack[2])
@@ -638,16 +639,16 @@ func fdReadOrPread(ctx context.Context, mod api.Module, stack []uint64, isPread
 	}
 
 	var nread uint32
-	for i := uint32(0); i < iovsCount; i++ {
-		iov := iovs + i*8
-		offset, ok := mem.ReadUint32Le(ctx, iov)
-		if !ok {
-			return 0, ErrnoFault
-		}
-		l, ok := mem.ReadUint32Le(ctx, iov+4)
-		if !ok {
-			return 0, ErrnoFault
-		}
+	iovsStop := iovsCount << 3 // iovsCount * 8
+	iovsBuf, ok := mem.Read(ctx, iovs, iovsStop)
+	if !ok {
+		return 0, ErrnoFault
+	}
+
+	for iovsPos := uint32(0); iovsPos < iovsStop; iovsPos += 8 {
+		offset := le.Uint32(iovsBuf[iovsPos:])
+		l := le.Uint32(iovsBuf[iovsPos+4:])
+
 		b, ok := mem.Read(ctx, offset, l)
 		if !ok {
 			return 0, ErrnoFault
@@ -1102,31 +1103,35 @@ var fdWrite = proxyResultParams(&wasm.HostFunc{
 }, fdWriteName)
 
 func fdWriteFn(ctx context.Context, mod api.Module, stack []uint64) (uint32, Errno) {
+	sysCtx := mod.(*wasm.CallContext).Sys
+	mem := mod.Memory()
+
 	fd := uint32(stack[0])
 	iovs := uint32(stack[1])
 	iovsCount := uint32(stack[2])
 
-	sysCtx := mod.(*wasm.CallContext).Sys
 	writer := internalsys.FdWriter(ctx, sysCtx, fd)
 	if writer == nil {
 		return 0, ErrnoBadf
 	}
 
 	var err error
 	var nwritten uint32
-	iovsBuf, ok := mod.Memory().Read(ctx, iovs, iovsCount<<3)
+	iovsStop := iovsCount << 3 // iovsCount * 8
+	iovsBuf, ok := mem.Read(ctx, iovs, iovsStop)
 	if !ok {
 		return 0, ErrnoFault
 	}
-	for iovsPos := uint32(0); iovsPos < (iovsCount << 3); iovsPos += 8 {
+
+	for iovsPos := uint32(0); iovsPos < iovsStop; iovsPos += 8 {
 		offset := le.Uint32(iovsBuf[iovsPos:])
 		l := le.Uint32(iovsBuf[iovsPos+4:])
 
 		var n int
 		if writer == io.Discard { // special-case default
 			n = int(l)
 		} else {
-			b, ok := mod.Memory().Read(ctx, offset, l)
+			b, ok := mem.Read(ctx, offset, l)
 			if !ok {
 				return 0, ErrnoFault
 			}

imports/wasi_snapshot_preview1/wasi_bench_test.go

@@ -49,8 +49,79 @@ func Benchmark_ArgsEnviron(b *testing.B) {
 	}
 }
 
+type money struct{}
+
+// Read implements io.Reader by returning endless '$'.
+func (money) Read(b []byte) (n int, err error) {
+	for i := range b {
+		b[i] = '$'
+	}
+
+	return len(b), nil
+}
+
+func Benchmark_fdRead(b *testing.B) {
+	r := wazero.NewRuntime(testCtx)
+	defer r.Close(testCtx)
+
+	mod, err := instantiateProxyModule(r, wazero.NewModuleConfig().WithStdin(money{}))
+	if err != nil {
+		b.Fatal(err)
+	}
+	fn := mod.ExportedFunction(fdReadName)
+
+	mod.Memory().Write(testCtx, 0, []byte{
+		32, 0, 0, 0, // = iovs[0].offset
+		8, 0, 0, 0, // = iovs[0].length
+		40, 0, 0, 0, // = iovs[1].offset
+		8, 0, 0, 0, // = iovs[1].length
+		48, 0, 0, 0, // = iovs[2].offset
+		16, 0, 0, 0, // = iovs[2].length
+		64, 0, 0, 0, // = iovs[3].offset
+		16, 0, 0, 0, // = iovs[3].length
+	})
+
+	benches := []struct {
+		name      string
+		iovs      uint32
+		iovsCount uint32
+	}{
+		{
+			name:      "1x8",
+			iovs:      0,
+			iovsCount: 1,
+		},
+		{
+			name:      "2x16",
+			iovs:      16,
+			iovsCount: 2,
+		},
+	}
+
+	for _, bb := range benches {
+		bc := bb
+
+		b.ReportAllocs()
+		b.Run(bc.name, func(b *testing.B) {
+			resultNread := uint32(128) // arbitrary offset
+
+			for i := 0; i < b.N; i++ {
+				results, err := fn.Call(testCtx, uint64(0), uint64(bc.iovs), uint64(bc.iovsCount), uint64(resultNread))
+				if err != nil {
+					b.Fatal(err)
+				}
+				errno := Errno(results[0])
+				if errno != 0 {
+					b.Fatal(ErrnoName(errno))
+				}
+			}
+		})
+	}
+}
+
 type writerFunc func(p []byte) (n int, err error)
 
+// Write implements io.Writer by calling writerFunc.
 func (f writerFunc) Write(p []byte) (n int, err error) {
 	return f(p)
 }