Make extension use RuntimeDone events

Author: doguscan-namal

go-example-logs-api-extension/logsapi/client.go

@@ -41,6 +41,13 @@ const (
 	Extension EventType = "extension"
 )
 
+type SubEventType string
+
+const (
+	// RuntimeDone event is sent when lambda function is finished it's execution
+	RuntimeDone SubEventType = "platform.runtimeDone"
+)
+
 // BufferingCfg is the configuration set for receiving logs from Logs API. Whichever of the conditions below is met first, the logs will be sent
 type BufferingCfg struct {
 	// MaxItems is the maximum number of events to be buffered in memory. (default: 10000, minimum: 1000, maximum: 10000)
@@ -86,11 +93,19 @@ type Destination struct {
 	Encoding   HttpEncoding `json:"encoding"`
 }
 
+type SchemaVersion string
+
+const (
+	SchemaVersion20210318 = "2021-03-18"
+	SchemaVersionLatest   = SchemaVersion20210318
+)
+
 // SubscribeRequest is the request body that is sent to Logs API on subscribe
 type SubscribeRequest struct {
-	EventTypes   []EventType  `json:"types"`
-	BufferingCfg BufferingCfg `json:"buffering"`
-	Destination  Destination  `json:"destination"`
+	SchemaVersion SchemaVersion `json:"schemaVersion"`
+	EventTypes    []EventType   `json:"types"`
+	BufferingCfg  BufferingCfg  `json:"buffering"`
+	Destination   Destination   `json:"destination"`
 }
 
 // SubscribeResponse is the response body that is received from Logs API on subscribe
@@ -103,9 +118,10 @@ func (c *Client) Subscribe(types []EventType, bufferingCfg BufferingCfg, destina
 
 	data, err := json.Marshal(
 		&SubscribeRequest{
-			EventTypes:   types,
-			BufferingCfg: bufferingCfg,
-			Destination:  destination,
+			SchemaVersion: SchemaVersionLatest,
+			EventTypes:    types,
+			BufferingCfg:  bufferingCfg,
+			Destination:   destination,
 		})
 	if err != nil {
 		return nil, errors.WithMessage(err, "failed to marshal SubscribeRequest")

go-example-logs-api-extension/main.go

@@ -4,17 +4,18 @@
 package main
 
 import (
+	"aws-lambda-extensions/go-example-logs-api-extension/agent"
+	"aws-lambda-extensions/go-example-logs-api-extension/extension"
+	"aws-lambda-extensions/go-example-logs-api-extension/logsapi"
 	"context"
 	"fmt"
+	"github.com/golang-collections/go-datastructures/queue"
+	log "github.com/sirupsen/logrus"
 	"os"
 	"os/signal"
 	"path"
+	"strings"
 	"syscall"
-
-	"aws-lambda-extensions/go-example-logs-api-extension/agent"
-	"aws-lambda-extensions/go-example-logs-api-extension/extension"
-	"github.com/golang-collections/go-datastructures/queue"
-	log "github.com/sirupsen/logrus"
 )
 
 // INITIAL_QUEUE_SIZE is the initial size set for the synchronous logQueue
@@ -54,14 +55,16 @@ func main() {
 	// and process the logs from main goroutine (consumer)
 	logQueue := queue.New(INITIAL_QUEUE_SIZE)
 	// Helper function to empty the log queue
-	flushLogQueue := func() {
-		for !logQueue.Empty() {
+	var logsStr string = ""
+	flushLogQueue := func(force bool) {
+		for !(logQueue.Empty() && (force || strings.Contains(logsStr, string(logsapi.RuntimeDone)))) {
 			logs, err := logQueue.Get(1)
 			if err != nil {
 				logger.Error(printPrefix, err)
 				return
 			}
-			err = logsApiLogger.PushLog(fmt.Sprintf("%v", logs[0]))
+			logsStr = fmt.Sprintf("%v", logs[0])
+			err = logsApiLogger.PushLog(logsStr)
 			if err != nil {
 				logger.Error(printPrefix, err)
 				return
@@ -98,11 +101,11 @@ func main() {
 				return
 			}
 			// Flush log queue in here after waking up
-			flushLogQueue()
+			flushLogQueue(false)
 			// Exit if we receive a SHUTDOWN event
 			if res.EventType == extension.Shutdown {
 				logger.Info(printPrefix, "Received SHUTDOWN event")
-				flushLogQueue()
+				flushLogQueue(true)
 				logsApiAgent.Shutdown()
 				logger.Info(printPrefix, "Exiting")
 				return