From 6f36b9a90057ab60becef0c1cb885fb67574b2df Mon Sep 17 00:00:00 2001
From: Max Englander <max@planetscale.com>
Date: Thu, 28 Sep 2023 18:20:09 +0100
Subject: [PATCH] go/cmd/vtbackup: wait for plugins to finish initializing
 (#14113)

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

diff --git a/go/cmd/vtbackup/cli/vtbackup.go b/go/cmd/vtbackup/cli/vtbackup.go
index 8e70e3c5c1d..3e48b75df4d 100644
--- a/go/cmd/vtbackup/cli/vtbackup.go
+++ b/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
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() {