Fetch Lambda Platform Metrics

apm-lambda-extension/e2e-testing/e2e_util.go

@@ -20,10 +20,6 @@ package e2eTesting
 import (
 	"archive/zip"
 	"bufio"
-	"bytes"
-	"compress/gzip"
-	"compress/zlib"
-	"elastic/apm-lambda-extension/extension"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -33,6 +29,8 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
+
+	"elastic/apm-lambda-extension/extension"
 )
 
 // GetEnvVarValueOrSetDefault retrieves the environment variable envVarName.
@@ -157,33 +155,7 @@ func GetDecompressedBytesFromRequest(req *http.Request) ([]byte, error) {
 	if req.Body != nil {
 		rawBytes, _ = ioutil.ReadAll(req.Body)
 	}
-
-	switch req.Header.Get("Content-Encoding") {
-	case "deflate":
-		reader := bytes.NewReader(rawBytes)
-		zlibreader, err := zlib.NewReader(reader)
-		if err != nil {
-			return nil, fmt.Errorf("could not create zlib.NewReader: %v", err)
-		}
-		bodyBytes, err := ioutil.ReadAll(zlibreader)
-		if err != nil {
-			return nil, fmt.Errorf("could not read from zlib reader using ioutil.ReadAll: %v", err)
-		}
-		return bodyBytes, nil
-	case "gzip":
-		reader := bytes.NewReader(rawBytes)
-		zlibreader, err := gzip.NewReader(reader)
-		if err != nil {
-			return nil, fmt.Errorf("could not create gzip.NewReader: %v", err)
-		}
-		bodyBytes, err := ioutil.ReadAll(zlibreader)
-		if err != nil {
-			return nil, fmt.Errorf("could not read from gzip reader using ioutil.ReadAll: %v", err)
-		}
-		return bodyBytes, nil
-	default:
-		return rawBytes, nil
-	}
+	return extension.GetUncompressedBytes(rawBytes, req.Header.Get("Content-Encoding"))
 }
 
 // GetFreePort is a function that queries the kernel and obtains an unused port.

apm-lambda-extension/extension/apm_server_transport.go

@@ -47,7 +47,7 @@ const (
 type ApmServerTransport struct {
 	sync.Mutex
 	bufferPool        sync.Pool
-	config            *extensionConfig
+	config            *Config
 	AgentDoneSignal   chan struct{}
 	dataChannel       chan AgentData
 	client            *http.Client
@@ -56,7 +56,7 @@ type ApmServerTransport struct {
 	gracePeriodTimer  *time.Timer
 }
 
-func InitApmServerTransport(config *extensionConfig) *ApmServerTransport {
+func InitApmServerTransport(config *Config) *ApmServerTransport {
 	var transport ApmServerTransport
 	transport.bufferPool = sync.Pool{New: func() interface{} {
 		return &bytes.Buffer{}
@@ -75,7 +75,7 @@ func InitApmServerTransport(config *extensionConfig) *ApmServerTransport {
 // StartBackgroundApmDataForwarding Receive agent data as it comes in and post it to the APM server.
 // Stop checking for, and sending agent data when the function invocation
 // has completed, signaled via a channel.
-func (transport *ApmServerTransport) ForwardApmData(ctx context.Context) error {
+func (transport *ApmServerTransport) ForwardApmData(ctx context.Context, metadataContainer *MetadataContainer) error {
 	if transport.status == Failing {
 		return nil
 	}
@@ -85,6 +85,12 @@ func (transport *ApmServerTransport) ForwardApmData(ctx context.Context) error {
 			Log.Debug("Invocation context cancelled, not processing any more agent data")
 			return nil
 		case agentData := <-transport.dataChannel:
+			if metadataContainer.Metadata == nil {
+				err := ProcessMetadata(agentData, metadataContainer)
+				if err != nil {
+					Log.Errorf("Error extracting metadata from agent payload %v", err)
+				}
+			}
 			if err := transport.PostToApmServer(ctx, agentData); err != nil {
 				return fmt.Errorf("error sending to APM server, skipping: %v", err)
 			}
@@ -151,16 +157,16 @@ func (transport *ApmServerTransport) PostToApmServer(ctx context.Context, agentD
 		r = buf
 	}
 
-	req, err := http.NewRequest("POST", transport.config.apmServerUrl+endpointURI, r)
+	req, err := http.NewRequest("POST", transport.config.ApmServerUrl+endpointURI, r)
 	if err != nil {
 		return fmt.Errorf("failed to create a new request when posting to APM server: %v", err)
 	}
 	req.Header.Add("Content-Encoding", encoding)
 	req.Header.Add("Content-Type", "application/x-ndjson")
-	if transport.config.apmServerApiKey != "" {
-		req.Header.Add("Authorization", "ApiKey "+transport.config.apmServerApiKey)
-	} else if transport.config.apmServerSecretToken != "" {
-		req.Header.Add("Authorization", "Bearer "+transport.config.apmServerSecretToken)
+	if transport.config.ApmServerApiKey != "" {
+		req.Header.Add("Authorization", "ApiKey "+transport.config.ApmServerApiKey)
+	} else if transport.config.ApmServerSecretToken != "" {
+		req.Header.Add("Authorization", "Bearer "+transport.config.ApmServerSecretToken)
 	}
 
 	Log.Debug("Sending data chunk to APM server")

apm-lambda-extension/extension/apm_server_transport_test.go

@@ -20,12 +20,13 @@ package extension
 import (
 	"compress/gzip"
 	"context"
-	"github.com/stretchr/testify/assert"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPostToApmServerDataCompressed(t *testing.T) {
@@ -66,8 +67,8 @@ func TestPostToApmServerDataCompressed(t *testing.T) {
 	}))
 	defer apmServer.Close()
 
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 	transport := InitApmServerTransport(&config)
 	err := transport.PostToApmServer(context.Background(), agentData)
@@ -111,16 +112,16 @@ func TestPostToApmServerDataNotCompressed(t *testing.T) {
 	}))
 	defer apmServer.Close()
 
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 	transport := InitApmServerTransport(&config)
 	err := transport.PostToApmServer(context.Background(), agentData)
 	assert.Equal(t, nil, err)
 }
 
 func TestGracePeriod(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 
 	transport.reconnectionCount = 0
 	val0 := transport.computeGracePeriod().Seconds()
@@ -156,7 +157,7 @@ func TestGracePeriod(t *testing.T) {
 }
 
 func TestSetHealthyTransport(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	transport.SetApmServerTransportState(context.Background(), Healthy)
 	assert.True(t, transport.status == Healthy)
 	assert.Equal(t, transport.reconnectionCount, -1)
@@ -165,15 +166,15 @@ func TestSetHealthyTransport(t *testing.T) {
 func TestSetFailingTransport(t *testing.T) {
 	// By explicitly setting the reconnection count to 0, we ensure that the grace period will not be 0
 	// and avoid a race between reaching the pending status and the test assertion.
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	transport.reconnectionCount = 0
 	transport.SetApmServerTransportState(context.Background(), Failing)
 	assert.True(t, transport.status == Failing)
 	assert.Equal(t, transport.reconnectionCount, 1)
 }
 
 func TestSetPendingTransport(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	transport.SetApmServerTransportState(context.Background(), Healthy)
 	transport.SetApmServerTransportState(context.Background(), Failing)
 	for {
@@ -186,15 +187,15 @@ func TestSetPendingTransport(t *testing.T) {
 }
 
 func TestSetPendingTransportExplicitly(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	transport.SetApmServerTransportState(context.Background(), Healthy)
 	transport.SetApmServerTransportState(context.Background(), Pending)
 	assert.True(t, transport.status == Healthy)
 	assert.Equal(t, transport.reconnectionCount, -1)
 }
 
 func TestSetInvalidTransport(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	transport.SetApmServerTransportState(context.Background(), Healthy)
 	transport.SetApmServerTransportState(context.Background(), "Invalid")
 	assert.True(t, transport.status == Healthy)
@@ -233,8 +234,8 @@ func TestEnterBackoffFromHealthy(t *testing.T) {
 			return
 		}
 	}))
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 	transport := InitApmServerTransport(&config)
 	transport.SetApmServerTransportState(context.Background(), Healthy)
@@ -286,8 +287,8 @@ func TestEnterBackoffFromFailing(t *testing.T) {
 	// Close the APM server early so that POST requests fail and that backoff is enabled
 	apmServer.Close()
 
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 
 	transport := InitApmServerTransport(&config)
@@ -339,8 +340,8 @@ func TestAPMServerRecovery(t *testing.T) {
 	}))
 	defer apmServer.Close()
 
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 
 	transport := InitApmServerTransport(&config)
@@ -392,8 +393,8 @@ func TestContinuedAPMServerFailure(t *testing.T) {
 	}))
 	apmServer.Close()
 
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 
 	transport := InitApmServerTransport(&config)
@@ -425,8 +426,8 @@ func BenchmarkPostToAPM(b *testing.B) {
 			return
 		}
 	}))
-	config := extensionConfig{
-		apmServerUrl: apmServer.URL + "/",
+	config := Config{
+		ApmServerUrl: apmServer.URL + "/",
 	}
 
 	transport := InitApmServerTransport(&config)

apm-lambda-extension/extension/client.go

@@ -23,6 +23,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"time"
 )
 
 // RegisterResponse is the body of the response for /register
@@ -34,6 +35,7 @@ type RegisterResponse struct {
 
 // NextEventResponse is the response for /event/next
 type NextEventResponse struct {
+	Timestamp          time.Time `json:"timestamp,omitempty"`
 	EventType          EventType `json:"eventType"`
 	DeadlineMs         int64     `json:"deadlineMs"`
 	RequestID          string    `json:"requestId"`

apm-lambda-extension/extension/http_server.go

@@ -29,9 +29,9 @@ func StartHttpServer(ctx context.Context, transport *ApmServerTransport) (agentD
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", handleInfoRequest(ctx, transport))
 	mux.HandleFunc("/intake/v2/events", handleIntakeV2Events(transport))
-	timeout := time.Duration(transport.config.dataReceiverTimeoutSeconds) * time.Second
+	timeout := time.Duration(transport.config.DataReceiverTimeoutSeconds) * time.Second
 	server := &http.Server{
-		Addr:           transport.config.dataReceiverServerPort,
+		Addr:           transport.config.DataReceiverServerPort,
 		Handler:        mux,
 		ReadTimeout:    timeout,
 		WriteTimeout:   timeout,

apm-lambda-extension/extension/http_server_test.go

@@ -51,12 +51,12 @@ func TestInfoProxy(t *testing.T) {
 	defer apmServer.Close()
 
 	// Create extension config and start the server
-	config := extensionConfig{
-		apmServerUrl:               apmServer.URL,
-		apmServerSecretToken:       "foo",
-		apmServerApiKey:            "bar",
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerUrl:               apmServer.URL,
+		ApmServerSecretToken:       "foo",
+		ApmServerApiKey:            "bar",
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 	agentDataServer, err := StartHttpServer(context.Background(), transport)
@@ -100,12 +100,12 @@ func TestInfoProxyErrorStatusCode(t *testing.T) {
 	defer apmServer.Close()
 
 	// Create extension config and start the server
-	config := extensionConfig{
-		apmServerUrl:               apmServer.URL,
-		apmServerSecretToken:       "foo",
-		apmServerApiKey:            "bar",
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerUrl:               apmServer.URL,
+		ApmServerSecretToken:       "foo",
+		ApmServerApiKey:            "bar",
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 
@@ -145,11 +145,11 @@ func Test_handleInfoRequest(t *testing.T) {
 `
 
 	// Create extension config
-	config := extensionConfig{
-		apmServerSecretToken:       "foo",
-		apmServerApiKey:            "bar",
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerSecretToken:       "foo",
+		ApmServerApiKey:            "bar",
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 
@@ -191,7 +191,7 @@ func (errReader) Read(_ []byte) (int, error) {
 }
 
 func Test_handleInfoRequestInvalidBody(t *testing.T) {
-	transport := InitApmServerTransport(&extensionConfig{})
+	transport := InitApmServerTransport(&Config{})
 	mux := http.NewServeMux()
 	urlPath := "/intake/v2/events"
 	mux.HandleFunc(urlPath, handleIntakeV2Events(transport))
@@ -211,10 +211,10 @@ func Test_handleIntakeV2EventsQueryParam(t *testing.T) {
 	defer apmServer.Close()
 
 	// Create extension config and start the server
-	config := extensionConfig{
-		apmServerUrl:               apmServer.URL,
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerUrl:               apmServer.URL,
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 	transport.AgentDoneSignal = make(chan struct{}, 1)
@@ -264,10 +264,10 @@ func Test_handleIntakeV2EventsNoQueryParam(t *testing.T) {
 	defer apmServer.Close()
 
 	// Create extension config and start the server
-	config := extensionConfig{
-		apmServerUrl:               apmServer.URL,
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerUrl:               apmServer.URL,
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 	transport.AgentDoneSignal = make(chan struct{}, 1)
@@ -307,10 +307,10 @@ func Test_handleIntakeV2EventsQueryParamEmptyData(t *testing.T) {
 	defer apmServer.Close()
 
 	// Create extension config and start the server
-	config := extensionConfig{
-		apmServerUrl:               apmServer.URL,
-		dataReceiverServerPort:     ":1234",
-		dataReceiverTimeoutSeconds: 15,
+	config := Config{
+		ApmServerUrl:               apmServer.URL,
+		DataReceiverServerPort:     ":1234",
+		DataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(&config)
 	transport.AgentDoneSignal = make(chan struct{}, 1)