From bc2178aa134f2bb08d7f77c3ddcabf5385c04f8a Mon Sep 17 00:00:00 2001
From: Max Englander <max@planetscale.com>
Date: Thu, 28 Sep 2023 11:31:55 +0100
Subject: [PATCH 1/2] go/cmd/vtbackup: wait for plugins to finish initializing

Signed-off-by: Max Englander <max@planetscale.com>
---
 go/cmd/vtbackup/cli/vtbackup.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/go/cmd/vtbackup/cli/vtbackup.go b/go/cmd/vtbackup/cli/vtbackup.go
index 8e70e3c5c1d..17ea0468821 100644
--- a/go/cmd/vtbackup/cli/vtbackup.go
+++ b/go/cmd/vtbackup/cli/vtbackup.go
@@ -229,7 +229,14 @@ func run(_ *cobra.Command, args []string) error {
 		<-ctx.Done()
 	}()
 
+	// Some plugins use OnRun to initialize. Use a channel to signal to
+	// ourselves that these plugins have been initialized.
+	runCh := make(chan struct{})
+	servenv.OnRun(func() {
+		close(runCh)
+	})
 	go servenv.RunDefault()
+	<-runCh
 
 	if detachedMode {
 		// this method will call os.Exit and kill this process

From 83932bc282beec1d91bf11d74ceea8e1549e8db1 Mon Sep 17 00:00:00 2001
From: Max Englander <max@planetscale.com>
Date: Thu, 28 Sep 2023 13:53:49 +0100
Subject: [PATCH 2/2] address pr feedback

Signed-off-by: Max Englander <max@planetscale.com>
---
 go/cmd/vtbackup/cli/vtbackup.go | 12 +++++-------
 go/stats/export.go              | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/go/cmd/vtbackup/cli/vtbackup.go b/go/cmd/vtbackup/cli/vtbackup.go
index 17ea0468821..3e48b75df4d 100644
--- a/go/cmd/vtbackup/cli/vtbackup.go
+++ b/go/cmd/vtbackup/cli/vtbackup.go
@@ -229,14 +229,12 @@ func run(_ *cobra.Command, args []string) error {
 		<-ctx.Done()
 	}()
 
-	// Some plugins use OnRun to initialize. Use a channel to signal to
-	// ourselves that these plugins have been initialized.
-	runCh := make(chan struct{})
-	servenv.OnRun(func() {
-		close(runCh)
-	})
 	go servenv.RunDefault()
-	<-runCh
+	// 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
diff --git a/go/stats/export.go b/go/stats/export.go
index 8bda85c87b2..58be67e13f9 100644
--- a/go/stats/export.go
+++ b/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() {