Batch lambda logs API data before sending to APM-Server

apmproxy/apmserver.go

@@ -40,27 +40,20 @@ type jsonError struct {
 	Document string `json:"document,omitempty"`
 }
 
-// ForwardApmData receives agent data as it comes in and posts it to the APM server.
-// Stop checking for, and sending agent data when the function invocation
+// ForwardApmData receives apm data as it comes in and posts it to the APM server.
+// Stop checking for, and sending apm data when the function invocation
 // has completed, signaled via a channel.
-func (c *Client) ForwardApmData(ctx context.Context, metadataContainer *MetadataContainer) error {
+func (c *Client) ForwardApmData(ctx context.Context) error {
 	if c.IsUnhealthy() {
 		return nil
 	}
 	for {
 		select {
 		case <-ctx.Done():
 			c.logger.Debug("Invocation context cancelled, not processing any more agent data")
-			return nil
-		case agentData := <-c.DataChannel:
-			if metadataContainer.Metadata == nil {
-				metadata, err := ProcessMetadata(agentData)
-				if err != nil {
-					return fmt.Errorf("failed to extract metadata from agent payload %w", err)
-				}
-				metadataContainer.Metadata = metadata
-			}
-			if err := c.PostToApmServer(ctx, agentData); err != nil {
+			return c.sendBatch(ctx)
+		case apmData := <-c.DataChannel:
+			if err := c.forwardAPMDataByType(ctx, apmData); err != nil {
 				return fmt.Errorf("error sending to APM server, skipping: %v", err)
 			}
 		}
@@ -76,13 +69,17 @@ func (c *Client) FlushAPMData(ctx context.Context) {
 	c.logger.Debug("Flush started - Checking for agent data")
 	for {
 		select {
-		case agentData := <-c.DataChannel:
+		case apmData := <-c.DataChannel:
 			c.logger.Debug("Flush in progress - Processing agent data")
-			if err := c.PostToApmServer(ctx, agentData); err != nil {
+			if err := c.forwardAPMDataByType(ctx, apmData); err != nil {
 				c.logger.Errorf("Error sending to APM server, skipping: %v", err)
 			}
 		default:
 			c.logger.Debug("Flush ended - No agent data on buffer")
+			// Flush any remaining data in batch
+			if err := c.sendBatch(ctx); err != nil {
+				c.logger.Errorf("Error sending to APM server, skipping: %v", err)
+			}
 			return
 		}
 	}
@@ -93,7 +90,7 @@ func (c *Client) FlushAPMData(ctx context.Context) {
 // The function compresses the APM agent data, if it's not already compressed.
 // It sets the APM transport status to failing upon errors, as part of the backoff
 // strategy.
-func (c *Client) PostToApmServer(ctx context.Context, agentData AgentData) error {
+func (c *Client) PostToApmServer(ctx context.Context, agentData APMData) error {
 	// todo: can this be a streaming or streaming style call that keeps the
 	//       connection open across invocations?
 	if c.IsUnhealthy() {
@@ -281,12 +278,12 @@ func (c *Client) ComputeGracePeriod() time.Duration {
 	return time.Duration((gracePeriodWithoutJitter + jitter*gracePeriodWithoutJitter) * float64(time.Second))
 }
 
-// EnqueueAPMData adds a AgentData struct to the agent data channel, effectively queueing for a send
+// EnqueueAPMData adds a apm data struct to the agent data channel, effectively queueing for a send
 // to the APM server.
-func (c *Client) EnqueueAPMData(agentData AgentData) {
+func (c *Client) EnqueueAPMData(apmData APMData) {
 	select {
-	case c.DataChannel <- agentData:
-		c.logger.Debug("Adding agent data to buffer to be sent to apm server")
+	case c.DataChannel <- apmData:
+		c.logger.Debug("Adding APM data to buffer to be sent to apm server")
 	default:
 		c.logger.Warn("Channel full: dropping a subset of agent data")
 	}
@@ -313,3 +310,42 @@ func (c *Client) WaitForFlush() <-chan struct{} {
 	defer c.flushMutex.Unlock()
 	return c.flushCh
 }
+
+func (c *Client) forwardAPMDataByType(ctx context.Context, apmData APMData) error {
+	switch apmData.Type {
+	case Agent:
+		if c.batch == nil {
+			metadata, err := ProcessMetadata(apmData)
+			if err != nil {
+				return fmt.Errorf("failed to extract metadata from agent payload %w", err)
+			}
+			c.batch = NewBatch(c.maxBatchSize, c.maxBatchAge, metadata)
+			// broadcast that metadata is available
+			close(c.metadataAvailable)
+		}
+		return c.PostToApmServer(ctx, apmData)
+	case Lambda:
+		if c.batch == nil {
+			// This state is not possible since we are pushing back on
+			// lambda logs API until metadata is available.
+			return errors.New("unexpected state, metadata not yet set")
+		}
+		if err := c.batch.Add(apmData); err != nil {
+			c.logger.Warnf("Dropping data due to error: %v", err)
+		}
+		if c.batch.ShouldShip() {
+			return c.sendBatch(ctx)
+		}
+		return nil
+	default:
+		return errors.New("invalid apm data type")
+	}
+}
+
+func (c *Client) sendBatch(ctx context.Context) error {
+	if c.batch == nil || c.batch.Count() == 0 {
+		return nil
+	}
+	defer c.batch.Reset()
+	return c.PostToApmServer(ctx, c.batch.ToAPMData())
+}