Handle both object and string LogEventRecord types in json Unmarshal

apm-lambda-extension/go.mod

@@ -5,5 +5,6 @@ go 1.14
 require (
 	github.com/google/go-cmp v0.5.6 // indirect
 	github.com/pkg/errors v0.9.1
+	github.com/stretchr/testify v1.7.0
 	gotest.tools v2.2.0+incompatible
 )

apm-lambda-extension/go.sum

@@ -1,8 +1,19 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.5.6 h1:BKbKCqvP6I+rmFHt06ZmyQtvB8xAkWdhFyr0ZUNZcxQ=
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
 gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=

apm-lambda-extension/logsapi/client.go

@@ -61,6 +61,7 @@ type SubEventType string
 const (
 	// RuntimeDone event is sent when lambda function is finished it's execution
 	RuntimeDone SubEventType = "platform.runtimeDone"
+	Fault       SubEventType = "platform.fault"
 )
 
 // BufferingCfg is the configuration set for receiving logs from Logs API. Whichever of the conditions below is met first, the logs will be sent

apm-lambda-extension/logsapi/http_listener.go

@@ -25,12 +25,15 @@ import (
 	"net/http"
 	"os"
 	"time"
+
+	"github.com/pkg/errors"
 )
 
 type LogEvent struct {
-	Time   time.Time      `json:"time"`
-	Type   string         `json:"type"`
-	Record LogEventRecord `json:"record"`
+	Time      time.Time       `json:"time"`
+	Type      string          `json:"type"`
+	RawRecord json.RawMessage `json:"record"`
+	Record    LogEventRecord
 }
 
 type LogEventRecord struct {
@@ -101,10 +104,27 @@ func (h *LogsAPIHttpListener) http_handler(w http.ResponseWriter, r *http.Reques
 	if err != nil {
 		log.Println("error unmarshaling log event:", err)
 	}
-	// Send the log events to the channel
-	for _, logEvent := range logEvents {
-		h.logChannel <- logEvent
+
+	for idx := range logEvents {
+		err = logEvents[idx].unmarshalRecord()
+		if err != nil {
+			log.Printf("Error unmarshalling log event: %+v", err)
+			continue
+		}
+		h.logChannel <- logEvents[idx]
+	}
+}
+
+func (le *LogEvent) unmarshalRecord() error {
+	if SubEventType(le.Type) != Fault {
+		record := LogEventRecord{}
+		err := json.Unmarshal([]byte(le.RawRecord), &record)
+		if err != nil {
+			return errors.New("Could not unmarshal log event raw record into record")
+		}
+		le.Record = record
 	}
+	return nil
 }
 
 func (s *LogsAPIHttpListener) Shutdown() {

apm-lambda-extension/logsapi/http_listener_test.go

@@ -18,8 +18,11 @@
 package logsapi
 
 import (
+	"encoding/json"
 	"os"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestListenOnAddressWithEnvVariable(t *testing.T) {
@@ -44,3 +47,71 @@ func TestListenOnAddressDefault(t *testing.T) {
 		t.Fail()
 	}
 }
+
+func Test_unmarshalRuntimeDoneRecordObject(t *testing.T) {
+	jsonBytes := []byte(`
+	{
+		"time": "2021-10-20T08:13:03.278Z",
+		"type": "platform.runtimeDone",
+		"record": {
+			"requestId": "61c0fdeb-f013-4f2a-b627-56278f5666b8"
+		}
+	}
+	`)
+
+	var le LogEvent
+	err := json.Unmarshal(jsonBytes, &le)
+	if err != nil {
+		t.Fail()
+	}
+
+	err = le.unmarshalRecord()
+	if err != nil {
+		t.Fail()
+	}
+
+	record := LogEventRecord{RequestId: "61c0fdeb-f013-4f2a-b627-56278f5666b8"}
+	assert.Equal(t, record, le.Record)
+}
+
+func Test_unmarshalFaultRecordString(t *testing.T) {
+	jsonBytes := []byte(`
+	{
+		"time": "2021-10-20T08:13:03.278Z",
+		"type": "platform.fault",
+		"record": "Unknown application error occurred"
+	}
+	`)
+
+	var le LogEvent
+	err := json.Unmarshal(jsonBytes, &le)
+	if err != nil {
+		t.Fail()
+	}
+
+	err = le.unmarshalRecord()
+	if err != nil {
+		t.Fail()
+	}
+
+	assert.Equal(t, "\"Unknown application error occurred\"", string(le.RawRecord))
+}
+
+func Test_unmarshalRecordError(t *testing.T) {
+	jsonBytes := []byte(`
+	{
+		"time": "2021-10-20T08:13:03.278Z",
+		"type": "platform.runtimeDone",
+		"record": "Unknown application error occurred"
+	}
+	`)
+
+	var le LogEvent
+	err := json.Unmarshal(jsonBytes, &le)
+	if err != nil {
+		t.Fail()
+	}
+
+	err = le.unmarshalRecord()
+	assert.Error(t, err)
+}

apm-lambda-extension/main.go

@@ -143,7 +143,7 @@ func main() {
 						log.Println("Function invocation is complete, not receiving any more log events")
 						return
 					case logEvent := <-logsChannel:
-						log.Printf("Received log event %v\n", logEvent)
+						log.Printf("Received log event %v\n", logEvent.Type)
 						// Check the logEvent for runtimeDone and compare the RequestID
 						// to the id that came in via the Next API
 						if logsapi.SubEventType(logEvent.Type) == logsapi.RuntimeDone {