Merge pull request #10 from livesession/feat/custom-monitored-resource

Custom monitored resources

Author: tam7t

stackdriver.go

@@ -55,6 +55,8 @@ type Sink struct {
 	prefix    string
 	taskInfo  *taskInfo
 
+	monitoredResource *monitoredrespb.MonitoredResource
+
 	mu        sync.Mutex
 	debugLogs bool
 }
@@ -104,6 +106,13 @@ type Config struct {
 	// will log additional information that is helpful when debugging errors.
 	// Optional. Defaults to false.
 	DebugLogs bool
+
+	// MonitoredResource identifies the machine/service/resource that is monitored.
+	// Different possible settings are defined here:
+	// https://cloud.google.com/monitoring/api/resources
+	//
+	// Setting a nil MonitoredResource will run a defaultMonitoredResource function.
+	MonitoredResource *monitoredrespb.MonitoredResource
 }
 
 type taskInfo struct {
@@ -123,6 +132,20 @@ type BucketFn func([]string) []float64
 // writing to stackdriver.
 type ExtractLabelsFn func([]string, string) ([]string, []metrics.Label, error)
 
+// defaultMonitoredResource returns default monitored resource
+func defaultMonitoredResource(taskInfo *taskInfo) *monitoredrespb.MonitoredResource {
+	return &monitoredrespb.MonitoredResource{
+		Type: "generic_task",
+		Labels: map[string]string{
+			"project_id": taskInfo.ProjectID,
+			"location":   taskInfo.Location,
+			"namespace":  taskInfo.Namespace,
+			"job":        taskInfo.Job,
+			"task_id":    taskInfo.TaskID,
+		},
+	}
+}
+
 // DefaultBucketer is the default BucketFn used to determing bucketing values
 // for metrics.
 func DefaultBucketer(key []string) []float64 {
@@ -211,6 +234,12 @@ func NewSink(client *monitoring.MetricClient, config *Config) *Sink {
 		s.taskInfo.TaskID = "go-" + strconv.Itoa(os.Getpid()) + "@" + hostname
 	}
 
+	if config.MonitoredResource != nil {
+		s.monitoredResource = config.MonitoredResource
+	} else {
+		s.monitoredResource = defaultMonitoredResource(s.taskInfo)
+	}
+
 	s.reset()
 
 	// run cancelable goroutine that reports on interval
@@ -291,16 +320,7 @@ func (s *Sink) report(ctx context.Context) {
 	end, rGauges, rCounters, rHistograms := s.deep()
 
 	// https://cloud.google.com/monitoring/api/resources
-	resource := &monitoredrespb.MonitoredResource{
-		Type: "generic_task",
-		Labels: map[string]string{
-			"project_id": s.taskInfo.ProjectID,
-			"location":   s.taskInfo.Location,
-			"namespace":  s.taskInfo.Namespace,
-			"job":        s.taskInfo.Job,
-			"task_id":    s.taskInfo.TaskID,
-		},
-	}
+	resource := s.monitoredResource
 
 	ts := []*monitoringpb.TimeSeries{}
 

stackdriver_test.go

@@ -30,6 +30,7 @@ import (
 	distributionpb "google.golang.org/genproto/googleapis/api/distribution"
 	"google.golang.org/genproto/googleapis/api/metric"
 	metricpb "google.golang.org/genproto/googleapis/api/metric"
+	monitoredrespb "google.golang.org/genproto/googleapis/api/monitoredres"
 	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/test/bufconn"
@@ -91,6 +92,12 @@ func BenchmarkReport10(b *testing.B)  { benchmarkCopy(10, 10, 10, b) }
 func BenchmarkReport50(b *testing.B)  { benchmarkCopy(50, 50, 50, b) }
 func BenchmarkReport100(b *testing.B) { benchmarkCopy(100, 100, 100, b) }
 
+func diffTest(t *testing.T, title string, x, y interface{}) {
+	if diff := cmp.Diff(x, y); diff != "" {
+		t.Errorf("%s mismatch (-want +got):\n%s", title, diff)
+	}
+}
+
 func sPtr(s string) *string {
 	return &s
 }
@@ -1091,3 +1098,68 @@ func diffCreateMsg(want, got *monitoringpb.CreateTimeSeriesRequest) string {
 	}
 	return out
 }
+
+func monitoredResourceDiff(t *testing.T, s *Sink, labels map[string]string) {
+	diffTest(t, "Monitored Resource labels", s.monitoredResource.GetLabels(), labels)
+}
+
+func TestCustomMonitorResource(t *testing.T) {
+	labels := map[string]string{
+		"project_id":     "project",
+		"location":       "zone",
+		"cluster_name":   "cluster",
+		"container_name": "container_name",
+		"namespace_name": "namespace_name",
+		"pod_name":       "pod_name",
+	}
+
+	sink := NewSink(nil, &Config{
+		ProjectID: "example_project",
+		Prefix:    sPtr(""),
+		MonitoredResource: &monitoredrespb.MonitoredResource{
+			Labels: labels,
+			Type:   "k8s_container",
+		},
+	})
+
+	monitoredResourceDiff(t, sink, labels)
+}
+
+func TestCustomMonitorResourceWithDefaultLabels(t *testing.T) {
+	sink := NewSink(nil, &Config{
+		ProjectID: "example_project",
+		Prefix:    sPtr(""),
+	})
+
+	labels := defaultMonitoredResource(sink.taskInfo).GetLabels()
+
+	monitoredResourceDiff(t, sink, labels)
+}
+
+func TestCustomMonitorResourceWithInvalidLabels(t *testing.T) {
+	labels := map[string]string{
+		"project_id":     "project",
+		"location":       "zone",
+		"cluster_name":   "cluster",
+		"container_name": "container_name",
+		"namespace_name": "namespace_name",
+		"pod_name":       "pod_name",
+	}
+
+	invalidLabels := map[string]string{
+		"project_id": "project",
+	}
+
+	sink := NewSink(nil, &Config{
+		ProjectID: "example_project",
+		Prefix:    sPtr(""),
+		MonitoredResource: &monitoredrespb.MonitoredResource{
+			Labels: labels,
+			Type:   "k8s_container",
+		},
+	})
+
+	if diff := cmp.Diff(sink.monitoredResource.GetLabels(), invalidLabels); diff == "" {
+		t.Error("Monitored Resource labels should not be equal")
+	}
+}