Merge pull request #2 from SumoLogic/develop

moving drainqueue to go routine

Author: SumoSourabh

CHANGELOG.md

@@ -23,6 +23,10 @@ This repository contains SumoLogic AWS Lambda extension.
    * To know more about how to use this extension follow docs [here](https://help.sumologic.com/03Send-Data/Collect-from-Other-Data-Sources/Collect_Logs_from_AWS_Lambda_using_Lambda_Extension).
    * [AWS Lambda Extensions API](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html)
 
+## Change Log
+
+For Full Change Log, please visit [Releases](https://github.com/SumoLogic/sumologic-lambda-extensions/releases) page.
+
 [github-build-badge]: https://github.com/SumoLogic/sumologic-lambda-extensions/workflows/build-and-test/badge.svg?branch=main
 
 [github-build]: https://github.com/SumoLogic/sumologic-lambda-extensions/actions?query=workflow%3Abuild-and-test

README.md

@@ -75,7 +75,7 @@ func (cfg *LambdaExtensionConfig) setDefaults() {
 	enableFailover := os.Getenv("SUMO_ENABLE_FAILOVER")
 	logTypes := os.Getenv("SUMO_LOG_TYPES")
 	if numRetry == "" {
-		cfg.NumRetry = 0
+		cfg.NumRetry = 3
 	}
 	if logLevel == "" {
 		cfg.LogLevel = logrus.InfoLevel

lambda-extensions/config/config.go

@@ -154,7 +154,7 @@ func (s *sumoLogicClient) FlushAll(msgQueue [][]byte) error {
 	return err
 }
 
-func (s *sumoLogicClient) createCWLogLinee(item map[string]interface{}) {
+func (s *sumoLogicClient) createCWLogLine(item map[string]interface{}) {
 
 	message, ok := item["record"].(map[string]interface{})
 	if ok {
@@ -198,7 +198,7 @@ func (s *sumoLogicClient) enhanceLogs(msg responseBody) {
 			}
 			item["message"] = strings.TrimSpace(message)
 		} else if ok && logType == "platform.report" {
-			s.createCWLogLinee(item)
+			s.createCWLogLine(item)
 		}
 	}
 }

lambda-extensions/sumoclient/sumoclient.go

@@ -6,7 +6,6 @@ import (
 	"os/signal"
 	"path/filepath"
 	"syscall"
-	"time"
 
 	"github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/utils"
 
@@ -83,54 +82,32 @@ func nextEvent(ctx context.Context) (*lambdaapi.NextEventResponse, error) {
 	return nextResponse, nil
 }
 
-func flushData(ctx context.Context, DeadlineMs int64) {
-	if config.EnableFailover {
-		consumer.FlushDataQueue()
-	} else {
-		consumer.DrainQueue(ctx, DeadlineMs)
-	}
-}
-
 // processEvents is - Will block until shutdown event is received or cancelled via the context..
 func processEvents(ctx context.Context) {
-	DeadlineMs, err := runTimeAPIInit()
+	_, err := runTimeAPIInit()
 	if err != nil {
 		logger.Error("Error during Registration: ", err.Error())
 		return
 	}
-	var totalMessagedProcessed int
-	startTime := time.Now()
 	// The For loop will continue till we recieve a shutdown event.
 	for {
 		select {
 		case <-ctx.Done():
-			flushData(ctx, 0)
+			consumer.FlushDataQueue(ctx)
 			return
 		default:
-			currentMessagedProcessed := consumer.DrainQueue(ctx, DeadlineMs)
-			messagesChanged, durationComplete := utils.TotalMessagesCountChanged(totalMessagedProcessed, totalMessagedProcessed+currentMessagedProcessed, config.ProcessingSleepTime, startTime)
-			totalMessagedProcessed = totalMessagedProcessed + currentMessagedProcessed
-			// Call the next event is we reach timeout or no new message are received based on sleep time.
-			if !utils.IsTimeRemaining(DeadlineMs) || durationComplete {
-				logger.Debugf("Total Messages: %v, Current Messages: %v, messages changes: %v, duration Complete: %v, start Time: %s, Sleep Time: %s", totalMessagedProcessed, currentMessagedProcessed, messagesChanged, durationComplete, startTime, config.ProcessingSleepTime)
-				logger.Info("Waiting for Run Time API event...")
-				// This statement will freeze lambda
-				nextResponse, err := nextEvent(ctx)
-				if err != nil {
-					logger.Error("Error during Next Event call: ", err.Error())
-					return
-				}
-				// Next invoke will start from here
-				logger.Infof("Received Next Event as %s", nextResponse.EventType)
-				if nextResponse.EventType == lambdaapi.Shutdown {
-					flushData(ctx, 0)
-					return
-				}
-				DeadlineMs = nextResponse.DeadlineMs
-				totalMessagedProcessed = 0
+			go consumer.DrainQueue(ctx)
+			// This statement will freeze lambda
+			nextResponse, err := nextEvent(ctx)
+			if err != nil {
+				logger.Error("Error during Next Event call: ", err.Error())
+				return
 			}
-			if messagesChanged {
-				startTime = time.Now()
+			// Next invoke will start from here
+			logger.Infof("Received Next Event as %s", nextResponse.EventType)
+			if nextResponse.EventType == lambdaapi.Shutdown {
+				consumer.FlushDataQueue(ctx)
+				return
 			}
 		}
 	}

lambda-extensions/sumologic-extension.go

@@ -5,7 +5,6 @@ import (
 	"compress/gzip"
 	"encoding/json"
 	"errors"
-	"time"
 )
 
 //------------------Retry Logic Code-------------------------------
@@ -72,25 +71,3 @@ func PrettyPrint(v interface{}) string {
 	}
 	return string(data)
 }
-
-// IsTimeRemaining is to check is the lambda is nearing its timeout.
-func IsTimeRemaining(deadtime int64) bool {
-	t := time.Unix(deadtime, 0)
-	dif := time.Now().Sub(t)
-	if dif.Seconds() <= 10 {
-		return false
-	}
-	return true
-}
-
-// TotalMessagesCountChanged is to check is the lambda function is creating any new logs or not.
-func TotalMessagesCountChanged(totalMessages, currentMessages int, duration time.Duration, startTime time.Time) (bool, bool) {
-	if totalMessages != 0 && totalMessages == currentMessages {
-		t := time.Now().Sub(startTime)
-		if t.Milliseconds() >= duration.Milliseconds() {
-			return false, true
-		}
-		return false, false
-	}
-	return true, false
-}

lambda-extensions/utils/utils.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"sync"
 
-	"github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/utils"
-
 	cfg "github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/config"
 	sumocli "github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/sumoclient"
 
@@ -14,8 +12,8 @@ import (
 
 // TaskConsumer exposing methods every consmumer should implement
 type TaskConsumer interface {
-	FlushDataQueue()
-	DrainQueue(context.Context, int64) int
+	FlushDataQueue(context.Context)
+	DrainQueue(context.Context) int
 }
 
 // sumoConsumer to drain log from dataQueue
@@ -37,23 +35,35 @@ func NewTaskConsumer(consumerQueue chan []byte, config *cfg.LambdaExtensionConfi
 }
 
 // FlushDataQueue drains the dataqueue commpletely
-func (sc *sumoConsumer) FlushDataQueue() {
-	var rawMsgArr [][]byte
-Loop:
-	for {
-		//Receives block when the buffer is empty.
-		select {
-		case rawmsg := <-sc.dataQueue:
-			rawMsgArr = append(rawMsgArr, rawmsg)
-		default:
-			err := sc.sumoclient.FlushAll(rawMsgArr)
-			if err != nil {
-				sc.logger.Debugln("Unable to flush DataQueue", err.Error())
-				// TODO: raise alert if flush fails
+func (sc *sumoConsumer) FlushDataQueue(ctx context.Context) {
+	if sc.config.EnableFailover {
+		var rawMsgArr [][]byte
+	Loop:
+		for {
+			//Receives block when the buffer is empty.
+			select {
+			case rawmsg := <-sc.dataQueue:
+				rawMsgArr = append(rawMsgArr, rawmsg)
+			default:
+				err := sc.sumoclient.FlushAll(rawMsgArr)
+				if err != nil {
+					sc.logger.Errorln("Unable to flush DataQueue", err.Error())
+					// putting back all the msg to the queue in case of failure
+					for _, msg := range rawMsgArr {
+						sc.dataQueue <- msg
+					}
+					// TODO: raise alert if flush fails
+				}
+				close(sc.dataQueue)
+				sc.logger.Debugf("DataQueue completely drained")
+				break Loop
 			}
-			close(sc.dataQueue)
-			sc.logger.Debugf("DataQueue completely drained")
-			break Loop
+		}
+	} else {
+		// calling drainqueue (during shutdown) if failover is not enabled
+		maxCallsNeededForCompleteDraining := (len(sc.dataQueue) / sc.config.MaxConcurrentRequests) + 1
+		for i := 0; i < maxCallsNeededForCompleteDraining; i++ {
+			sc.DrainQueue(ctx)
 		}
 	}
 
@@ -64,17 +74,19 @@ func (sc *sumoConsumer) consumeTask(ctx context.Context, wg *sync.WaitGroup, raw
 	err := sc.sumoclient.SendLogs(ctx, rawmsg)
 	if err != nil {
 		sc.logger.Error("Error during Send Logs to Sumo Logic.", err.Error())
+		// putting back the msg to the queue in case of failure
+		sc.dataQueue <- rawmsg
 		// TODO: raise alert if send logs fails
 	}
 	return
 }
 
-func (sc *sumoConsumer) DrainQueue(ctx context.Context, deadtimems int64) int {
+func (sc *sumoConsumer) DrainQueue(ctx context.Context) int {
 	wg := new(sync.WaitGroup)
 	//sc.logger.Debug("Consuming data from dataQueue")
 	counter := 0
 Loop:
-	for i := 0; i < sc.config.MaxConcurrentRequests && len(sc.dataQueue) != 0 && utils.IsTimeRemaining(deadtimems); i++ {
+	for i := 0; i < sc.config.MaxConcurrentRequests && len(sc.dataQueue) != 0; i++ {
 		//Receives block when the buffer is empty.
 		select {
 		case rawmsg := <-sc.dataQueue:

lambda-extensions/workers/consumer.go

@@ -79,6 +79,6 @@ for region in "${AWS_REGIONS[@]}"; do
     echo "Layer Arn: arn:aws:lambda:${region}:<accountId>:layer:${layer_name}:${layer_version} deployed to Region ${region}"
 
     echo "Setting public permissions for layer version: ${layer_version}"
-    # aws lambda add-layer-version-permission --layer-name ${layer_name}  --statement-id ${layer_name}-prod --version-number $layer_version --principal '*' --action lambda:GetLayerVersion --region ${region}
-    aws lambda add-layer-version-permission --layer-name ${layer_name}  --statement-id ${layer_name}-dev --version-number ${layer_version} --principal '956882708938' --action lambda:GetLayerVersion --region ${region}
+    aws lambda add-layer-version-permission --layer-name ${layer_name}  --statement-id ${layer_name}-prod --version-number $layer_version --principal '*' --action lambda:GetLayerVersion --region ${region}
+    # aws lambda add-layer-version-permission --layer-name ${layer_name}  --statement-id ${layer_name}-dev --version-number ${layer_version} --principal '956882708938' --action lambda:GetLayerVersion --region ${region}
 done