Merge pull request #11 from SumoLogic/hpal_SUMO-175893_delay_fix

Delay Fix in Lambda Extension

Author: himanshu219

README.md

@@ -13,11 +13,11 @@ The Sumo Logic lambda extension is available as an AWS public Layer. The latest
 
 For x86_64 use:
 
-    arn:aws:lambda:<AWS_REGION>:956882708938:layer:sumologic-extension-amd64:2
+    arn:aws:lambda:<AWS_REGION>:956882708938:layer:sumologic-extension-amd64:3
 
 For arm64 use:
 
-    arn:aws:lambda:<AWS_REGION>:956882708938:layer:sumologic-extension-arm64:2
+    arn:aws:lambda:<AWS_REGION>:956882708938:layer:sumologic-extension-arm64:3
 
 
 - AWS_REGION - Replace with your AWS Lambda Region.
@@ -60,3 +60,31 @@ For Full Change Log, please visit [Releases](https://github.com/SumoLogic/sumolo
 [github-release-badge]: https://img.shields.io/github/release/sumologic/sumologic-lambda-extensions/all.svg?label=release
 
 [github-release]: https://github.com/sumologic/sumologic-lambda-extensions/releases/latest
+
+
+## Compiling
+   
+   `env GOOS=darwin go build -o "sumologic-extension" "lambda-extensions/sumologic-extension.go"`
+
+
+## Building
+   This script assumes you have aws cli already configured.
+
+   - Go to scripts folder
+   - Export Profile export AWS_PROFILE=<sumo content profile>
+   - Change the layer_name variable in zip.sh to avoid replacing the prod.
+   - Run below command
+     `sh zip.sh`
+
+## Testing
+
+   1> Unit Testing locally
+      
+    - Go to root folder and run "go test  ./..."
+
+    - Go to lambda-extensions folder and run "go test  ./..."
+
+   2> Testing with Lambda function
+
+   Add the layer arn generated from build command output to your lambda function by following instructions in [docs](https://help.sumologic.com/03Send-Data/Collect-from-Other-Data-Sources/Collect_AWS_Lambda_Logs_using_an_Extension).
+

lambda-extensions/config/config.go

@@ -37,7 +37,7 @@ type LambdaExtensionConfig struct {
 	SourceCategoryOverride string
 }
 
-var validLogTypes = []string{"platform", "function", "extension"}
+var validLogTypes = []string{"platform", "function"}
 
 // GetConfig to get config instance
 func GetConfig() (*LambdaExtensionConfig, error) {

lambda-extensions/lambdaapi/logsapiclient.go

@@ -12,8 +12,8 @@ const (
 	logsURL = "2020-08-15/logs"
 	// Subscription Body Constants. Subscribe to platform logs and receive them on ${local_ip}:4243 via HTTP protocol.
 	timeoutMs    = 1000
-	maxBytes     = 262144
-	maxItems     = 1000
+	maxBytes     = 1048576
+	maxItems     = 10000
 	receiverPort = 4243
 )
 
@@ -22,9 +22,10 @@ func (client *Client) SubscribeToLogsAPI(ctx context.Context, logEvents []string
 	URL := client.baseURL + logsURL
 
 	reqBody, err := json.Marshal(map[string]interface{}{
-		"destination": map[string]interface{}{"protocol": "HTTP", "URI": fmt.Sprintf("http://sandbox:%v", receiverPort)},
-		"types":       logEvents,
-		"buffering":   map[string]interface{}{"timeoutMs": timeoutMs, "maxBytes": maxBytes, "maxItems": maxItems},
+		"destination":   map[string]interface{}{"protocol": "HTTP", "URI": fmt.Sprintf("http://sandbox:%v", receiverPort)},
+		"types":         logEvents,
+		"buffering":     map[string]interface{}{"timeoutMs": timeoutMs, "maxBytes": maxBytes, "maxItems": maxItems},
+		"schemaVersion": "2021-03-18",
 	})
 	if err != nil {
 		return nil, err

lambda-extensions/sumoclient/sumoclient.go

@@ -23,6 +23,7 @@ var isColdStart = true
 // LogSender interface which needs to be implemented to send logs
 type LogSender interface {
 	SendLogs(context.Context, []byte) error
+	SendAllLogs(context.Context, [][]byte) error
 	FlushAll([][]byte) error
 }
 
@@ -279,8 +280,61 @@ func (s *sumoLogicClient) SendLogs(ctx context.Context, rawmsg []byte) error {
 	return nil
 }
 
+func (s *sumoLogicClient) SendAllLogs(ctx context.Context, allMessages [][]byte) error {
+	if (len(allMessages) == 0) {
+		s.logger.Debugf("SendAllLogs: No messages to send")
+		return nil
+	}
+
+	s.logger.Debugf("SendAllLogs: Attempting to send %d payloads from dataqueue to SumoLogic", len(allMessages))
+
+	var errorCount int = 0
+	var totalitems int = 0
+	var payload responseBody
+	for _, rawmsg := range allMessages {
+		// converting to arr of maps
+		msgArr, err := s.transformBytesToArrayOfMap(rawmsg)
+		if err != nil {
+			s.logger.Error("SendAllLogs: Error in transforming bytes to array of struct", err.Error())
+			errorCount++
+			continue
+		}
+
+		if len(msgArr) > 0 {
+			// enhancing logs
+			s.enhanceLogs(msgArr)
+			totalitems += len(msgArr)
+			// converting back to string
+			for _, item := range msgArr {
+				payload = append(payload, item)
+			}
+		}
+	}
+	s.logger.Debugf("SendAllLogs: Enhanced TotalLogItems - %d \n", totalitems)
+	// converting back to chunks of string
+	chunks, err := s.createChunks(payload)
+	if err != nil {
+		return fmt.Errorf("SendAllLogs: CreateChunks failed - %v", err)
+	}
+	for _, strobj := range chunks {
+		err := s.postToSumo(ctx, &strobj)
+		if err != nil {
+			errorCount++
+		}
+	}
+	if errorCount > 0 {
+		err = fmt.Errorf("SendAllLogs: Errors during postToSumo - %d", errorCount)
+		return err
+	} else {
+		s.logger.Debugf("SendAllLogs: Sent TotalChunks - %d \n", totalitems)
+	}
+
+	return nil
+}
+
 func (s *sumoLogicClient) postToSumo(ctx context.Context, logStringToSend *string) error {
-	s.logger.Debug("Attempting to send to Sumo Endpoint")
+
+	s.logger.Debug("postToSumo: Attempting to send to Sumo Endpoint")
 
 	// compressing here because Sumo recommends payload size of 1MB before compression
 	bytedata := utils.Compress(logStringToSend)
@@ -295,39 +349,39 @@ func (s *sumoLogicClient) postToSumo(ctx context.Context, logStringToSend *strin
 		defer response.Body.Close()
 	}
 	if (err != nil) || (response.StatusCode != 200 && response.StatusCode != 302 && response.StatusCode < 500) {
-		s.logger.Errorf("Not able to post statuscode:  %v %v\n", err, response)
+		s.logger.Errorf("postToSumo: Not able to post statuscode -  %v %v\n", err, response)
 		err := utils.Retry(func(attempt int) (bool, error) {
-			s.logger.Debugf("Waiting for %v ms for retry attempt: %v\n", s.config.RetrySleepTime, attempt)
+			s.logger.Debugf("postToSumo: Waiting for %v ms for retry attempt - %v\n", s.config.RetrySleepTime, attempt)
 			time.Sleep(s.config.RetrySleepTime)
 			buf := createBuffer()
 			retryResponse, errRetry := s.makeRequest(ctx, buf)
 			if (errRetry != nil) || (retryResponse.StatusCode != 200 && retryResponse.StatusCode != 302 && retryResponse.StatusCode < 500) {
 				if errRetry == nil {
 					errRetry = fmt.Errorf("statuscode %v", retryResponse.StatusCode)
 				}
-				s.logger.Error("Not able to post: ", errRetry)
+				s.logger.Error("postToSumo: Not able to post - ", errRetry)
 				return attempt < s.config.MaxRetryAttempts, errRetry
 			} else if retryResponse.StatusCode == 200 {
-				s.logger.Debugf("Post of logs successful after retry %v attempts\n", attempt)
+				s.logger.Debugf("postToSumo: Post of logs successful after retry %v attempts\n", attempt)
 				return true, nil
 			}
 			return attempt < s.config.MaxRetryAttempts, errRetry
 		}, s.config.NumRetry)
 		if err != nil {
-			s.logger.Error("Finished retrying Error: ", err)
+			s.logger.Error("postToSumo: Finished retrying Error - ", err)
 			if s.config.EnableFailover {
 				buf = createBuffer()
 				err := s.failoverHandler(buf)
 				if err != nil {
-					s.logger.Errorf("Dropping messages as post to S3 failed: %v\n", err)
+					s.logger.Errorf("postToSumo: Dropping messages as post to S3 failed - %v\n", err)
 					return err
 				}
 			} else {
-				s.logger.Info("Dropping messages as no failover enabled.")
+				s.logger.Info("postToSumo: Dropping messages as no failover enabled.")
 			}
 		}
 	} else if response.StatusCode == 200 {
-		s.logger.Debugf("Post of logs successful")
+		s.logger.Debugf("postToSumo: Post of logs successful")
 	}
 
 	return nil

lambda-extensions/sumoclient/sumoclient_test.go

@@ -119,6 +119,8 @@ func TestSumoClient(t *testing.T) {
 	var reportLogs = []byte(`[{"record":{"metrics":{"billedDurationMs":120000,"durationMs":122066.85,"maxMemoryUsedMB":74,"memorySizeMB":128},"requestId":"fcea12d9-e0b4-43b2-a9a2-04d04519539f"},"time":"2020-11-02T20:33:16.536Z","type":"platform.report"}]`)
 	assertEqual(t, client.SendLogs(ctx, reportLogs), nil, "SendLogs should not generate error")
 
+	t.Log("\ntesting SendAllLogs\n======================")
+	assertEqual(t, client.SendAllLogs(ctx, multiplelargedata), nil, "SendAllLogs should not generate error")
 	//Todo remove this function from sumologic-extension
 	// t.Log("\ntesting sumo if no s3 failover\n=================")
 	// config.EnableFailover = false

lambda-extensions/sumologic-extension.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"runtime"
 	"syscall"
 
 	"github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/utils"
@@ -22,12 +23,18 @@ var (
 	extensionClient = lambdaapi.NewClient(os.Getenv("AWS_LAMBDA_RUNTIME_API"), extensionName)
 	logger          = logrus.New().WithField("Name", extensionName)
 )
+
 var producer workers.TaskProducer
 var consumer workers.TaskConsumer
 var config *cfg.LambdaExtensionConfig
 var dataQueue chan []byte
 
 func init() {
+	Formatter := new(logrus.TextFormatter)
+	Formatter.TimestampFormat = "2006-01-02T15:04:05.999999999Z07:00"
+	Formatter.FullTimestamp = true
+	logger.Logger.SetFormatter(Formatter)
+
 	logger.Logger.SetOutput(os.Stdout)
 
 	// Creating config and performing validation
@@ -89,14 +96,27 @@ func processEvents(ctx context.Context) {
 		logger.Error("Error during Registration: ", err.Error())
 		return
 	}
+
 	// The For loop will continue till we recieve a shutdown event.
 	for {
 		select {
 		case <-ctx.Done():
 			consumer.FlushDataQueue(ctx)
 			return
 		default:
-			go consumer.DrainQueue(ctx)
+			logger.Infof("Calling DrainQueue from processEvents")
+			for {
+				runtime_done := consumer.DrainQueue(ctx)
+
+				if runtime_done == 1 {
+					logger.Infof("Exiting DrainQueueLoop: Runtime is Done")
+					break
+				} else {
+					logger.Debugf("switching to other go routine")
+					runtime.Gosched()
+				}
+			}
+
 			// This statement will freeze lambda
 			nextResponse, err := nextEvent(ctx)
 			if err != nil {
@@ -106,7 +126,7 @@ func processEvents(ctx context.Context) {
 			// Next invoke will start from here
 			logger.Infof("Received Next Event as %s", nextResponse.EventType)
 			if nextResponse.EventType == lambdaapi.Shutdown {
-				consumer.FlushDataQueue(ctx)
+				consumer.DrainQueue(ctx)
 				return
 			}
 		}

lambda-extensions/workers/consumer.go

@@ -2,6 +2,8 @@ package workers
 
 import (
 	"context"
+	"fmt"
+	"strings"
 	"sync"
 
 	cfg "github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/config"
@@ -10,6 +12,13 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+type SubEventType string
+
+const (
+	// RuntimeDone event is sent when lambda function is finished it's execution
+	RuntimeDone SubEventType = "platform.runtimeDone"
+)
+
 // TaskConsumer exposing methods every consmumer should implement
 type TaskConsumer interface {
 	FlushDataQueue(context.Context)
@@ -82,24 +91,38 @@ func (sc *sumoConsumer) consumeTask(ctx context.Context, wg *sync.WaitGroup, raw
 }
 
 func (sc *sumoConsumer) DrainQueue(ctx context.Context) int {
-	wg := new(sync.WaitGroup)
 	//sc.logger.Debug("Consuming data from dataQueue")
-	counter := 0
+
+	var rawMsgArr [][]byte
+	var logsStr string = ""
+	var runtime_done int = 0
 Loop:
-	for i := 0; i < sc.config.MaxConcurrentRequests && len(sc.dataQueue) != 0; i++ {
+	for {
 		//Receives block when the buffer is empty.
 		select {
 		case rawmsg := <-sc.dataQueue:
-			counter++
-			wg.Add(1)
-			go sc.consumeTask(ctx, wg, rawmsg)
+			rawMsgArr = append(rawMsgArr, rawmsg)
+			logsStr = fmt.Sprintf("%s", rawmsg)
+			sc.logger.Debugf("DrainQueue: logsStr: %s", logsStr)
+			if strings.Contains(logsStr, string(RuntimeDone)) {
+				runtime_done = 1
+			}
+
 		default:
-			sc.logger.Debugf("DataQueue completely drained")
+			err := sc.sumoclient.SendAllLogs(ctx, 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
+			} else {
+				sc.logger.Debugf("DrainQueue: DataQueue completely drained")
+			}
 			break Loop
 		}
-
 	}
-	//sc.logger.Debugf("Waiting for %d consumer to finish their tasks", counter)
-	wg.Wait()
-	return counter
+	sc.logger.Debugf("DrainQueue: Runtime done or not? %d", runtime_done)
+	return runtime_done
 }

lambda-extensions/workers/producer.go

@@ -48,14 +48,17 @@ func (httpServer *httpServer) logsHandler(writer http.ResponseWriter, request *h
 	}
 	switch request.Method {
 	case "POST":
+		defer request.Body.Close()
 		reqBody, err := ioutil.ReadAll(request.Body)
 		if err != nil {
 			// TODO: raise alert if read fails
 			httpServer.logger.Error("Read from Logs API failed: ", err.Error())
 		}
-		httpServer.logger.Debug("Producing data into dataQueue")
+
+		httpServer.logger.Debugf("Producing data into dataQueue - %d \n", len(reqBody))
 		payload := []byte(reqBody)
 		// Sends to a buffered channel block only when the buffer is full
 		httpServer.dataQueue <- payload
+		writer.WriteHeader(http.StatusOK)
 	}
 }

scripts/zip.sh

@@ -89,7 +89,7 @@ for arch in "${ARCHITECTURES[@]}"; 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}-prod --version-number $layer_version --principal '*' --action lambda:GetLayerVersion --region ${region} --profile ${AWS_PROFILE}
       # 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