github.com/ncruces/go-sqlite3@v0.15.1-0.20240520133447-53eef1510ff0/vfs/tests/speedtest1/speedtest1_test.go (about)

     1  package speedtest1
     2  
     3  import (
     4  	"bytes"
     5  	"compress/bzip2"
     6  	"context"
     7  	"crypto/rand"
     8  	"flag"
     9  	"io"
    10  	"os"
    11  	"path/filepath"
    12  	"runtime"
    13  	"strconv"
    14  	"strings"
    15  	"testing"
    16  
    17  	_ "embed"
    18  
    19  	"github.com/tetratelabs/wazero"
    20  	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
    21  
    22  	"github.com/ncruces/go-sqlite3/internal/util"
    23  	"github.com/ncruces/go-sqlite3/vfs"
    24  	_ "github.com/ncruces/go-sqlite3/vfs/adiantum"
    25  	_ "github.com/ncruces/go-sqlite3/vfs/memdb"
    26  )
    27  
    28  //go:embed testdata/speedtest1.wasm.bz2
    29  var compressed string
    30  
    31  var (
    32  	rt      wazero.Runtime
    33  	module  wazero.CompiledModule
    34  	output  bytes.Buffer
    35  	options []string
    36  )
    37  
    38  func TestMain(m *testing.M) {
    39  	initFlags()
    40  
    41  	ctx := context.Background()
    42  	rt = wazero.NewRuntime(ctx)
    43  	wasi_snapshot_preview1.MustInstantiate(ctx, rt)
    44  	env := vfs.ExportHostFunctions(rt.NewHostModuleBuilder("env"))
    45  	_, err := env.Instantiate(ctx)
    46  	if err != nil {
    47  		panic(err)
    48  	}
    49  
    50  	if !strings.HasPrefix(compressed, "BZh") {
    51  		panic("Please use Git LFS to clone this repo: https://git-lfs.com/")
    52  	}
    53  	binary, err := io.ReadAll(bzip2.NewReader(strings.NewReader(compressed)))
    54  	if err != nil {
    55  		panic(err)
    56  	}
    57  
    58  	module, err = rt.CompileModule(ctx, binary)
    59  	if err != nil {
    60  		panic(err)
    61  	}
    62  
    63  	code := m.Run()
    64  	defer os.Exit(code)
    65  	io.Copy(os.Stderr, &output)
    66  }
    67  
    68  func initFlags() {
    69  	i := 1
    70  	options = append(options, "speedtest1")
    71  	for _, arg := range os.Args[1:] {
    72  		switch {
    73  		case strings.HasPrefix(arg, "-test."):
    74  			// keep test flags
    75  			os.Args[i] = arg
    76  			i++
    77  		default:
    78  			// collect everything else
    79  			options = append(options, arg)
    80  		}
    81  	}
    82  	os.Args = os.Args[:i]
    83  	flag.Parse()
    84  }
    85  
    86  func Benchmark_speedtest1(b *testing.B) {
    87  	output.Reset()
    88  	ctx := util.NewContext(context.Background())
    89  	name := filepath.Join(b.TempDir(), "test.db")
    90  	args := append(options, "--size", strconv.Itoa(b.N), name)
    91  	cfg := wazero.NewModuleConfig().
    92  		WithArgs(args...).WithName("speedtest1").
    93  		WithStdout(&output).WithStderr(&output).
    94  		WithSysWalltime().WithSysNanotime().WithSysNanosleep().
    95  		WithOsyield(runtime.Gosched).
    96  		WithRandSource(rand.Reader)
    97  	mod, err := rt.InstantiateModule(ctx, module, cfg)
    98  	if err != nil {
    99  		b.Fatal(err)
   100  	}
   101  	mod.Close(ctx)
   102  }
   103  
   104  func Benchmark_adiantum(b *testing.B) {
   105  	output.Reset()
   106  	ctx := util.NewContext(context.Background())
   107  	name := "file:" + filepath.Join(b.TempDir(), "test.db") +
   108  		"?hexkey=e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
   109  	args := append(options, "--vfs", "adiantum", "--size", strconv.Itoa(b.N), name)
   110  	cfg := wazero.NewModuleConfig().
   111  		WithArgs(args...).WithName("speedtest1").
   112  		WithStdout(&output).WithStderr(&output).
   113  		WithSysWalltime().WithSysNanotime().WithSysNanosleep().
   114  		WithOsyield(runtime.Gosched).
   115  		WithRandSource(rand.Reader)
   116  	mod, err := rt.InstantiateModule(ctx, module, cfg)
   117  	if err != nil {
   118  		b.Fatal(err)
   119  	}
   120  	mod.Close(ctx)
   121  }