Refactor Transport State Struct

Author: jlvoiseux

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

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

apm-lambda-extension/extension/apm_server.go

@@ -36,10 +36,10 @@ var bufferPool = sync.Pool{New: func() interface{} {
 	return &bytes.Buffer{}
 }}
 
-type ApmServerTransportStatusType string
-
 // Constants for the state of the transport used in
 // the backoff implementation.
+type ApmServerTransportStatusType string
+
 const (
 	Failing ApmServerTransportStatusType = "Failing"
 	Pending ApmServerTransportStatusType = "Pending"
@@ -48,34 +48,45 @@ const (
 
 // A struct to track the state and status of sending
 // to the APM server. Used in the backoff implementation.
-type ApmServerTransportStateType struct {
+type ApmServerTransport struct {
+	sync.Pool
 	sync.Mutex
+	ctx               context.Context
+	config            *extensionConfig
+	DataChannel       chan AgentData
+	Client            *http.Client
 	Status            ApmServerTransportStatusType
 	ReconnectionCount int
 	GracePeriodTimer  *time.Timer
 }
 
-// The status of transport to the APM server.
-//
-// This instance of the ApmServerTransportStateType is public for use in tests.
-var ApmServerTransportState = ApmServerTransportStateType{
-	Status:            Healthy,
-	ReconnectionCount: -1,
+func InitApmServerTransport(ctx context.Context, config *extensionConfig) *ApmServerTransport {
+	var transport ApmServerTransport
+	transport.DataChannel = make(chan AgentData, 100)
+	transport.Client = &http.Client{
+		Timeout:   time.Duration(config.DataForwarderTimeoutSeconds) * time.Second,
+		Transport: http.DefaultTransport.(*http.Transport).Clone(),
+	}
+	transport.config = config
+	transport.ctx = ctx
+	transport.Status = Healthy
+	transport.ReconnectionCount = -1
+	return &transport
 }
 
-func StartBackgroundSending(ctx context.Context, agentDataChannel chan AgentData, client *http.Client, config *extensionConfig, funcDone chan struct{}, backgroundDataSendWg *sync.WaitGroup) {
+func StartBackgroundSending(transport *ApmServerTransport, funcDone chan struct{}, backgroundDataSendWg *sync.WaitGroup) {
 	go func() {
 		defer backgroundDataSendWg.Done()
-		if !IsTransportStatusHealthyOrPending() {
+		if transport.Status == Failing {
 			return
 		}
 		for {
 			select {
 			case <-funcDone:
 				Log.Debug("Received signal that function has completed, not processing any more agent data")
 				return
-			case agentData := <-agentDataChannel:
-				if err := PostToApmServer(client, agentData, config, ctx); err != nil {
+			case agentData := <-transport.DataChannel:
+				if err := PostToApmServer(transport, agentData); err != nil {
 					Log.Errorf("Error sending to APM server, skipping: %v", err)
 					return
 				}
@@ -89,10 +100,10 @@ func StartBackgroundSending(ctx context.Context, agentDataChannel chan AgentData
 // 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 PostToApmServer(client *http.Client, agentData AgentData, config *extensionConfig, ctx context.Context) error {
+func PostToApmServer(transport *ApmServerTransport, agentData AgentData) error {
 	// todo: can this be a streaming or streaming style call that keeps the
 	//       connection open across invocations?
-	if !IsTransportStatusHealthyOrPending() {
+	if transport.Status == Failing {
 		return errors.New("transport status is unhealthy")
 	}
 
@@ -122,89 +133,81 @@ func PostToApmServer(client *http.Client, agentData AgentData, config *extension
 		r = buf
 	}
 
-	req, err := http.NewRequest("POST", 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 config.apmServerApiKey != "" {
-		req.Header.Add("Authorization", "ApiKey "+config.apmServerApiKey)
-	} else if config.apmServerSecretToken != "" {
-		req.Header.Add("Authorization", "Bearer "+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")
-	resp, err := client.Do(req)
+	resp, err := transport.Client.Do(req)
 	if err != nil {
-		SetApmServerTransportState(Failing, ctx)
+		SetApmServerTransportState(transport, Failing)
 		return fmt.Errorf("failed to post to APM server: %v", err)
 	}
 
 	//Read the response body
 	defer resp.Body.Close()
 	body, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		SetApmServerTransportState(Failing, ctx)
+		SetApmServerTransportState(transport, Failing)
 		return fmt.Errorf("failed to read the response body after posting to the APM server")
 	}
 
-	SetApmServerTransportState(Healthy, ctx)
+	SetApmServerTransportState(transport, Healthy)
 	Log.Debug("Transport status set to healthy")
 	Log.Debugf("APM server response body: %v", string(body))
 	Log.Debugf("APM server response status code: %v", resp.StatusCode)
 	return nil
 }
 
-// IsTransportStatusHealthyOrPending returns true if the APM server transport status is
-// healthy or pending, and false otherwise.
-//
-// This function is public for use in tests.
-func IsTransportStatusHealthyOrPending() bool {
-	return ApmServerTransportState.Status != Failing
-}
-
 // SetApmServerTransportState takes a state of the APM server transport and updates
 // the current state of the transport. For a change to a failing state, the grace period
 // is calculated and a go routine is started that waits for that period to complete
 // before changing the status to "pending". This would allow a subsequent send attempt
 // to the APM server.
 //
 // This function is public for use in tests.
-func SetApmServerTransportState(status ApmServerTransportStatusType, ctx context.Context) {
+func SetApmServerTransportState(transport *ApmServerTransport, status ApmServerTransportStatusType) {
 	switch status {
 	case Healthy:
-		ApmServerTransportState.Lock()
-		ApmServerTransportState.Status = status
-		Log.Debugf("APM Server Transport status set to %s", status)
-		ApmServerTransportState.ReconnectionCount = -1
-		ApmServerTransportState.Unlock()
+		transport.Lock()
+		transport.Status = status
+		Log.Debugf("APM Server Transport status set to %s", transport.Status)
+		transport.ReconnectionCount = -1
+		transport.Unlock()
 	case Failing:
-		ApmServerTransportState.Lock()
-		ApmServerTransportState.Status = status
-		Log.Debugf("APM Server Transport status set to %s", status)
-		ApmServerTransportState.ReconnectionCount++
-		ApmServerTransportState.GracePeriodTimer = time.NewTimer(computeGracePeriod())
-		Log.Debugf("Grace period entered, reconnection count : %d", ApmServerTransportState.ReconnectionCount)
+		transport.Lock()
+		transport.Status = status
+		Log.Debugf("APM Server Transport status set to %s", transport.Status)
+		transport.ReconnectionCount++
+		transport.GracePeriodTimer = time.NewTimer(computeGracePeriod(transport))
+		Log.Debugf("Grace period entered, reconnection count : %d", transport.ReconnectionCount)
 		go func() {
 			select {
-			case <-ApmServerTransportState.GracePeriodTimer.C:
+			case <-transport.GracePeriodTimer.C:
 				Log.Debug("Grace period over - timer timed out")
-			case <-ctx.Done():
+			case <-transport.ctx.Done():
 				Log.Debug("Grace period over - context done")
 			}
-			ApmServerTransportState.Status = Pending
-			Log.Debugf("APM Server Transport status set to %s", status)
-			ApmServerTransportState.Unlock()
+			transport.Status = Pending
+			Log.Debugf("APM Server Transport status set to %s", transport.Status)
+			transport.Unlock()
 		}()
 	default:
 		Log.Errorf("Cannot set APM Server Transport status to %s", status)
 	}
 }
 
 // ComputeGracePeriod https://github.com/elastic/apm/blob/main/specs/agents/transport.md#transport-errors
-func computeGracePeriod() time.Duration {
-	gracePeriodWithoutJitter := math.Pow(math.Min(float64(ApmServerTransportState.ReconnectionCount), 6), 2)
+func computeGracePeriod(transport *ApmServerTransport) time.Duration {
+	gracePeriodWithoutJitter := math.Pow(math.Min(float64(transport.ReconnectionCount), 6), 2)
 	jitter := rand.Float64()/5 - 0.1
 	return time.Duration((gracePeriodWithoutJitter + jitter*gracePeriodWithoutJitter) * float64(time.Second))
 }