Remove state variables

Author: jlvoiseux

apm-lambda-extension/e2e-testing/.e2e_test_config

@@ -1,2 +1,2 @@
-RUN_E2E_TESTS=true
+RUN_E2E_TESTS=false
 ELASTIC_APM_LOG_LEVEL=info

apm-lambda-extension/extension/apm_server.go

@@ -32,10 +32,6 @@ import (
 	"time"
 )
 
-var bufferPool = sync.Pool{New: func() interface{} {
-	return &bytes.Buffer{}
-}}
-
 // Constants for the state of the transport used in
 // the backoff implementation.
 type ApmServerTransportStatusType string
@@ -49,10 +45,11 @@ const (
 // A struct to track the state and status of sending
 // to the APM server. Used in the backoff implementation.
 type ApmServerTransport struct {
-	sync.Pool
 	sync.Mutex
+	bufferPool        sync.Pool
 	ctx               context.Context
 	config            *extensionConfig
+	AgentDoneSignal   chan struct{}
 	DataChannel       chan AgentData
 	Client            *http.Client
 	Status            ApmServerTransportStatusType
@@ -62,6 +59,10 @@ type ApmServerTransport struct {
 
 func InitApmServerTransport(ctx context.Context, config *extensionConfig) *ApmServerTransport {
 	var transport ApmServerTransport
+	transport.bufferPool = sync.Pool{New: func() interface{} {
+		return &bytes.Buffer{}
+	}}
+	transport.AgentDoneSignal = make(chan struct{}, 1)
 	transport.DataChannel = make(chan AgentData, 100)
 	transport.Client = &http.Client{
 		Timeout:   time.Duration(config.DataForwarderTimeoutSeconds) * time.Second,
@@ -74,7 +75,10 @@ func InitApmServerTransport(ctx context.Context, config *extensionConfig) *ApmSe
 	return &transport
 }
 
-func StartBackgroundSending(transport *ApmServerTransport, funcDone chan struct{}, backgroundDataSendWg *sync.WaitGroup) {
+// 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 StartBackgroundApmDataForwarding(transport *ApmServerTransport, funcDone chan struct{}, backgroundDataSendWg *sync.WaitGroup) {
 	go func() {
 		defer backgroundDataSendWg.Done()
 		if transport.Status == Failing {
@@ -115,10 +119,10 @@ func PostToApmServer(transport *ApmServerTransport, agentData AgentData) error {
 		r = bytes.NewReader(agentData.Data)
 	} else {
 		encoding = "gzip"
-		buf := bufferPool.Get().(*bytes.Buffer)
+		buf := transport.bufferPool.Get().(*bytes.Buffer)
 		defer func() {
 			buf.Reset()
-			bufferPool.Put(buf)
+			transport.bufferPool.Put(buf)
 		}()
 		gw, err := gzip.NewWriterLevel(buf, gzip.BestSpeed)
 		if err != nil {

apm-lambda-extension/extension/http_server.go

@@ -23,36 +23,34 @@ import (
 	"time"
 )
 
-var agentDataServer *http.Server
-
 // StartHttpServer starts the server listening for APM agent data.
-func StartHttpServer(transport *ApmServerTransport) (err error) {
+func StartHttpServer(transport *ApmServerTransport) (agentDataServer *http.Server, err error) {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", handleInfoRequest(transport))
-	mux.HandleFunc("/intake/v2/events", handleIntakeV2Events(transport.DataChannel))
+	mux.HandleFunc("/intake/v2/events", handleIntakeV2Events(transport))
 	timeout := time.Duration(transport.config.dataReceiverTimeoutSeconds) * time.Second
-	agentDataServer = &http.Server{
+	server := &http.Server{
 		Addr:           transport.config.dataReceiverServerPort,
 		Handler:        mux,
 		ReadTimeout:    timeout,
 		WriteTimeout:   timeout,
 		MaxHeaderBytes: 1 << 20,
 	}
 
-	ln, err := net.Listen("tcp", agentDataServer.Addr)
+	ln, err := net.Listen("tcp", server.Addr)
 	if err != nil {
 		return
 	}
 
 	go func() {
-		Log.Infof("Extension listening for apm data on %s", agentDataServer.Addr)
-		if err = agentDataServer.Serve(ln); err != nil {
+		Log.Infof("Extension listening for apm data on %s", server.Addr)
+		if err = server.Serve(ln); err != nil {
 			if err.Error() == "http: Server closed" {
 				Log.Debug(err)
 			} else {
 				Log.Errorf("Error upon APM data server start : %v", err)
 			}
 		}
 	}()
-	return nil
+	return server, nil
 }

apm-lambda-extension/extension/http_server_test.go

@@ -59,7 +59,8 @@ func TestInfoProxy(t *testing.T) {
 		dataReceiverTimeoutSeconds: 15,
 	}
 	transport := InitApmServerTransport(context.Background(), &config)
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -108,7 +109,8 @@ func TestInfoProxyErrorStatusCode(t *testing.T) {
 	}
 	transport := InitApmServerTransport(context.Background(), &config)
 
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -152,7 +154,8 @@ func Test_handleInfoRequest(t *testing.T) {
 	transport := InitApmServerTransport(context.Background(), &config)
 
 	// Start extension server
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -188,10 +191,10 @@ func (errReader) Read(_ []byte) (int, error) {
 }
 
 func Test_handleInfoRequestInvalidBody(t *testing.T) {
-	testChan := make(chan AgentData)
+	transport := InitApmServerTransport(context.Background(), &extensionConfig{})
 	mux := http.NewServeMux()
 	urlPath := "/intake/v2/events"
-	mux.HandleFunc(urlPath, handleIntakeV2Events(testChan))
+	mux.HandleFunc(urlPath, handleIntakeV2Events(transport))
 	req := httptest.NewRequest(http.MethodGet, urlPath, errReader(0))
 	recorder := httptest.NewRecorder()
 
@@ -202,8 +205,6 @@ func Test_handleInfoRequestInvalidBody(t *testing.T) {
 func Test_handleIntakeV2EventsQueryParam(t *testing.T) {
 	body := []byte(`{"metadata": {}`)
 
-	AgentDoneSignal = make(chan struct{})
-
 	// Create apm server and handler
 	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))
@@ -217,7 +218,8 @@ func Test_handleIntakeV2EventsQueryParam(t *testing.T) {
 	}
 	transport := InitApmServerTransport(context.Background(), &config)
 
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -245,7 +247,7 @@ func Test_handleIntakeV2EventsQueryParam(t *testing.T) {
 	defer timer.Stop()
 
 	select {
-	case <-AgentDoneSignal:
+	case <-transport.AgentDoneSignal:
 		<-transport.DataChannel
 	case <-timer.C:
 		t.Log("Timed out waiting for server to send FuncDone signal")
@@ -269,7 +271,8 @@ func Test_handleIntakeV2EventsNoQueryParam(t *testing.T) {
 	}
 	transport := InitApmServerTransport(context.Background(), &config)
 
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -298,8 +301,6 @@ func Test_handleIntakeV2EventsNoQueryParam(t *testing.T) {
 func Test_handleIntakeV2EventsQueryParamEmptyData(t *testing.T) {
 	body := []byte(``)
 
-	AgentDoneSignal = make(chan struct{})
-
 	// Create apm server and handler
 	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))
@@ -313,7 +314,8 @@ func Test_handleIntakeV2EventsQueryParamEmptyData(t *testing.T) {
 	}
 	transport := InitApmServerTransport(context.Background(), &config)
 
-	if err := StartHttpServer(transport); err != nil {
+	agentDataServer, err := StartHttpServer(transport)
+	if err != nil {
 		t.Fail()
 		return
 	}
@@ -341,7 +343,7 @@ func Test_handleIntakeV2EventsQueryParamEmptyData(t *testing.T) {
 	defer timer.Stop()
 
 	select {
-	case <-AgentDoneSignal:
+	case <-transport.AgentDoneSignal:
 	case <-timer.C:
 		t.Log("Timed out waiting for server to send FuncDone signal")
 		t.Fail()

apm-lambda-extension/extension/process_events.go

@@ -19,13 +19,15 @@ package extension
 
 import (
 	"encoding/json"
+	"net/http"
 )
 
 // ProcessShutdown processes the Shutdown event received from the
 // Lambda runtime API.
-func ProcessShutdown() {
+func ProcessShutdown(agentDataServer *http.Server, logsServer *http.Server) {
 	Log.Info("Received SHUTDOWN event, exiting")
 	agentDataServer.Close()
+	logsServer.Close()
 }
 
 // FlushAPMData reads all the apm data in the apm data channel and sends it to the APM server.

apm-lambda-extension/extension/route_handlers.go

@@ -30,7 +30,6 @@ type AgentData struct {
 	ContentEncoding string
 }
 
-var AgentDoneSignal chan struct{}
 var currApmServerTransport *ApmServerTransport
 
 // URL: http://server/
@@ -77,7 +76,7 @@ func reverseProxyErrorHandler(res http.ResponseWriter, req *http.Request, err er
 }
 
 // URL: http://server/intake/v2/events
-func handleIntakeV2Events(agentDataChan chan AgentData) func(w http.ResponseWriter, r *http.Request) {
+func handleIntakeV2Events(transport *ApmServerTransport) func(w http.ResponseWriter, r *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 
 		Log.Debug("Handling APM Data Intake")
@@ -95,11 +94,11 @@ func handleIntakeV2Events(agentDataChan chan AgentData) func(w http.ResponseWrit
 				ContentEncoding: r.Header.Get("Content-Encoding"),
 			}
 
-			EnqueueAPMData(agentDataChan, agentData)
+			EnqueueAPMData(transport.DataChannel, agentData)
 		}
 
 		if len(r.URL.Query()["flushed"]) > 0 && r.URL.Query()["flushed"][0] == "true" {
-			AgentDoneSignal <- struct{}{}
+			transport.AgentDoneSignal <- struct{}{}
 		}
 
 		w.WriteHeader(http.StatusAccepted)

apm-lambda-extension/logsapi/route_handlers.go

@@ -24,7 +24,7 @@ import (
 	"time"
 )
 
-func handleLogEventsRequest(out chan LogEvent) func(w http.ResponseWriter, r *http.Request) {
+func handleLogEventsRequest(transport *LogsTransport) func(w http.ResponseWriter, r *http.Request) {
 
 	return func(w http.ResponseWriter, r *http.Request) {
 		var logEvents []LogEvent
@@ -40,7 +40,7 @@ func handleLogEventsRequest(out chan LogEvent) func(w http.ResponseWriter, r *ht
 				w.WriteHeader(http.StatusInternalServerError)
 				continue
 			}
-			out <- logEvents[idx]
+			transport.LogsChannel <- logEvents[idx]
 		}
 	}
 }