feat(common.shim): Add batch to shim

plugins/common/shim/batch_metrics.go

@@ -0,0 +1,34 @@
+package shim
+
+import (
+	"github.com/influxdata/telegraf"
+	"sync"
+)
+
+type batchMetrics struct {
+	metrics []telegraf.Metric
+	wg      *sync.WaitGroup
+	mu      *sync.RWMutex
+}
+
+func (bm *batchMetrics) add(metric telegraf.Metric) {
+	bm.mu.Lock()
+	defer bm.mu.Unlock()
+
+	bm.metrics = append(bm.metrics, metric)
+}
+
+func (bm *batchMetrics) clear() {
+	bm.mu.Lock()
+	defer bm.mu.Unlock()
+
+	bm.wg.Add(-len(bm.metrics))
+	bm.metrics = bm.metrics[:0]
+}
+
+func (bm *batchMetrics) len() int {
+	bm.mu.RLock()
+	defer bm.mu.RUnlock()
+
+	return len(bm.metrics)
+}

plugins/common/shim/batch_metrics_test.go

@@ -0,0 +1,68 @@
+package shim
+
+import (
+	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/testutil"
+	"github.com/stretchr/testify/require"
+	"sync"
+	"testing"
+)
+
+func TestBatchMetricsAdd(t *testing.T) {
+	var wg sync.WaitGroup
+	var mu sync.RWMutex
+
+	bm := &batchMetrics{
+		metrics: make([]telegraf.Metric, 0),
+		wg:      &wg,
+		mu:      &mu,
+	}
+
+	metric := testutil.TestMetric(101, "metric1")
+
+	bm.add(metric)
+
+	testutil.RequireMetricsEqual(t, []telegraf.Metric{metric}, bm.metrics)
+}
+
+func TestBatchMetricsClear(t *testing.T) {
+	var wg sync.WaitGroup
+	var mu sync.RWMutex
+
+	wg.Add(2)
+
+	bm := &batchMetrics{
+		metrics: make([]telegraf.Metric, 0),
+		wg:      &wg,
+		mu:      &mu,
+	}
+
+	metric1 := testutil.TestMetric(101, "metric1")
+	metric2 := testutil.TestMetric(102, "metric2")
+
+	bm.add(metric1)
+	bm.add(metric2)
+
+	require.Len(t, bm.metrics, 2)
+	bm.clear()
+
+	require.Empty(t, bm.metrics)
+}
+
+func TestBatchMetricsLen(t *testing.T) {
+	var wg sync.WaitGroup
+	var mu sync.RWMutex
+
+	bm := &batchMetrics{
+		metrics: make([]telegraf.Metric, 0),
+		wg:      &wg,
+		mu:      &mu,
+	}
+
+	require.Empty(t, bm.metrics)
+
+	metric := testutil.TestMetric(101, "metric1")
+	bm.add(metric)
+
+	require.Equal(t, 1, bm.len())
+}

plugins/common/shim/goshim.go

@@ -39,6 +39,9 @@ type Shim struct {
 	Processor telegraf.StreamingProcessor
 	Output    telegraf.Output
 
+	BatchSize    int
+	BatchTimeout time.Duration
+
 	log telegraf.Logger
 
 	// streams
@@ -56,11 +59,13 @@ type Shim struct {
 // New creates a new shim interface
 func New() *Shim {
 	return &Shim{
-		metricCh: make(chan telegraf.Metric, 1),
-		stdin:    os.Stdin,
-		stdout:   os.Stdout,
-		stderr:   os.Stderr,
-		log:      logger.New("", "", ""),
+		BatchSize:    1,
+		BatchTimeout: 10 * time.Second,
+		metricCh:     make(chan telegraf.Metric, 1),
+		stdin:        os.Stdin,
+		stdout:       os.Stdout,
+		stderr:       os.Stderr,
+		log:          logger.New("", "", ""),
 	}
 }
 

plugins/common/shim/output.go

@@ -3,6 +3,9 @@ package shim
 import (
 	"bufio"
 	"fmt"
+	"os"
+	"sync"
+	"time"
 
 	"github.com/influxdata/telegraf"
 	"github.com/influxdata/telegraf/models"
@@ -36,19 +39,59 @@ func (s *Shim) RunOutput() error {
 	}
 	defer s.Output.Close()
 
-	var m telegraf.Metric
+	mCh := make(chan telegraf.Metric)
+	done := make(chan struct{})
+	batch := batchMetrics{wg: &sync.WaitGroup{}, mu: &sync.RWMutex{}}
+
+	go func() {
+		timer := time.NewTimer(s.BatchTimeout)
+		defer timer.Stop()
+
+		for {
+			select {
+			case m := <-mCh:
+				batch.add(m)
+				if batch.len() >= s.BatchSize {
+					if err = s.Output.Write(batch.metrics); err != nil {
+						fmt.Fprintf(os.Stderr, "Failed to write metrics: %s\n", err)
+					}
+					batch.clear()
+					timer.Reset(s.BatchTimeout)
+				}
+			case <-timer.C:
+				if batch.len() > 0 {
+					if err = s.Output.Write(batch.metrics); err != nil {
+						fmt.Fprintf(os.Stderr, "Failed to write metrics: %s\n", err)
+					}
+					batch.clear()
+				}
+				timer.Reset(s.BatchTimeout)
+			case <-done:
+				if batch.len() > 0 {
+					if err = s.Output.Write(batch.metrics); err != nil {
+						fmt.Fprintf(os.Stderr, "Failed to write remaining metrics: %s\n", err)
+					}
+				}
+				return
+			}
+		}
+	}()
 
 	scanner := bufio.NewScanner(s.stdin)
 	for scanner.Scan() {
-		m, err = parser.ParseLine(scanner.Text())
+		m, err := parser.ParseLine(scanner.Text())
 		if err != nil {
 			fmt.Fprintf(s.stderr, "Failed to parse metric: %s\n", err)
 			continue
 		}
-		if err = s.Output.Write([]telegraf.Metric{m}); err != nil {
-			fmt.Fprintf(s.stderr, "Failed to write metric: %s\n", err)
-		}
+
+		batch.wg.Add(1)
+		mCh <- m
 	}
 
+	batch.wg.Wait()
+
+	close(done)
+
 	return nil
 }