go/cmd/vtbackup: wait for plugins to finish initializing (#14113)

Signed-off-by: Max Englander <max@planetscale.com>

go/cmd/vtbackup/cli/vtbackup.go

@@ -230,6 +230,11 @@ func run(_ *cobra.Command, args []string) error {
 	}()
 
 	go servenv.RunDefault()
+	// Some stats plugins use OnRun to initialize. Wait for them to finish
+	// initializing before continuing, so we don't lose any stats.
+	if err := stats.AwaitBackend(ctx); err != nil {
+		return fmt.Errorf("failed to await stats backend: %w", err)
+	}
 
 	if detachedMode {
 		// this method will call os.Exit and kill this process

go/stats/export.go

@@ -29,6 +29,7 @@ package stats
 
 import (
 	"bytes"
+	"context"
 	"expvar"
 	"fmt"
 	"strconv"
@@ -45,6 +46,7 @@ var (
 	emitStats         bool
 	statsEmitPeriod   = 60 * time.Second
 	statsBackend      string
+	statsBackendInit  = make(chan struct{})
 	combineDimensions string
 	dropVariables     string
 )
@@ -209,6 +211,18 @@ var pushBackends = make(map[string]PushBackend)
 var pushBackendsLock sync.Mutex
 var once sync.Once
 
+func AwaitBackend(ctx context.Context) error {
+	if statsBackend == "" {
+		return nil
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-statsBackendInit:
+		return nil
+	}
+}
+
 // RegisterPushBackend allows modules to register PushBackend implementations.
 // Should be called on init().
 func RegisterPushBackend(name string, backend PushBackend) {
@@ -218,6 +232,9 @@ func RegisterPushBackend(name string, backend PushBackend) {
 		log.Fatalf("PushBackend %s already exists; can't register the same name multiple times", name)
 	}
 	pushBackends[name] = backend
+	if name == statsBackend {
+		close(statsBackendInit)
+	}
 	if emitStats {
 		// Start a single goroutine to emit stats periodically
 		once.Do(func() {