PushMetricSetV2WithContext

metricbeat/mb/builders.go

@@ -244,6 +244,11 @@ func mustImplementFetcher(ms MetricSet) error {
 	if _, ok := ms.(PushMetricSetV2); ok {
 		ifcs = append(ifcs, "PushMetricSetV2")
 	}
+
+	if _, ok := ms.(PushMetricSetV2WithContext); ok {
+		ifcs = append(ifcs, "PushMetricSetV2WithContext")
+	}
+
 	switch len(ifcs) {
 	case 0:
 		return fmt.Errorf("MetricSet '%s/%s' does not implement an event "+

metricbeat/mb/mb.go

@@ -22,6 +22,7 @@ to implement Modules and their associated MetricSets.
 package mb
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"time"
@@ -217,6 +218,15 @@ type PushMetricSetV2 interface {
 	Run(r PushReporterV2)
 }
 
+// PushMetricSetV2Context is a MetricSet that pushes events (rather than pulling them
+// periodically via a Fetch callback). Run is invoked to start the event
+// subscription and it should block until the MetricSet is ready to stop or
+// the context is closed.
+type PushMetricSetV2WithContext interface {
+	MetricSet
+	Run(ctx context.Context, r ReporterV2)
+}
+
 // HostData contains values parsed from the 'host' configuration. Other
 // configuration data like protocols, usernames, and passwords may also be
 // used to construct this HostData data.

metricbeat/mb/module/wrapper.go

@@ -18,6 +18,7 @@
 package module
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"sync"
@@ -191,6 +192,8 @@ func (msw *metricSetWrapper) run(done <-chan struct{}, out chan<- beat.Event) {
 		ms.Run(reporter.V1())
 	case mb.PushMetricSetV2:
 		ms.Run(reporter.V2())
+	case mb.PushMetricSetV2WithContext:
+		ms.Run(&channelContext{done}, reporter.V2())
 	case mb.EventFetcher, mb.EventsFetcher,
 		mb.ReportingMetricSet, mb.ReportingMetricSetV2, mb.ReportingMetricSetV2Error:
 		msw.startPeriodicFetching(reporter)
@@ -313,6 +316,23 @@ func (r *eventReporter) V1() mb.PushReporter {
 }
 func (r *eventReporter) V2() mb.PushReporterV2 { return reporterV2{r} }
 
+// channelContext implements context.Context by wrapping a channel
+type channelContext struct {
+	done <-chan struct{}
+}
+
+func (r *channelContext) Deadline() (time.Time, bool) { return time.Time{}, false }
+func (r *channelContext) Done() <-chan struct{}       { return r.done }
+func (r *channelContext) Err() error {
+	select {
+	case <-r.done:
+		return context.Canceled
+	default:
+		return nil
+	}
+}
+func (r *channelContext) Value(key interface{}) interface{} { return nil }
+
 // reporterV1 wraps V2 to provide a v1 interface.
 type reporterV1 struct {
 	v2     mb.PushReporterV2

metricbeat/mb/testing/modules.go

@@ -54,6 +54,7 @@ that Metricbeat does it and with the same validations.
 package testing
 
 import (
+	"context"
 	"sync"
 	"testing"
 	"time"
@@ -301,7 +302,21 @@ func NewPushMetricSetV2(t testing.TB, config interface{}) mb.PushMetricSetV2 {
 
 	pushMetricSet, ok := metricSet.(mb.PushMetricSetV2)
 	if !ok {
-		t.Fatal("MetricSet does not implement PushMetricSet")
+		t.Fatal("MetricSet does not implement PushMetricSetV2")
+	}
+
+	return pushMetricSet
+}
+
+// NewPushMetricSetV2WithContext instantiates a new PushMetricSetV2WithContext
+// using the given configuration. The ModuleFactory and MetricSetFactory are
+// obtained from the global Registry.
+func NewPushMetricSetV2WithContext(t testing.TB, config interface{}) mb.PushMetricSetV2WithContext {
+	metricSet := NewMetricSet(t, config)
+
+	pushMetricSet, ok := metricSet.(mb.PushMetricSetV2WithContext)
+	if !ok {
+		t.Fatal("MetricSet does not implement PushMetricSetV2WithContext")
 	}
 
 	return pushMetricSet
@@ -345,7 +360,7 @@ func (r *capturingPushReporterV2) Done() <-chan struct{} {
 // time and returns all of the events and errors that occur during that period.
 func RunPushMetricSetV2(timeout time.Duration, waitEvents int, metricSet mb.PushMetricSetV2) []mb.Event {
 	var (
-		r      = &capturingPushReporterV2{doneC: make(chan struct{}), eventsC: make(chan mb.Event)}
+		r      = &capturingPushReporterV2{eventsC: make(chan mb.Event)}
 		wg     sync.WaitGroup
 		events []mb.Event
 	)
@@ -387,3 +402,28 @@ func RunPushMetricSetV2(timeout time.Duration, waitEvents int, metricSet mb.Push
 	wg.Wait()
 	return events
 }
+
+// RunPushMetricSetV2WithContext run the given push metricset for the specific amount of
+// time and returns all of the events that occur during that period.
+func RunPushMetricSetV2WithContext(timeout time.Duration, waitEvents int, metricSet mb.PushMetricSetV2WithContext) []mb.Event {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	r := &capturingPushReporterV2{eventsC: make(chan mb.Event)}
+
+	go metricSet.Run(ctx, r)
+
+	var events []mb.Event
+	for {
+		select {
+		case <-ctx.Done():
+			// Timeout
+			return events
+		case e := <-r.eventsC:
+			events = append(events, e)
+			if len(events) >= waitEvents {
+				cancel()
+				return events
+			}
+		}
+	}
+	return events
+}

metricbeat/module/docker/event/event.go

@@ -76,8 +76,7 @@ func New(base mb.BaseMetricSet) (mb.MetricSet, error) {
 }
 
 // Run listens for docker events and reports them
-func (m *MetricSet) Run(reporter mb.PushReporterV2) {
-	ctx, cancel := context.WithCancel(context.Background())
+func (m *MetricSet) Run(ctx context.Context, reporter mb.ReporterV2) {
 	options := types.EventsOptions{
 		Since: fmt.Sprintf("%d", time.Now().Unix()),
 	}
@@ -100,16 +99,15 @@ func (m *MetricSet) Run(reporter mb.PushReporterV2) {
 				time.Sleep(1 * time.Second)
 				break WATCH
 
-			case <-reporter.Done():
+			case <-ctx.Done():
 				m.logger.Debug("docker", "event watcher stopped")
-				cancel()
 				return
 			}
 		}
 	}
 }
 
-func (m *MetricSet) reportEvent(reporter mb.PushReporterV2, event events.Message) {
+func (m *MetricSet) reportEvent(reporter mb.ReporterV2, event events.Message) {
 	time := time.Unix(event.Time, 0)
 
 	attributes := make(map[string]string, len(event.Actor.Attributes))

metricbeat/module/docker/event/event_integration_test.go

@@ -20,6 +20,7 @@
 package event
 
 import (
+	"context"
 	"io"
 	"os"
 	"testing"
@@ -28,19 +29,18 @@ import (
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/client"
-	"golang.org/x/net/context"
 
 	"github.com/elastic/beats/auditbeat/core"
 	"github.com/elastic/beats/metricbeat/mb"
 	mbtest "github.com/elastic/beats/metricbeat/mb/testing"
 )
 
 func TestData(t *testing.T) {
-	ms := mbtest.NewPushMetricSetV2(t, getConfig())
+	ms := mbtest.NewPushMetricSetV2WithContext(t, getConfig())
 	var events []mb.Event
 	done := make(chan interface{})
 	go func() {
-		events = mbtest.RunPushMetricSetV2(10*time.Second, 1, ms)
+		events = mbtest.RunPushMetricSetV2WithContext(10*time.Second, 1, ms)
 		close(done)
 	}()