From 65ffff4282d9059b44b9abf8dcf95b8b9e670ca4 Mon Sep 17 00:00:00 2001
From: Max Englander <max@planetscale.com>
Date: Wed, 28 Sep 2022 17:21:14 -0400
Subject: [PATCH] expose /metrics from vtbackup http --port

Signed-off-by: Max Englander <max@planetscale.com>
---
 go/cmd/vtbackup/plugin_prometheusbackend.go | 31 +++++++++++++
 go/cmd/vtbackup/vtbackup.go                 | 51 +++++++++++++--------
 2 files changed, 64 insertions(+), 18 deletions(-)
 create mode 100644 go/cmd/vtbackup/plugin_prometheusbackend.go

diff --git a/go/cmd/vtbackup/plugin_prometheusbackend.go b/go/cmd/vtbackup/plugin_prometheusbackend.go
new file mode 100644
index 00000000000..de4ecbb5e9f
--- /dev/null
+++ b/go/cmd/vtbackup/plugin_prometheusbackend.go
@@ -0,0 +1,31 @@
+/*
+Copyright 2022 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+// This plugin imports Prometheus to allow for instrumentation
+// with the Prometheus client library
+
+import (
+	"vitess.io/vitess/go/stats/prometheusbackend"
+	"vitess.io/vitess/go/vt/servenv"
+)
+
+func init() {
+	servenv.OnRun(func() {
+		prometheusbackend.Init("vtbackup")
+	})
+}
diff --git a/go/cmd/vtbackup/vtbackup.go b/go/cmd/vtbackup/vtbackup.go
index 59344f6c1b9..f4a02f661d6 100644
--- a/go/cmd/vtbackup/vtbackup.go
+++ b/go/cmd/vtbackup/vtbackup.go
@@ -63,8 +63,8 @@ import (
 	"fmt"
 	"math"
 	"math/big"
+	mathrand "math/rand"
 	"os"
-	"os/signal"
 	"strings"
 	"syscall"
 	"time"
@@ -112,11 +112,12 @@ var (
 	initShard          string
 	concurrency        = 4
 	// mysqlctld-like flags
-	mysqlPort     = 3306
-	mysqlSocket   string
-	mysqlTimeout  = 5 * time.Minute
-	initDBSQLFile string
-	detachedMode  bool
+	mysqlPort        = 3306
+	mysqlSocket      string
+	mysqlTimeout     = 5 * time.Minute
+	initDBSQLFile    string
+	detachedMode     bool
+	keepAliveTimeout = 0 * time.Second
 )
 
 func registerFlags(fs *pflag.FlagSet) {
@@ -137,17 +138,32 @@ func registerFlags(fs *pflag.FlagSet) {
 	fs.DurationVar(&mysqlTimeout, "mysql_timeout", mysqlTimeout, "how long to wait for mysqld startup")
 	fs.StringVar(&initDBSQLFile, "init_db_sql_file", initDBSQLFile, "path to .sql file to run after mysql_install_db")
 	fs.BoolVar(&detachedMode, "detach", detachedMode, "detached mode - run backups detached from the terminal")
+	fs.DurationVar(&keepAliveTimeout, "keep-alive-timeout", keepAliveTimeout, "Wait until timeout elapses after a successful backup before shutting down.")
 }
 
 func init() {
+	mathrand.Seed(time.Now().UnixNano())
+	servenv.RegisterDefaultFlags()
+	servenv.RegisterFlags()
+	dbconfigs.RegisterFlags(dbconfigs.All...)
+	mysqlctl.RegisterFlags()
 	servenv.OnParse(registerFlags)
 }
 
 func main() {
 	defer exit.Recover()
-	dbconfigs.RegisterFlags(dbconfigs.All...)
-	mysqlctl.RegisterFlags()
+	servenv.Init()
 	servenv.ParseFlags("vtbackup")
+	ctx, cancel := context.WithCancel(context.Background())
+	servenv.OnClose(func() {
+		cancel()
+	})
+	defer func() {
+		servenv.ExitChan <- syscall.SIGTERM
+		<-ctx.Done()
+	}()
+	go servenv.RunDefault()
+
 	if detachedMode {
 		// this method will call os.Exit and kill this process
 		cmd.DetachFromTerminalAndExit()
@@ -160,16 +176,6 @@ func main() {
 		exit.Return(1)
 	}
 
-	// Catch SIGTERM and SIGINT so we get a chance to clean up.
-	ctx, cancel := context.WithCancel(context.Background())
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-	go func() {
-		sig := <-sigChan
-		log.Infof("Cancelling due to signal: %v", sig)
-		cancel()
-	}()
-
 	// Open connection backup storage.
 	backupStorage, err := backupstorage.GetBackupStorage()
 	if err != nil {
@@ -202,6 +208,15 @@ func main() {
 		log.Errorf("Couldn't prune old backups: %v", err)
 		exit.Return(1)
 	}
+
+	if keepAliveTimeout > 0 {
+		log.Infof("Backup was successful, waiting %s before exiting (or until context expires).", keepAliveTimeout)
+		select {
+		case <-time.After(keepAliveTimeout):
+		case <-ctx.Done():
+		}
+	}
+	log.Info("Exiting.")
 }
 
 func takeBackup(ctx context.Context, topoServer *topo.Server, backupStorage backupstorage.BackupStorage) error {