Repository: atlassian/gostatsd
Pull request #422 (2021-11-20 10:14:43)

Ensure forwarder flushes all queued data
Does what it says on the tin :) 

There is a breaking change that it now enforces that a forwarder is single use since there is no means of unclosing a channel.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ddcac70..8296b4a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+unreleased
+------
+- Forces http metrics forwarder to flush all data queued on shutdown
+
 33.0.2
 ------
 - Fixing release process after removing `-sym` releases in prior tag
diff --git a/metric_consolidator.go b/metric_consolidator.go
index 6a9ad2f..39923a2 100644
--- a/metric_consolidator.go
+++ b/metric_consolidator.go
@@ -8,8 +8,9 @@ import (
 )
 
 // MetricConsolidator will consolidate metrics randomly in to a slice of MetricMaps, and either send the slice to
-// the provided channel, or make them available synchronously through Drain/Fill. Run can also be started in a long
-// running goroutine to perform flushing, or Flush can be called externally to trigger the channel send.
+// the provided channel, or make them available synchronously through Drain[WithContext]/Fill. Run can also be
+// started in a long running goroutine to perform flushing, or Flush can be called externally to trigger the channel
+// send.
 //
 // Used to consolidate metrics such as:
 // - counter[name=x, value=1]
@@ -37,24 +38,31 @@ func NewMetricConsolidator(spots int, flushInterval time.Duration, sink chan<- [
 	return mc
 }
 
-func (mc *MetricConsolidator) Run(ctx context.Context) {
-	t := clock.NewTicker(ctx, mc.flushInterval)
-	defer t.Stop()
+	func (mc *MetricConsolidator) Run(ctx context.Context) {
+		clck := clock.FromContext(ctx)
+		t := clck.NewTicker(mc.flushInterval)
+		defer t.Stop()
 
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case <-t.C:
-			mc.Flush(ctx)
+		for {
+			select {
+			case <-ctx.Done():
+				mc.Flush()
+				return
+			case <-t.C:
+				mc.Flush()
+			}
 		}
 	}
+
+// Drain will collect all the MetricMaps in the MetricConsolidator and return them.
+func (mc *MetricConsolidator) Drain() []*MetricMap {
+	return mc.DrainWithContext(context.Background())
 }
 
-// Drain will collect all the MetricMaps in the MetricConsolidator and return them.  If the
-// context.Context is canceled before everything can be collected, they are returned to the
-// MetricConsolidator and nil is returned.
-func (mc *MetricConsolidator) Drain(ctx context.Context) []*MetricMap {
+// DrainWithContext will collect all the MetricMaps in the MetricConsolidator and return them.
+// If the context.Context is canceled before everything can be collected, they are returned to
+// the MetricConsolidator and nil is returned.
+func (mc *MetricConsolidator) DrainWithContext(ctx context.Context) []*MetricMap {
 	mms := make([]*MetricMap, 0, cap(mc.maps))
 	for i := 0; i < cap(mc.maps); i++ {
 		select {
@@ -74,24 +82,15 @@ func (mc *MetricConsolidator) Drain(ctx context.Context) []*MetricMap {
 
 // Flush will collect all the MetricMaps in to a slice, send them to the channel provided, then
 // create new MetricMaps for new metrics to land in.  Not thread-safe.
-func (mc *MetricConsolidator) Flush(ctx context.Context) {
-	mms := mc.Drain(ctx)
-	if mms == nil {
-		return
-	}
-
+func (mc *MetricConsolidator) Flush() {
 	// Send the collected data to the sink before putting new maps in place.  This allows back-pressure
 	// to propagate through the system, if the sink can't keep up.
-	select {
-	case mc.sink <- mms:
-	case <-ctx.Done():
-	}
-
+	mc.sink <- mc.Drain()
 	mc.Fill()
 }
 
-// Fill re-populates the MetricConsolidator with empty MetricMaps, it is the pair to Drain and
-// must be called after a successful Drain, must not be called after a failed Drain.
+// Fill re-populates the MetricConsolidator with empty MetricMaps, it is the pair to Drain[WithContext] and
+// must be called after a successful Drain[WithContext], must not be called after a failed DrainWithContext.
 func (mc *MetricConsolidator) Fill() {
 	for i := 0; i < cap(mc.maps); i++ {
 		mc.maps <- NewMetricMap()
diff --git a/pkg/stats/statser_internal.go b/pkg/stats/statser_internal.go
index f062568..07f70c1 100644
--- a/pkg/stats/statser_internal.go
+++ b/pkg/stats/statser_internal.go
@@ -44,7 +44,7 @@ func NewInternalStatser(tags gostatsd.Tags, namespace string, hostname gostatsd.
 }
 
 func (is *InternalStatser) NotifyFlush(ctx context.Context, d time.Duration) {
-	mms := is.consolidator.Drain(ctx)
+	mms := is.consolidator.DrainWithContext(ctx)
 	if mms == nil {
 		// context is canceled
 		return
diff --git a/pkg/statsd/handler_http_forwarder_v2.go b/pkg/statsd/handler_http_forwarder_v2.go
index e5aa57d..89fb243 100644
--- a/pkg/statsd/handler_http_forwarder_v2.go
+++ b/pkg/statsd/handler_http_forwarder_v2.go
@@ -50,12 +50,14 @@ type HttpForwarderHandlerV2 struct {
 	maxRequestElapsedTime time.Duration
 	metricsSem            chan struct{}
 	client                *http.Client
-	consolidator          *gostatsd.MetricConsolidator
-	consolidatedMetrics   <-chan []*gostatsd.MetricMap
 	eventWg               sync.WaitGroup
 	compress              bool
 	headers               map[string]string
 	dynHeaderNames        []string
+
+	done                chan struct{}
+	consolidator        *gostatsd.MetricConsolidator
+	consolidatedMetrics chan []*gostatsd.MetricMap
 }
 
 // NewHttpForwarderHandlerV2FromViper returns a new http API client.
@@ -173,6 +175,7 @@ func NewHttpForwarderHandlerV2(
 		client:                httpClient.Client,
 		headers:               headers,
 		dynHeaderNames:        dynHeaderNamesWithColon,
+		done:                  make(chan struct{}),
 	}, nil
 }
 
@@ -182,7 +185,14 @@ func (hfh *HttpForwarderHandlerV2) EstimatedTags() int {
 
 // DispatchMetricMap dispatches a metric map to the MetricConsolidator
 func (hfh *HttpForwarderHandlerV2) DispatchMetricMap(ctx context.Context, mm *gostatsd.MetricMap) {
-	hfh.consolidator.ReceiveMetricMap(mm)
+	select {
+	case <-ctx.Done():
+		hfh.logger.WithError(ctx.Err()).Debug("Context is cancelled")
+	case <-hfh.done:
+		hfh.logger.Debug("Forwarder has shutdown and unable to accept metrics")
+	default:
+		hfh.consolidator.ReceiveMetricMap(mm)
+	}
 }
 
 func (hfh *HttpForwarderHandlerV2) RunMetricsContext(ctx context.Context) {
@@ -217,28 +227,39 @@ func (hfh *HttpForwarderHandlerV2) emitMetrics(statser stats.Statser) {
 
 func (hfh *HttpForwarderHandlerV2) Run(ctx context.Context) {
 	var wg wait.Group
-	defer wg.Wait()
-	wg.StartWithContext(ctx, hfh.consolidator.Run)
-
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case metricMaps := <-hfh.consolidatedMetrics:
+	wg.Start(func() {
+		for metricMaps := range hfh.consolidatedMetrics {
 			mergedMetricMap := mergeMaps(metricMaps)
 			mms := mergedMetricMap.SplitByTags(hfh.dynHeaderNames)
 			for dynHeaderTags, mm := range mms {
-				if !hfh.acquireSem(ctx) {
-					return
+				if mm.IsEmpty() {
+					continue
 				}
+				hfh.acquireSem()
 				postId := atomic.AddUint64(&hfh.postId, 1) - 1
 				go func(postId uint64, metricMap *gostatsd.MetricMap, dynHeaderTags string) {
-					hfh.postMetrics(ctx, metricMap, dynHeaderTags, postId)
+					hfh.postMetrics(context.Background(), metricMap, dynHeaderTags, postId)
 					hfh.releaseSem()
 				}(postId, mm, dynHeaderTags)
 			}
 		}
-	}
+		for i := 0; i < cap(hfh.metricsSem); i++ {
+			hfh.acquireSem()
+		}
+	})
+	wg.StartWithContext(ctx, func(ctx context.Context) {
+		hfh.consolidator.Run(ctx)
+		hfh.Close()
+	})
+
+	wg.Wait()
+	hfh.logger.Debug("Shutdown http forwarder")
+
+}
+
+func (hfh *HttpForwarderHandlerV2) Close() {
+	close(hfh.done)
+	close(hfh.consolidatedMetrics)
 }
 
 func mergeMaps(maps []*gostatsd.MetricMap) *gostatsd.MetricMap {
@@ -249,13 +270,8 @@ func mergeMaps(maps []*gostatsd.MetricMap) *gostatsd.MetricMap {
 	return mm
 }
 
-func (hfh *HttpForwarderHandlerV2) acquireSem(ctx context.Context) bool {
-	select {
-	case <-ctx.Done():
-		return false
-	case <-hfh.metricsSem:
-		return true
-	}
+func (hfh *HttpForwarderHandlerV2) acquireSem() {
+	<-hfh.metricsSem // will potentially block
 }
 
 func (hfh *HttpForwarderHandlerV2) releaseSem() {


Test patch:
diff --git a/metric_consolidator_test.go b/metric_consolidator_test.go
index 7b6035a..4a32568 100644
--- a/metric_consolidator_test.go
+++ b/metric_consolidator_test.go
@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
-	"github.com/tilinna/clock"
 )
 
 func TestConsolidation(t *testing.T) {
@@ -16,9 +15,6 @@ func TestConsolidation(t *testing.T) {
 	ctxTest, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	mockClock := clock.NewMock(time.Unix(0, 0))
-	ctxClock := clock.Context(ctxTest, mockClock)
-
 	ch := make(chan []*MetricMap, 1)
 	mc := NewMetricConsolidator(2, 1*time.Second, ch)
 
@@ -38,7 +34,7 @@ func TestConsolidation(t *testing.T) {
 	}
 	mc.ReceiveMetrics([]*Metric{m1})
 	mc.ReceiveMetrics([]*Metric{m2})
-	mc.Flush(ctxClock)
+	mc.Flush()
 
 	var mm []*MetricMap
 	select {
diff --git a/pkg/statsd/handler_http_forwarder_v2_test.go b/pkg/statsd/handler_http_forwarder_v2_test.go
index e38cc28..9176645 100644
--- a/pkg/statsd/handler_http_forwarder_v2_test.go
+++ b/pkg/statsd/handler_http_forwarder_v2_test.go
@@ -1,14 +1,24 @@
 package statsd
 
 import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/ash2k/stager/wait"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/tilinna/clock"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/atlassian/gostatsd"
+	"github.com/atlassian/gostatsd/internal/fixtures"
 	"github.com/atlassian/gostatsd/pb"
 	"github.com/atlassian/gostatsd/pkg/transport"
 )
@@ -207,6 +217,83 @@ func BenchmarkHttpForwarderV2TranslateAll(b *testing.B) {
 	}
 }
 
+func TestForwardingData(t *testing.T) {
+	t.Parallel()
+	var called, havePineapples, haveDerpinton uint64
+	s := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		atomic.AddUint64(&called, 1)
+
+		buf, err := io.ReadAll(r.Body)
+		require.NoError(t, err, "Must not error when reading request")
+		defer r.Body.Close()
+
+		var data pb.RawMessageV2
+		require.NoError(t, proto.Unmarshal(buf, &data))
+
+		counters, ok := data.GetCounters()["pineapples"]
+		if ok {
+			atomic.AddUint64(&havePineapples, 1)
+			val, ok := counters.GetTagMap()["derpinton"]
+			if ok {
+				atomic.AddUint64(&haveDerpinton, 1)
+				assert.Equal(t, int64(10), val.GetValue())
+			}
+		}
+
+	}))
+	t.Cleanup(s.Close)
+
+	log := logrus.New().WithField("testcase", t.Name())
+	pool := transport.NewTransportPool(log, viper.New())
+	forwarder, err := NewHttpForwarderHandlerV2(
+		log,
+		"default",
+		s.URL,
+		1,
+		1,
+		false,
+		100*time.Millisecond, // maxRequestElapsedTime
+		100*time.Millisecond, // flushInterval
+		map[string]string{},
+		[]string{},
+		pool,
+	)
+	require.NoError(t, err, "Must not error when creating the forwarder")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	mockClock := clock.NewMock(time.Unix(1000, 0))
+	ctx = clock.Context(ctx, mockClock)
+	t.Cleanup(cancel)
+
+	for i := 0; i < 10; i++ {
+		forwarder.DispatchMetricMap(ctx, &gostatsd.MetricMap{
+			Counters: gostatsd.Counters{
+				"pineapples": map[string]gostatsd.Counter{
+					"derpinton": {
+						PerSecond: 0.1,
+						Value:     1,
+						Source:    gostatsd.UnknownSource,
+						Timestamp: gostatsd.Nanotime(mockClock.Now().Nanosecond()),
+						Tags:      gostatsd.Tags{},
+					},
+				},
+			},
+		})
+	}
+	var wg wait.Group
+	wg.StartWithContext(ctx, forwarder.Run)
+	fixtures.NextStep(ctx, mockClock)
+
+	cancel()
+
+	wg.Wait()
+
+	assert.Greater(t, atomic.LoadUint64(&called), uint64(0), "Handler must have been called")
+	assert.EqualValues(t, 1, atomic.LoadUint64(&havePineapples))
+	assert.EqualValues(t, 1, atomic.LoadUint64(&haveDerpinton))
+	assert.Equal(t, 0, mockClock.Len(), "Must have closed all event handlers")
+}
+
 func TestHttpForwarderV2New(t *testing.T) {
 	logger := logrus.New()
 	pool := transport.NewTransportPool(logger, viper.New())