Signal that function is complete if query param flushed=true is received

Author: estolfo

apm-lambda-extension/extension/http_server_test.go

@@ -18,11 +18,13 @@
 package extension
 
 import (
+	"bytes"
 	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"gotest.tools/assert"
 )
@@ -78,3 +80,54 @@ func TestInfoProxy(t *testing.T) {
 		resp.Body.Close()
 	}
 }
+
+func Test_handleIntakeV2EventsFlushed(t *testing.T) {
+	body := []byte(`{"metadata": {}`)
+
+	FuncDone = make(chan struct{})
+
+	// Create apm server and handler
+	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	}))
+	defer apmServer.Close()
+
+	// Create extension config and start the server
+	dataChannel := make(chan AgentData, 100)
+	config := extensionConfig{
+		apmServerUrl:               apmServer.URL,
+		dataReceiverServerPort:     ":1234",
+		dataReceiverTimeoutSeconds: 15,
+	}
+
+	StartHttpServer(dataChannel, &config)
+	defer agentDataServer.Close()
+
+	hosts, _ := net.LookupHost("localhost")
+	url := "http://" + hosts[0] + ":1234/intake/v2/events?flushed=true"
+
+	// Create a request to send to the extension
+	req, err := http.NewRequest("POST", url, bytes.NewReader(body))
+	if err != nil {
+		t.Logf("Could not create request")
+	}
+
+	// Send the request to the extension
+	client := &http.Client{}
+	go func() {
+		_, err := client.Do(req)
+		if err != nil {
+			t.Logf("Error fetching %s, [%v]", agentDataServer.Addr, err)
+			t.Fail()
+		}
+	}()
+
+	timer := time.NewTimer(1 * time.Second)
+	defer timer.Stop()
+
+	select {
+	case <-FuncDone:
+	case <-timer.C:
+		t.Log("Timed out waiting for server to send FuncDone signal")
+		t.Fail()
+	}
+}

apm-lambda-extension/extension/route_handlers.go

@@ -28,6 +28,8 @@ type AgentData struct {
 	ContentEncoding string
 }
 
+var FuncDone chan struct{}
+
 // URL: http://server/
 func handleInfoRequest(apmServerUrl string) func(w http.ResponseWriter, r *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
@@ -97,5 +99,9 @@ func handleIntakeV2Events(agentDataChan chan AgentData) func(w http.ResponseWrit
 		}
 		log.Println("Adding agent data to buffer to be sent to apm server")
 		agentDataChan <- agentData
+
+		if r.URL.Query()["flushed"][0] == "true" {
+			FuncDone <- struct{}{}
+		}
 	}
 }

apm-lambda-extension/main.go

@@ -60,12 +60,17 @@ func main() {
 	// pulls ELASTIC_ env variable into globals for easy access
 	config := extension.ProcessEnv()
 
-	// setup http server to receive data from agent
-	// and get a channel to listen for that data
+	// Create a channel to buffer apm agent data
 	agentDataChannel := make(chan extension.AgentData, 100)
 
+	// Start http server to receive data from agent
 	extension.StartHttpServer(agentDataChannel, config)
 
+	// Create a client to use for sending data to the apm server
+	client := &http.Client{
+		Transport: http.DefaultTransport.(*http.Transport).Clone(),
+	}
+
 	// Make channel for collecting logs and create a HTTP server to listen for them
 	logsChannel := make(chan logsapi.LogEvent)
 
@@ -74,7 +79,7 @@ func main() {
 		extensionClient.ExtensionID,
 		[]logsapi.EventType{logsapi.Platform})
 	if err != nil {
-		log.Printf("Could not subscribe to the logs API. Will instead flush APM data 100ms before the function deadline.")
+		log.Printf("Could not subscribe to the logs API.")
 	} else {
 		logsAPIListener, err := logsapi.NewLogsAPIHttpListener(logsChannel)
 		if err != nil {
@@ -88,9 +93,6 @@ func main() {
 		}
 	}
 
-	client := &http.Client{
-		Transport: http.DefaultTransport.(*http.Transport).Clone(),
-	}
 	for {
 		select {
 		case <-ctx.Done():
@@ -107,30 +109,27 @@ func main() {
 			}
 			log.Printf("Received event: %v\n", extension.PrettyPrint(event))
 
+			// Make a channel for signaling that all apm data has been received
+			extension.FuncDone = make(chan struct{})
+
+			// Flush any APM data, in case waiting for the FuncDone signal timed out,
+			// the agent data wasn't available yet, and we got to the next event
+			extension.FlushAPMData(client, agentDataChannel, config)
+
 			// A shutdown event indicates the execution environment is shutting down.
 			// This is usually due to inactivity.
 			if event.EventType == extension.Shutdown {
 				extension.ProcessShutdown()
 				return
 			}
 
-			// Flush any APM data, in case waiting for the runtimeDone event timed out,
-			// the agent data wasn't available yet, and we got to the next event
-			extension.FlushAPMData(client, agentDataChannel, config)
-
-			// Make a channel for signaling that a runtimeDone event has been received
-			runtimeDone := make(chan struct{})
-
-			// Make a channel for signaling that that function invocation has completed
-			funcInvocDone := make(chan struct{})
-
 			// 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.
 			go func() {
 				for {
 					select {
-					case <-funcInvocDone:
+					case <-extension.FuncDone:
 						log.Println("Function invocation is complete, not receiving any more agent data")
 						return
 					case agentData := <-agentDataChannel:
@@ -147,7 +146,7 @@ func main() {
 			go func() {
 				for {
 					select {
-					case <-funcInvocDone:
+					case <-extension.FuncDone:
 						log.Println("Function invocation is complete, not receiving any more log events")
 						return
 					case logEvent := <-logsChannel:
@@ -157,7 +156,7 @@ func main() {
 						if logsapi.SubEventType(logEvent.Type) == logsapi.RuntimeDone {
 							if logEvent.Record.RequestId == event.RequestID {
 								log.Printf("Received runtimeDone event %v", logEvent)
-								runtimeDone <- struct{}{}
+								extension.FuncDone <- struct{}{}
 								return
 							} else {
 								log.Println("Log API runtimeDone event request id didn't match")
@@ -167,7 +166,7 @@ func main() {
 				}
 			}()
 
-			// Calculate how long to wait for a runtimeDone event
+			// Calculate how long to wait for a FuncDone signal
 			flushDeadlineMs := event.DeadlineMs - 100
 			durationUntilFlushDeadline := time.Until(time.Unix(flushDeadlineMs/1000, 0))
 
@@ -176,8 +175,8 @@ func main() {
 			defer timer.Stop()
 
 			select {
-			case <-runtimeDone:
-				log.Println("Received runtimeDone event signal")
+			case <-extension.FuncDone:
+				log.Println("Received signal that function is complete")
 			case <-timer.C:
 				log.Println("Time expired waiting for runtimeDone event")
 			}
@@ -186,7 +185,7 @@ func main() {
 			extension.FlushAPMData(client, agentDataChannel, config)
 
 			// Signal that the function invocation has completed
-			close(funcInvocDone)
+			close(extension.FuncDone)
 		}
 	}
 }