Fix race condition when we init lambda

There's a couple of problems here:

1. We subscribe a telemetry endpoint before the service is ready.  If we
   receive a flush to that endpoint, then we try to perform a flush and
   get an NPE.

2. The access to the flusher is unprotected across multiple goroutines.  We
   just use an RWMutex for this, because it's rarely accessed and doesn't
   need anything fancy.

3. There's a race condition in the test because the nop write is (was) async.

Author: squizzling

internal/awslambda/extension/manager.go

@@ -104,21 +104,9 @@ func (m *manager) Run(parent context.Context) error {
 		return err
 	}
 
-	// Start telemetry server
 	ctx, cancel := context.WithCancel(parent)
 	defer cancel()
 
-	if m.telemetryEnabled() {
-		wg.StartWithContext(ctx, func(c context.Context) {
-			chErrs <- m.telemetryServer.Start(c)
-		})
-
-		if err := m.subscribeToTelemetry(ctx); err != nil {
-			m.log.WithError(err).Error("error subscribing to lambda telemetry endpoint")
-			return err
-		}
-	}
-
 	// Start gostatsd server
 	wg.StartWithContext(ctx, func(c context.Context) {
 		err := m.server.Run(c)
@@ -130,6 +118,18 @@ func (m *manager) Run(parent context.Context) error {
 		}
 	})
 
+	// Start telemetry server
+	if m.telemetryEnabled() {
+		wg.StartWithContext(ctx, func(c context.Context) {
+			chErrs <- m.telemetryServer.Start(c)
+		})
+
+		if err := m.subscribeToTelemetry(ctx); err != nil {
+			m.log.WithError(err).Error("error subscribing to lambda telemetry endpoint")
+			return err
+		}
+	}
+
 	select {
 	case err := <-chErrs:
 		// In the event that the lambda finished early before

internal/flush/flush_coordinator.go

@@ -1,5 +1,9 @@
 package flush
 
+import (
+	"sync"
+)
+
 type Coordinator interface {
 	Flushable
 	NotifyFlush()
@@ -14,13 +18,20 @@ type Flushable interface {
 var _ Coordinator = (*coordinator)(nil)
 
 type coordinator struct {
-	t         Flushable
 	flushChan chan struct{}
+
+	mu sync.RWMutex
+	t  Flushable
 }
 
 func NewFlushCoordinator() Coordinator {
+	// TODO: Because of the overall architecture, we don't have the Flusher
+	//       at the time we create the Coordinator.  It would be good to
+	//       re-architect the world for proper dependency injection, however
+	//       that is a non-trivial chunk of work.  So for now, it's just TODO.
 	return &coordinator{
 		flushChan: make(chan struct{}, 1),
+		t:         noopFlusher{},
 	}
 }
 
@@ -29,14 +40,24 @@ func (fm *coordinator) NotifyFlush() {
 }
 
 func (fm *coordinator) Flush() {
-	fm.t.Flush()
+	if t := fm.getTarget(); t != nil {
+		t.Flush()
+	}
 }
 
 func (fm *coordinator) WaitForFlush() {
 	<-fm.flushChan
 }
 
+func (fm *coordinator) getTarget() Flushable {
+	fm.mu.RLock()
+	defer fm.mu.RUnlock()
+	return fm.t
+}
+
 func (fm *coordinator) RegisterFlushable(t Flushable) {
+	fm.mu.Lock()
+	defer fm.mu.Unlock()
 	fm.t = t
 }
 
@@ -53,3 +74,7 @@ func (n *noopFlushCoordinator) NotifyFlush() {}
 func (n *noopFlushCoordinator) RegisterFlushable(Flushable) {}
 
 func (n *noopFlushCoordinator) WaitForFlush() {}
+
+type noopFlusher struct{}
+
+func (noopFlusher) Flush() {}

pkg/statsd/handler_http_forwarder_v2.go

@@ -308,7 +308,7 @@ func (hfh *HttpForwarderHandlerV2) sendNop(ctx context.Context) {
 
 func (hfh *HttpForwarderHandlerV2) Run(ctx context.Context) {
 	var wg wait.Group
-	wg.StartWithContext(ctx, hfh.sendNop)
+	hfh.sendNop(ctx)
 	wg.Start(func() {
 		for metricMaps := range hfh.consolidatedMetrics {
 			hfh.acquireMergingSem()

pkg/statsd/handler_http_forwarder_v2_test.go

@@ -494,7 +494,9 @@ func TestManualFlush(t *testing.T) {
 
 	wg.Wait()
 
-	assert.Equal(t, uint64(2), atomic.LoadUint64(&ts.called), "Handler must have been called")
+	// First is for the "nop" (see HttpForwarderHandlerV2.Run + HttpForwarderHandlerV2.sendNop)
+	// Second is for the actual flush.
+	assert.Equal(t, uint64(1+1), atomic.LoadUint64(&ts.called), "Handler must have been called")
 	assert.EqualValues(t, 1, atomic.LoadUint64(&ts.pineappleCount))
 	assert.EqualValues(t, 1, atomic.LoadUint64(&ts.derpCount))
 	assert.EqualValues(t, 10, atomic.LoadInt64(&ts.derpValue))