Warm the MySQL ready pool before EnsureTemplate returns

mysql/clone.go

@@ -16,6 +16,12 @@ import (
 // then views, triggers, procedures and functions from their SHOW CREATE
 // forms. Tables named _testdb_* hold testdb's own bookkeeping inside a
 // template and are not copied.
+//
+// Deliberately not copied: events (a scheduler job firing inside a test's
+// clone would make the test nondeterministic) and grants (tests connect as
+// the admin user). DEFINER clauses are dropped, so objects belong to
+// whoever runs the clone. A foreign key that references a table in another
+// database keeps pointing at that database.
 
 // bookkeepingPrefix names the tables testdb keeps inside a template.
 const bookkeepingPrefix = "_testdb_"

mysql/mysql.go

@@ -19,7 +19,9 @@
 // or a mysql://user:pass@host:port/db?params URL. The user needs CREATE and
 // DROP DATABASE and the right to kill other sessions (CONNECTION_ADMIN).
 //
-// Templates are not available for MySQL yet; see testdb.TemplateProvider.
+// Templates (EnsureTemplate) build the schema once and hand each test a
+// clone prepared ahead of time, which is where MySQL's speed comes from;
+// see Template and WithReadyDatabases.
 package mysql
 
 import (

mysql/pool.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log"
 	"sync"
+	"time"
 
 	"github.com/bashhack/testdb"
 )
@@ -61,6 +62,29 @@ func startReadyPool(tpl *Template, size, workers int) (*readyPool, error) {
 	return pool, nil
 }
 
+// waitFull blocks until the pool holds its full number of prepared
+// databases, the workers have given up, or ctx ends.
+func (p *readyPool) waitFull(ctx context.Context) error {
+	ticker := time.NewTicker(10 * time.Millisecond)
+	defer ticker.Stop()
+	for {
+		if len(p.ready) == cap(p.ready) {
+			return nil
+		}
+		p.mu.Lock()
+		failed := p.failure != nil
+		p.mu.Unlock()
+		if failed {
+			return nil // already logged; tests will clone on demand
+		}
+		select {
+		case <-ticker.C:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
 // take returns a prepared database, or false without waiting.
 func (p *readyPool) take() (string, bool) {
 	select {

mysql/template.go

@@ -71,6 +71,7 @@ type templateConfig struct {
 	ready        int
 	readyWorkers int
 	readyPrefix  string
+	warmup       bool
 }
 
 // WithFingerprint identifies the template's content, typically a hash of the
@@ -103,13 +104,24 @@ func WithTemplateAdminDSN(dsn string) TemplateOption {
 // default is runtime.GOMAXPROCS(0), which is also go test's default
 // -parallel. 0 disables the pool: every test clones on demand.
 //
-// The pool keeps up as long as tests take clones no faster than the
-// workers (see WithReadyWorkers) can make them. When it runs dry, a test
-// clones on demand rather than wait.
+// EnsureTemplate fills the pool once before returning (see
+// WithReadyWarmup), so the first tests find it full. After that it keeps
+// up as long as tests take clones no faster than the workers (see
+// WithReadyWorkers) can make them; the pool size is also the number of
+// back-to-back quick tests it can absorb. When it runs dry, a test clones
+// on demand rather than wait.
 func WithReadyDatabases(n int) TemplateOption {
 	return func(c *templateConfig) { c.ready = n }
 }
 
+// WithReadyWarmup sets whether EnsureTemplate waits for the pool to fill
+// before returning. On by default: it costs one round of parallel clones,
+// once, and means every test from the first finds a prepared database.
+// Turn it off to start tests immediately and let the pool fill alongside.
+func WithReadyWarmup(on bool) TemplateOption {
+	return func(c *templateConfig) { c.warmup = on }
+}
+
 // WithReadyWorkers sets how many clones are made in parallel to keep the
 // pool full. The server handles concurrent clones well, so the default is
 // the pool size.
@@ -154,7 +166,7 @@ func EnsureTemplate(ctx context.Context, name string, setup SetupFunc, opts ...T
 	if err := validateTemplateName(name); err != nil {
 		return nil, fmt.Errorf("mysql.EnsureTemplate: %w", err)
 	}
-	cfg := templateConfig{ready: runtime.GOMAXPROCS(0), readyPrefix: "test"}
+	cfg := templateConfig{ready: runtime.GOMAXPROCS(0), readyPrefix: "test", warmup: true}
 	for _, opt := range opts {
 		opt(&cfg)
 	}
@@ -180,6 +192,14 @@ func EnsureTemplate(ctx context.Context, name string, setup SetupFunc, opts ...T
 			return tpl, nil
 		}
 		tpl.pool = pool
+		if cfg.warmup {
+			if err := pool.waitFull(ctx); err != nil {
+				// ctx ended mid warm-up: nobody gets the handle, so let go of
+				// the template properly
+				err = errors.Join(err, tpl.Release(context.WithoutCancel(ctx)))
+				return nil, fmt.Errorf("mysql.EnsureTemplate %q: warm up prepared databases: %w", name, err)
+			}
+		}
 	}
 	return tpl, nil
 }

mysql/template_benchmark_test.go

@@ -121,7 +121,7 @@ func BenchmarkTemplateReady(b *testing.B) {
 				schemaDir := writeBenchSchema(b)
 				tpl := ensureBenchTemplate(b, "bench_mysql_tpl_ready", schemaDir,
 					mysql.WithReadyDatabases(8), mysql.WithReadyWorkers(workers), mysql.WithReadyPrefix("bench_ready"))
-				time.Sleep(6 * time.Second) // let the pool fill once
+				// EnsureTemplate warms the pool before returning
 				pooled := 0
 				b.ReportAllocs()
 				for b.Loop() {

mysql/template_test.go

@@ -280,7 +280,10 @@ func TestTemplateReadyPool(t *testing.T) {
 		}
 	})
 
-	waitFor(t, "the pool to fill", func() bool { return tpl.Ready() == 2 })
+	// Warm-up is on by default: the pool is full when EnsureTemplate returns
+	if n := tpl.Ready(); n != 2 {
+		t.Fatalf("Ready() = %d right after EnsureTemplate, want 2 (warm-up)", n)
+	}
 
 	// A test takes a prepared database: the name carries the pool's prefix,
 	// the clone is complete, and the pool refills
@@ -310,6 +313,62 @@ func TestTemplateReadyPool(t *testing.T) {
 	t.Logf("5 rapid takes: %d served on demand, %d from the pool", onDemand, 5-onDemand)
 }
 
+func TestTemplateReadyWarmupOff(t *testing.T) {
+	requireMySQL(t)
+	ctx := context.Background()
+	name := uniqueTemplateName(t, "tpl_nowarm")
+
+	start := time.Now()
+	tpl, err := mysql.EnsureTemplate(ctx, name, mysql.SQLFiles(schemaDir),
+		mysql.WithReadyDatabases(2), mysql.WithReadyWarmup(false), mysql.WithReadyPrefix("nowarm"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { release(t, tpl) })
+	t.Logf("EnsureTemplate without warm-up returned in %s with %d ready", time.Since(start), tpl.Ready())
+
+	// The pool fills alongside, and tests work meanwhile
+	db := mysql.Setup(t, tpl.Option())
+	if !tableExists(t, db, "users") {
+		t.Fatal("clone lacks the schema")
+	}
+	waitFor(t, "the pool to fill in the background", func() bool { return tpl.Ready() == 2 })
+}
+
+func TestTemplateWarmupCancelled(t *testing.T) {
+	requireMySQL(t)
+	admin := adminDB(t)
+	name := uniqueTemplateName(t, "tpl_warmcancel")
+
+	// A warm-up one worker cannot finish in time: 256 clones against a
+	// deadline a few clones long. Whether the deadline lands during the
+	// build or during the warm-up, the outcome must be the same.
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+	tpl, err := mysql.EnsureTemplate(ctx, name, mysql.SQLFiles(schemaDir),
+		mysql.WithReadyDatabases(256), mysql.WithReadyWorkers(1), mysql.WithReadyPrefix("warmcancel"))
+	if err == nil {
+		release(t, tpl)
+		t.Fatal("256 clones in 500 ms: the warm-up should not have finished")
+	}
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("expected a deadline error, got %v", err)
+	}
+	// Nothing was left behind: the template and its prepared databases were
+	// released on the way out
+	if databaseExists(t, admin, name) {
+		t.Errorf("template %s left behind after a cancelled warm-up", name)
+	}
+	var leftovers int
+	if err := admin.QueryRowContext(context.Background(),
+		"SELECT COUNT(*) FROM information_schema.SCHEMATA WHERE SCHEMA_NAME LIKE 'warmcancel\\_%'").Scan(&leftovers); err != nil {
+		t.Fatal(err)
+	}
+	if leftovers != 0 {
+		t.Errorf("%d prepared databases left behind after a cancelled warm-up", leftovers)
+	}
+}
+
 func TestTemplateNameValidation(t *testing.T) {
 	ctx := context.Background()
 	if _, err := mysql.EnsureTemplate(ctx, "", nil); err == nil {