Implement elasticsearch bulk indexing

Signed-off-by: Markus Blaschke <mblaschke82@gmail.com>

Author: mblaschke

README.md

@@ -19,6 +19,7 @@ Configuration
 | `PAGERDUTY_MAX_CONNECTIONS`             | `4`                         | Maximum numbers of HTTP connections to PagerDuty API                     |
 | `ELASTICSEARCH_ADDRESS`                 | none, required              | ElasticSearch cluster addresses (multiple)                               |
 | `ELASTICSEARCH_INDEX`                   | `pagerduty`                 | Name of ElasticSearch index                                              |
+| `ELASTICSEARCH_BATCH_COUNT`             | `50`                        | Number of documents which should be indexed in one request                |
 | `ELASTICSEARCH_RETRY_COUNT`             | `5`                         | ElasticSearch request retry count                                        |
 | `ELASTICSEARCH_RETRY_DELAY`             | `5s`                        | ElasticSearch request delay for reach retry                              |
 

exporter.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/PagerDuty/go-pagerduty"
@@ -18,8 +17,8 @@ type (
 	PagerdutyElasticsearchExporter struct {
 		scrapeTime *time.Duration
 
-		elasticSearchClient    *elasticsearch.Client
-		elasticsearchIndexName string
+		elasticSearchClient     *elasticsearch.Client
+		elasticsearchIndexName  string
 		elasticsearchBatchCount int
 		elasticsearchRetryCount int
 		elasticsearchRetryDelay time.Duration
@@ -30,21 +29,21 @@ type (
 		prometheus struct {
 			incident         *prometheus.CounterVec
 			incidentLogEntry *prometheus.CounterVec
-			esRequestTotal *prometheus.CounterVec
-			esRequestRetries  *prometheus.CounterVec
-			duration          *prometheus.GaugeVec
+			esRequestTotal   *prometheus.CounterVec
+			esRequestRetries *prometheus.CounterVec
+			duration         *prometheus.GaugeVec
 		}
 	}
 
 	ElasticsearchIncident struct {
-		DocumentID  string `json:"_id,omitempty"`
+		DocumentID string `json:"_id,omitempty"`
 		Timestamp  string `json:"@timestamp,omitempty"`
 		IncidentId string `json:"@incident,omitempty"`
 		*pagerduty.Incident
 	}
 
 	ElasticsearchIncidentLog struct {
-		DocumentID  string `json:"_id,omitempty"`
+		DocumentID string `json:"_id,omitempty"`
 		Timestamp  string `json:"@timestamp,omitempty"`
 		IncidentId string `json:"@incident,omitempty"`
 		*pagerduty.LogEntry
@@ -130,6 +129,10 @@ func (e *PagerdutyElasticsearchExporter) ConnectElasticsearch(cfg elasticsearch.
 	e.elasticsearchIndexName = indexName
 }
 
+func (e *PagerdutyElasticsearchExporter) SetElasticsearchBatchCount(batchCount int) {
+	e.elasticsearchBatchCount = batchCount
+}
+
 func (e *PagerdutyElasticsearchExporter) SetElasticsearchRetry(retryCount int, retryDelay time.Duration) {
 	e.elasticsearchRetryCount = retryCount
 	e.elasticsearchRetryDelay = retryDelay
@@ -168,8 +171,19 @@ func (e *PagerdutyElasticsearchExporter) runScrape() {
 	wgProcess.Add(1)
 	go func() {
 		defer wgProcess.Done()
+
+		bulkIndexRequests := []*esapi.IndexRequest{}
 		for esIndexRequest := range esIndexRequestChannel {
-			e.doESIndexRequest(esIndexRequest)
+			bulkIndexRequests = append(bulkIndexRequests, esIndexRequest)
+
+			if len(bulkIndexRequests) >= e.elasticsearchBatchCount {
+				e.doESIndexRequestBulk(bulkIndexRequests)
+				bulkIndexRequests = []*esapi.IndexRequest{}
+			}
+		}
+
+		if len(bulkIndexRequests) >= 1 {
+			e.doESIndexRequestBulk(bulkIndexRequests)
 		}
 	}()
 
@@ -245,16 +259,53 @@ func (e *PagerdutyElasticsearchExporter) indexIncidentLogEntry(incident pagerdut
 	callback <- &req
 }
 
-func (e *PagerdutyElasticsearchExporter) doESIndexRequest(req *esapi.IndexRequest) {
+type (
+	BulkMetaIndex struct {
+		Index BulkMetaIndexIndex `json:"index,omitempty"`
+	}
+
+	BulkMetaIndexIndex struct {
+		Id    string `json:"_id,omitempty"`
+		Type  string `json:"_type,omitempty"`
+		Index string `json:"_index,omitempty"`
+	}
+)
+
+func (e *PagerdutyElasticsearchExporter) doESIndexRequestBulk(bulkRequests []*esapi.IndexRequest) {
+	var buf bytes.Buffer
+	newline := []byte("\n")
+
 	var err error
-	var res *esapi.Response
+	var resp *esapi.Response
 
 	for i := 0; i < e.elasticsearchRetryCount; i++ {
-		e.prometheus.esRequestTotal.WithLabelValues().Inc()
+		for _, indexRequest := range bulkRequests {
+			// generate bulk index action line
+			meta := BulkMetaIndex{
+				Index: BulkMetaIndexIndex{
+					Id:    indexRequest.DocumentID,
+					Type:  indexRequest.DocumentType,
+					Index: indexRequest.Index,
+				},
+			}
+			metaJson, _ := json.Marshal(meta)
+
+			// generate document line
+			document := new(bytes.Buffer)
+			document.ReadFrom(indexRequest.Body)
+
+			// generate index line
+			buf.Grow(len(metaJson) + len(newline) + document.Len() + len(newline))
+			buf.Write(metaJson)
+			buf.Write(newline)
+			buf.Write(document.Bytes())
+			buf.Write(newline)
+		}
 
-		res, err = req.Do(context.Background(), e.elasticSearchClient)
-		if err == nil {
-			res.Body.Close()
+		e.prometheus.esRequestTotal.WithLabelValues().Inc()
+		resp, err = e.elasticSearchClient.Bulk(bytes.NewReader(buf.Bytes()))
+		if err == nil && resp.StatusCode == http.StatusOK {
+			resp.Body.Close()
 
 			// success
 			return

main.go

@@ -36,15 +36,16 @@ var opts struct {
 	ScrapeTime time.Duration `long:"scrape-time"  env:"SCRAPE_TIME"   description:"Scrape time (time.duration)" default:"5m"`
 
 	// PagerDuty settings
-	PagerDutyAuthToken      string        `long:"pagerduty.authtoken"        env:"PAGERDUTY_AUTH_TOKEN"  description:"PagerDuty auth token" required:"true"`
-	PagerDutySince          time.Duration `long:"pagerduty.date-range"       env:"PAGERDUTY_DATE_RANGE"  description:"PagerDuty date range" default:"168h"`
-	PagerDutyMaxConnections int           `long:"pagerduty.max-connections"  env:"PAGERDUTY_MAX_CONNECTIONS"                    description:"Maximum numbers of TCP connections to PagerDuty API (concurrency)" default:"4"`
+	PagerDutyAuthToken      string        `long:"pagerduty.authtoken"        env:"PAGERDUTY_AUTH_TOKEN"        description:"PagerDuty auth token" required:"true"`
+	PagerDutySince          time.Duration `long:"pagerduty.date-range"       env:"PAGERDUTY_DATE_RANGE"        description:"PagerDuty date range" default:"168h"`
+	PagerDutyMaxConnections int           `long:"pagerduty.max-connections"  env:"PAGERDUTY_MAX_CONNECTIONS"   description:"Maximum numbers of TCP connections to PagerDuty API (concurrency)" default:"4"`
 
 	// ElasticSearch settings
 	ElasticsearchAddresses  []string      `long:"elasticsearch.address"      env:"ELASTICSEARCH_ADDRESS"  delim:" "  description:"ElasticSearch urls" required:"true"`
 	ElasticsearchIndex      string        `long:"elasticsearch.index"        env:"ELASTICSEARCH_INDEX"               description:"ElasticSearch index name" default:"pagerduty"`
-	ElasticsearchRetryCount int           `long:"elasticsearch.retry-count"  env:"ELASTICSEARCH_RETRY_COUNT"         description:"ElasticSearch request retry count" default:"5"`
-	ElasticsearchRetryDelay time.Duration `long:"elasticsearch.retry-delay"  env:"ELASTICSEARCH_RETRY_DELAY"         description:"ElasticSearch request delay for reach retry" default:"5s"`
+	ElasticsearchBatchCount int           `long:"elasticsearch.batch-count"  env:"ELASTICSEARCH_BATCH_COUNT"         description:"Number of documents which should be indexed in one request"  default:"50"`
+	ElasticsearchRetryCount int           `long:"elasticsearch.retry-count"  env:"ELASTICSEARCH_RETRY_COUNT"         description:"ElasticSearch request retry count"                           default:"5"`
+	ElasticsearchRetryDelay time.Duration `long:"elasticsearch.retry-delay"  env:"ELASTICSEARCH_RETRY_DELAY"         description:"ElasticSearch request delay for reach retry"                 default:"5s"`
 }
 
 func main() {
@@ -87,6 +88,7 @@ func main() {
 		Addresses: opts.ElasticsearchAddresses,
 	}
 	exporter.ConnectElasticsearch(cfg, opts.ElasticsearchIndex)
+	exporter.SetElasticsearchBatchCount(opts.ElasticsearchBatchCount)
 	exporter.SetElasticsearchRetry(opts.ElasticsearchRetryCount, opts.ElasticsearchRetryDelay)
 	exporter.Run()