Experimental: Concurrently flush records in goroutines

fluent-bit-kinesis.go

@@ -22,10 +22,21 @@ import (
 
 	"github.com/aws/amazon-kinesis-firehose-for-fluent-bit/plugins"
 	"github.com/aws/amazon-kinesis-streams-for-fluent-bit/kinesis"
+	kinesisAPI "github.com/aws/aws-sdk-go/service/kinesis"
 	"github.com/fluent/fluent-bit-go/output"
 	"github.com/sirupsen/logrus"
 )
 
+const (
+	// Kinesis API Limit https://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#Kinesis.PutRecords
+	maximumRecordsPerPut = 500
+)
+
+const (
+	retries              = 6
+	concurrentRetryLimit = 4
+)
+
 var (
 	pluginInstances []*kinesis.OutputPlugin
 )
@@ -107,19 +118,65 @@ func FLBPluginInit(ctx unsafe.Pointer) int {
 
 //export FLBPluginFlushCtx
 func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int {
-	var count int
+	kinesisOutput := getPluginInstance(ctx)
+
+	curRetries := kinesisOutput.GetConcurrentRetries()
+	if curRetries > concurrentRetryLimit {
+		logrus.Infof("[kinesis] flush returning retry, too many concurrent retries (%d)\n", curRetries)
+		return output.FLB_RETRY
+	}
+
+	events, count, retCode := unpackRecords(kinesisOutput, data, length)
+	if retCode != output.FLB_OK {
+		logrus.Errorf("[kinesis] failed to unpackRecords\n")
+		return retCode
+	}
+	go flushWithRetries(kinesisOutput, tag, count, events, retries)
+	return output.FLB_OK
+}
+
+func flushWithRetries(kinesisOutput *kinesis.OutputPlugin, tag *C.char, count int, records []*kinesisAPI.PutRecordsRequestEntry, retries int) {
+	var retCode, tries int
+
+	currentRetries := kinesisOutput.GetConcurrentRetries()
+
+	for tries = 0; tries < retries; tries++ {
+		if currentRetries > 0 {
+			// Wait if other goroutines are retrying, as well as implement a progressive backoff
+			time.Sleep(time.Duration((2^currentRetries)*100) * time.Millisecond)
+		}
+
+		logrus.Debugf("[kinesis] Sending (%p) (%d) records, currentRetries=(%d)", records, len(records), currentRetries)
+		retCode = pluginConcurrentFlush(kinesisOutput, tag, count, &records)
+		if retCode != output.FLB_RETRY {
+			break
+		}
+		currentRetries = kinesisOutput.AddConcurrentRetries(1)
+		logrus.Infof("[kinesis] Going to retry with (%p) (%d) records, currentRetries=(%d)", records, len(records), currentRetries)
+	}
+	if tries > 0 {
+		kinesisOutput.AddConcurrentRetries(int64(-tries))
+	}
+	if retCode == output.FLB_ERROR {
+		logrus.Errorf("[kinesis] Failed to flush (%d) records with error", len(records))
+	}
+	if retCode == output.FLB_RETRY {
+		logrus.Errorf("[kinesis] Failed flush (%d) records after retries %d", len(records), retries)
+	}
+}
+
+func unpackRecords(kinesisOutput *kinesis.OutputPlugin, data unsafe.Pointer, length C.int) ([]*kinesisAPI.PutRecordsRequestEntry, int, int) {
 	var ret int
 	var ts interface{}
 	var timestamp time.Time
 	var record map[interface{}]interface{}
+	count := 0
+
+	records := make([]*kinesisAPI.PutRecordsRequestEntry, 0, maximumRecordsPerPut)
 
 	// Create Fluent Bit decoder
 	dec := output.NewDecoder(data, int(length))
 
-	kinesisOutput := getPluginInstance(ctx)
-	fluentTag := C.GoString(tag)
-	logrus.Debugf("[kinesis %d] Found logs with tag: %s\n", kinesisOutput.PluginID, fluentTag)
-
 	for {
 		//Extract Record
 		ret, ts, record = output.GetRecord(dec)
@@ -138,16 +195,23 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int
 			timestamp = time.Now()
 		}
 
-		retCode := kinesisOutput.AddRecord(record, &timestamp)
+		retCode := kinesisOutput.AddRecord(&records, record, &timestamp)
 		if retCode != output.FLB_OK {
-			return retCode
+			return nil, 0, retCode
 		}
+
 		count++
 	}
-	err := kinesisOutput.Flush()
-	if err != nil {
-		logrus.Errorf("[kinesis %d] %v\n", kinesisOutput.PluginID, err)
-		return output.FLB_ERROR
+
+	return records, count, output.FLB_OK
+}
+
+func pluginConcurrentFlush(kinesisOutput *kinesis.OutputPlugin, tag *C.char, count int, records *[]*kinesisAPI.PutRecordsRequestEntry) int {
+	fluentTag := C.GoString(tag)
+	logrus.Debugf("[kinesis %d] Found logs with tag: %s\n", kinesisOutput.PluginID, fluentTag)
+	retCode := kinesisOutput.Flush(records)
+	if retCode != output.FLB_OK {
+		return retCode
 	}
 	logrus.Debugf("[kinesis %d] Processed %d events with tag %s\n", kinesisOutput.PluginID, count, fluentTag)
 
@@ -157,12 +221,9 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int
 //export FLBPluginExit
 func FLBPluginExit() int {
 	// Before final exit, call Flush() for all the instances of the Output Plugin
-	for i := range pluginInstances {
-		err := pluginInstances[i].Flush()
-		if err != nil {
-			logrus.Errorf("[kinesis %d] %v\n", pluginInstances[i].PluginID, err)
-		}
-	}
+	// for i := range pluginInstances {
+	// 	pluginInstances[i].Flush(records)
+	// }
 
 	return output.FLB_OK
 }