From 7ee944037899ab1c59ff317c5b476276bf4aa1cf Mon Sep 17 00:00:00 2001
From: Evan Anderson <evana@vmware.com>
Date: Wed, 5 Feb 2020 14:43:51 -0800
Subject: [PATCH] Add methods to extract stats.Option

---
 stats/record.go      | 28 ++++++++++++++++++++
 stats/record_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/stats/record.go b/stats/record.go
index ad4691184..c1192b516 100644
--- a/stats/record.go
+++ b/stats/record.go
@@ -31,6 +31,15 @@ func init() {
 	}
 }
 
+// ResolvedOptions can be used to extract the current tags and measurements
+// from context and stats arguments when using custom workers to export stats
+// to a separate exporter.
+type ResolvedOptions struct {
+	Attachments metricdata.Attachments
+	Tags        *tag.Map
+	Measures    []Measurement
+}
+
 type recordOptions struct {
 	attachments  metricdata.Attachments
 	mutators     []tag.Mutator
@@ -84,6 +93,23 @@ func RecordWithTags(ctx context.Context, mutators []tag.Mutator, ms ...Measureme
 	return RecordWithOptions(ctx, WithTags(mutators...), WithMeasurements(ms...))
 }
 
+// ResolveOptions determines the full set of Tags, Measurements, etc from the
+// provided Options and context.Context. This is mostly useful when using
+// multiple exporters.
+func ResolveOptions(ctx context.Context, ros ...Options) (*ResolvedOptions, error) {
+	o := createRecordOption(ros...)
+
+	if len(o.mutators) > 0 {
+		var err error
+		if ctx, err = tag.New(ctx, o.mutators...); err != nil {
+			return nil, err
+		}
+	}
+	return &ResolvedOptions{Tags: tag.FromContext(ctx),
+		Measures:    o.measurements,
+		Attachments: o.attachments}, nil
+}
+
 // RecordWithOptions records measurements from the given options (if any) against context
 // and tags and attachments in the options (if any).
 // If there are any tags in the context, measurements will be tagged with them.
@@ -92,6 +118,8 @@ func RecordWithOptions(ctx context.Context, ros ...Options) error {
 	if len(o.measurements) == 0 {
 		return nil
 	}
+	// This could use ResolveOptions, but it does additional work to
+	// short-circuit if there are no metrics that need to be exported.
 	recorder := internal.DefaultRecorder
 	if recorder == nil {
 		return nil
diff --git a/stats/record_test.go b/stats/record_test.go
index 93a652200..eeb7a885d 100644
--- a/stats/record_test.go
+++ b/stats/record_test.go
@@ -56,6 +56,7 @@ func TestRecordWithAttachments(t *testing.T) {
 	if err := view.Register(v); err != nil {
 		log.Fatalf("Failed to register views: %v", err)
 	}
+	defer view.Unregister(v)
 
 	attachments := map[string]interface{}{metricdata.AttachmentKeySpanContext: spanCtx}
 	stats.RecordWithOptions(context.Background(), stats.WithAttachments(attachments), stats.WithMeasurements(m.M(12)))
@@ -93,3 +94,65 @@ func TestRecordWithAttachments(t *testing.T) {
 func cmpExemplar(got, want *metricdata.Exemplar) string {
 	return cmp.Diff(got, want, cmpopts.IgnoreFields(metricdata.Exemplar{}, "Timestamp"), cmpopts.IgnoreUnexported(metricdata.Exemplar{}))
 }
+
+func TestResolveOptions(t *testing.T) {
+	k1 := tag.MustNewKey("k1")
+	k2 := tag.MustNewKey("k2")
+	m1 := stats.Int64("TestResolveOptions/m1", "", stats.UnitDimensionless)
+	m2 := stats.Int64("TestResolveOptions/m2", "", stats.UnitDimensionless)
+	v := []*view.View{{
+		Name:        "test_view",
+		TagKeys:     []tag.Key{k1, k2},
+		Measure:     m1,
+		Aggregation: view.Distribution(5, 10),
+	}, {
+		Name:        "second_view",
+		TagKeys:     []tag.Key{k1},
+		Measure:     m2,
+		Aggregation: view.Count(),
+	}}
+	view.SetReportingPeriod(100 * time.Millisecond)
+	if err := view.Register(v...); err != nil {
+		t.Fatalf("Failed to register view: %v", err)
+	}
+	defer view.Unregister(v...)
+
+	attachments := map[string]interface{}{metricdata.AttachmentKeySpanContext: spanCtx}
+	ctx, err := tag.New(context.Background(), tag.Insert(k1, "foo"), tag.Insert(k2, "foo"))
+	if err != nil {
+		t.Fatalf("Failed to set context: %v", err)
+	}
+	ro, err := stats.ResolveOptions(ctx,
+		stats.WithTags(tag.Upsert(k1, "bar"), tag.Insert(k2, "bar")),
+		stats.WithAttachments(attachments),
+		stats.WithMeasurements(m1.M(12), m2.M(5)))
+	if err != nil {
+		t.Fatalf("Failed to resolve data point: %v", err)
+	}
+
+	s, ok := ro.Attachments[metricdata.AttachmentKeySpanContext]
+	if !ok || s != spanCtx {
+		t.Errorf("Unexpected SpanContext: want %v, got %v", spanCtx, s)
+	}
+	if len(ro.Attachments) != 1 {
+		t.Errorf("Expected only one attachment (SpanContext), got %v", ro.Attachments)
+	}
+
+	if len(ro.Measures) != 2 {
+		t.Errorf("Expected two measurements, got %v", ro.Measures)
+	}
+	mWant := []stats.Measurement{m1.M(12), m2.M(5)}
+	if ro.Measures[0] != mWant[0] || ro.Measures[1] != mWant[1] {
+		t.Errorf("Unexpected measurements: want %v, got %v", mWant, ro.Measures)
+	}
+
+	// k2 was Insert() ed, and shouldn't update the value that was in the supplied context.
+	tCtx, err := tag.New(context.Background(), tag.Insert(k1, "bar"), tag.Insert(k2, "foo"))
+	if err != nil {
+		t.Fatalf("Failed to construct tWant: %v", err)
+	}
+	tWant := tag.FromContext(tCtx)
+	if ro.Tags.String() != tWant.String() {
+		t.Errorf("Unexpected tags: want %v, got %v", tWant, ro.Tags)
+	}
+}