Address review comments

Author: lahsivjar

app/app.go

@@ -76,7 +76,7 @@ func New(ctx context.Context, opts ...ConfigOption) (*App, error) {
 		}
 
 		subscriptionLogStreams := []logsapi.SubscriptionType{logsapi.Platform}
-		if !c.disableFunctionLogSubscription {
+		if c.enableFunctionLogSubscription {
 			subscriptionLogStreams = append(subscriptionLogStreams, logsapi.Function)
 		}
 
@@ -85,7 +85,7 @@ func New(ctx context.Context, opts ...ConfigOption) (*App, error) {
 			logsapi.WithListenerAddress(addr),
 			logsapi.WithLogBuffer(100),
 			logsapi.WithLogger(app.logger),
-			logsapi.WithLogsAPISubscriptionTypes(subscriptionLogStreams...),
+			logsapi.WithSubscriptionTypes(subscriptionLogStreams...),
 		)
 		if err != nil {
 			return nil, err

app/config.go

@@ -20,13 +20,13 @@ package app
 import "github.com/aws/aws-sdk-go-v2/aws"
 
 type appConfig struct {
-	awsLambdaRuntimeAPI            string
-	awsConfig                      aws.Config
-	extensionName                  string
-	disableLogsAPI                 bool
-	disableFunctionLogSubscription bool
-	logLevel                       string
-	logsapiAddr                    string
+	awsLambdaRuntimeAPI           string
+	awsConfig                     aws.Config
+	extensionName                 string
+	disableLogsAPI                bool
+	enableFunctionLogSubscription bool
+	logLevel                      string
+	logsapiAddr                   string
 }
 
 // ConfigOption is used to configure the lambda extension
@@ -55,12 +55,12 @@ func WithoutLogsAPI() ConfigOption {
 	}
 }
 
-// WithoutFunctionLogSubscription disables the logs api subscription
-// to function log stream. This option will only work if LogsAPI is
-// not disabled by the WithoutLogsAPI config option.
-func WithoutFunctionLogSubscription() ConfigOption {
+// WithFunctionLogSubscription enables the logs api subscription
+// to function log stream. This option will only work if LogsAPI
+// is not disabled by the WithoutLogsAPI config option.
+func WithFunctionLogSubscription() ConfigOption {
 	return func(c *appConfig) {
-		c.disableFunctionLogSubscription = true
+		c.enableFunctionLogSubscription = true
 	}
 }
 

logsapi/client_test.go

@@ -104,7 +104,7 @@ func TestSubscribe(t *testing.T) {
 			}))
 			defer s.Close()
 
-			cOpts := append(tc.opts, logsapi.WithLogsAPIBaseURL(s.URL), logsapi.WithLogsAPISubscriptionTypes(logsapi.Platform))
+			cOpts := append(tc.opts, logsapi.WithLogsAPIBaseURL(s.URL), logsapi.WithSubscriptionTypes(logsapi.Platform))
 			c, err := logsapi.NewClient(cOpts...)
 			require.NoError(t, err)
 
@@ -147,7 +147,7 @@ func TestSubscribeAWSRequest(t *testing.T) {
 				tc.opts,
 				logsapi.WithLogsAPIBaseURL(s.URL),
 				logsapi.WithLogBuffer(1),
-				logsapi.WithLogsAPISubscriptionTypes(logsapi.Platform, logsapi.Function),
+				logsapi.WithSubscriptionTypes(logsapi.Platform, logsapi.Function),
 			)
 			c, err := logsapi.NewClient(cOpts...)
 			require.NoError(t, err)

logsapi/event.go

@@ -100,21 +100,19 @@ func (lc *Client) ProcessLogs(
 					lc.logger.Debug("Log API runtimeDone event request id didn't match")
 				}
 			case FunctionLog:
-				if metadataContainer != nil && len(metadataContainer.Metadata) > 0 {
-					lc.logger.Debug("Received function log")
-					processedLog, err := ProcessFunctionLog(
-						metadataContainer,
-						platformStartReqID,
-						invokedFnArn,
-						logEvent,
-					)
-					if err != nil {
-						lc.logger.Errorf("Error processing function log : %v", err)
-					} else {
-						apmClient.EnqueueAPMData(processedLog)
-					}
+				// TODO: @lahsivjar Buffer logs and send batches of data to APM-Server.
+				// Buffering should account for metadata being available before sending.
+				lc.logger.Debug("Received function log")
+				processedLog, err := ProcessFunctionLog(
+					metadataContainer,
+					platformStartReqID,
+					invokedFnArn,
+					logEvent,
+				)
+				if err != nil {
+					lc.logger.Errorf("Error processing function log : %v", err)
 				} else {
-					lc.logger.Warn("Function log received before metadata populated, quietly dropping log")
+					apmClient.EnqueueAPMData(processedLog)
 				}
 			}
 		case <-ctx.Done():

logsapi/functionlogs.go

@@ -18,6 +18,8 @@
 package logsapi
 
 import (
+	"errors"
+
 	"github.com/elastic/apm-aws-lambda/apmproxy"
 	"go.elastic.co/apm/v2/model"
 	"go.elastic.co/fastjson"
@@ -83,14 +85,18 @@ func (lc logContainer) MarshalFastJSON(json *fastjson.Writer) error {
 	return nil
 }
 
-// ProcessFunctionLog consumes extension event, agent metadata and log
-// event from Lambda logs API to create a payload for APM server
+// ProcessFunctionLog consumes agent metadata and log event from Lambda
+// logs API to create a payload for APM server.
 func ProcessFunctionLog(
 	metadataContainer *apmproxy.MetadataContainer,
 	requestID string,
 	invokedFnArn string,
 	log LogEvent,
 ) (apmproxy.AgentData, error) {
+	if metadataContainer == nil || len(metadataContainer.Metadata) == 0 {
+		return apmproxy.AgentData{}, errors.New("metadata is required")
+	}
+
 	lc := logContainer{
 		Log: &logLine{
 			Timestamp: model.Time(log.Time),
@@ -108,11 +114,11 @@ func ProcessFunctionLog(
 		return apmproxy.AgentData{}, err
 	}
 
-	var logData []byte
-	if metadataContainer.Metadata != nil {
-		logData = append(metadataContainer.Metadata, []byte("\n")...)
-	}
+	capacity := len(metadataContainer.Metadata) + jsonWriter.Size() + 1
+	logData := make([]byte, len(metadataContainer.Metadata), capacity)
+	copy(logData, metadataContainer.Metadata)
 
+	logData = append(logData, '\n')
 	logData = append(logData, jsonWriter.Bytes()...)
 	return apmproxy.AgentData{Data: logData}, nil
 }

logsapi/metrics.go

@@ -63,7 +63,6 @@ func (mc MetricsContainer) MarshalFastJSON(json *fastjson.Writer) error {
 }
 
 func ProcessPlatformReport(metadataContainer *apmproxy.MetadataContainer, functionData *extension.NextEventResponse, platformReport LogEvent) (apmproxy.AgentData, error) {
-	var metricsData []byte
 	metricsContainer := MetricsContainer{
 		Metrics: &model.Metrics{},
 	}
@@ -101,10 +100,11 @@ func ProcessPlatformReport(metadataContainer *apmproxy.MetadataContainer, functi
 		return apmproxy.AgentData{}, err
 	}
 
-	if metadataContainer.Metadata != nil {
-		metricsData = append(metadataContainer.Metadata, []byte("\n")...)
-	}
+	capacity := len(metadataContainer.Metadata) + jsonWriter.Size() + 1
+	metricsData := make([]byte, len(metadataContainer.Metadata), capacity)
+	copy(metricsData, metadataContainer.Metadata)
 
+	metricsData = append(metricsData, '\n')
 	metricsData = append(metricsData, jsonWriter.Bytes()...)
 	return apmproxy.AgentData{Data: metricsData}, nil
 }

logsapi/option.go

@@ -49,8 +49,8 @@ func WithLogger(logger *zap.SugaredLogger) ClientOption {
 	}
 }
 
-// WithLogsAPISubscriptionTypes sets the logstreams that the Logs API will subscribe to.
-func WithLogsAPISubscriptionTypes(types ...SubscriptionType) ClientOption {
+// WithSubscriptionTypes sets the logstreams that the Logs API will subscribe to.
+func WithSubscriptionTypes(types ...SubscriptionType) ClientOption {
 	return func(c *Client) {
 		c.logsAPISubscriptionTypes = types
 	}

main.go

@@ -47,9 +47,14 @@ func main() {
 		app.WithAWSConfig(cfg),
 	}
 
-	disableFnLog, _ := strconv.ParseBool(os.Getenv("ELASTIC_DISABLE_FUNCTION_LOG_SUBSCRIPTION"))
-	if disableFnLog {
-		appConfigs = append(appConfigs, app.WithoutFunctionLogSubscription())
+	captureLogs, err := strconv.ParseBool(os.Getenv("ELASTIC_APM_LAMBDA_CAPTURE_LOGS"))
+	// Default capture function logs to true
+	if err != nil {
+		captureLogs = true
+	}
+
+	if captureLogs {
+		appConfigs = append(appConfigs, app.WithFunctionLogSubscription())
 	}
 
 	app, err := app.New(ctx, appConfigs...)