From c83415e79a8a934294776f9162b46147e53781fa Mon Sep 17 00:00:00 2001
From: Jamie Wilkinson <jaq@spacepants.org>
Date: Mon, 16 Oct 2023 09:17:04 +1100
Subject: [PATCH] Add the waitgroup correctly to exporter and cleanly shutdown.

Disable the exporter loop in compile-only and one shot modes.
---
 internal/exporter/export.go | 42 ++++++++++++++++++++++++-------------
 internal/mtail/mtail.go     |  5 -----
 internal/mtail/options.go   |  4 +++-
 3 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/internal/exporter/export.go b/internal/exporter/export.go
index 17f019df6..35d4e8f3e 100644
--- a/internal/exporter/export.go
+++ b/internal/exporter/export.go
@@ -30,15 +30,16 @@ var (
 
 // Exporter manages the export of metrics to passive and active collectors.
 type Exporter struct {
-	ctx           context.Context
-	wg            sync.WaitGroup
-	store         *metrics.Store
-	pushInterval  time.Duration
-	hostname      string
-	omitProgLabel bool
-	emitTimestamp bool
-	pushTargets   []pushOptions
-	initDone      chan struct{}
+	ctx            context.Context
+	wg             sync.WaitGroup
+	store          *metrics.Store
+	pushInterval   time.Duration
+	hostname       string
+	omitProgLabel  bool
+	emitTimestamp  bool
+	exportDisabled bool
+	pushTargets    []pushOptions
+	initDone       chan struct{}
 }
 
 // Option configures a new Exporter.
@@ -75,6 +76,13 @@ func PushInterval(opt time.Duration) Option {
 	}
 }
 
+func DisableExport() Option {
+	return func(e *Exporter) error {
+		e.exportDisabled = true
+		return nil
+	}
+}
+
 // New creates a new Exporter.
 func New(ctx context.Context, wg *sync.WaitGroup, store *metrics.Store, options ...Option) (*Exporter, error) {
 	if store == nil {
@@ -112,13 +120,15 @@ func New(ctx context.Context, wg *sync.WaitGroup, store *metrics.Store, options
 	}
 	e.StartMetricPush()
 
-	// This routine manages shutdown of the Exporter.  TODO(jaq): This doesn't
-	// happen before mtail returns because of how context cancellation is set
-	// up..  How can we tie this shutdown in before mtail exits?  Should
-	// exporter be merged with httpserver?
+	wg.Add(1)
+	// This routine manages shutdown of the Exporter.
 	go func() {
+		defer wg.Done()
 		<-e.initDone
-		<-e.ctx.Done()
+		// Wait for the context to be completed before waiting for subroutines.
+		if !e.exportDisabled {
+			<-e.ctx.Done()
+		}
 		e.wg.Wait()
 	}()
 	return e, nil
@@ -213,6 +223,10 @@ func (e *Exporter) PushMetrics() {
 
 // StartMetricPush pushes metrics to the configured services each interval.
 func (e *Exporter) StartMetricPush() {
+	if e.exportDisabled {
+		glog.Info("Export loop disabled.")
+		return
+	}
 	if len(e.pushTargets) == 0 {
 		return
 	}
diff --git a/internal/mtail/mtail.go b/internal/mtail/mtail.go
index 33b0b07bd..8f1ac0c2f 100644
--- a/internal/mtail/mtail.go
+++ b/internal/mtail/mtail.go
@@ -62,11 +62,6 @@ func (m *Server) initRuntime() (err error) {
 
 // initExporter sets up an Exporter for this Server.
 func (m *Server) initExporter() (err error) {
-	if m.oneShot {
-		// This is a hack to avoid a race in test, but assume that in oneshot
-		// mode we don't want to export anything.
-		return nil
-	}
 	m.e, err = exporter.New(m.ctx, &m.wg, m.store, m.eOpts...)
 	if err != nil {
 		return err
diff --git a/internal/mtail/options.go b/internal/mtail/options.go
index 1941d8f5a..71dc062b0 100644
--- a/internal/mtail/options.go
+++ b/internal/mtail/options.go
@@ -164,6 +164,7 @@ var OneShot = &niladicOption{
 	func(m *Server) error {
 		m.rOpts = append(m.rOpts, runtime.ErrorsAbort())
 		m.tOpts = append(m.tOpts, tailer.OneShot)
+		m.eOpts = append(m.eOpts, exporter.DisableExport())
 		m.oneShot = true
 		return nil
 	},
@@ -173,6 +174,7 @@ var OneShot = &niladicOption{
 var CompileOnly = &niladicOption{
 	func(m *Server) error {
 		m.rOpts = append(m.rOpts, runtime.CompileOnly())
+		m.eOpts = append(m.eOpts, exporter.DisableExport())
 		m.compileOnly = true
 		return nil
 	},
@@ -186,7 +188,7 @@ var DumpAst = &niladicOption{
 	},
 }
 
-// DumpAstTypes instructs the Server's copmiler to print the AST after type checking.
+// DumpAstTypes instructs the Server's compiler to print the AST after type checking.
 var DumpAstTypes = &niladicOption{
 	func(m *Server) error {
 		m.rOpts = append(m.rOpts, runtime.DumpAstTypes())