clean up env var check

Author: DylanLovesCoffee

ddlambda.go

@@ -220,12 +220,14 @@ func initializeListeners(cfg *Config) []wrapper.HandlerListener {
 	if strings.EqualFold(logLevel, "debug") || (cfg != nil && cfg.DebugLogging) {
 		logger.SetLogLevel(logger.LevelDebug)
 	}
-	extensionManager := extension.BuildExtensionManager()
+	traceConfig := cfg.toTraceConfig()
+	extensionManager := extension.BuildExtensionManager(traceConfig.UniversalInstrumentation)
 	isExtensionRunning := extensionManager.IsExtensionRunning()
+	metricsConfig := cfg.toMetricsConfig(isExtensionRunning)
 
 	// Wrap the handler with listeners that add instrumentation for traces and metrics.
-	tl := trace.MakeListener(cfg.toTraceConfig(), extensionManager)
-	ml := metrics.MakeListener(cfg.toMetricsConfig(isExtensionRunning), extensionManager)
+	tl := trace.MakeListener(traceConfig, extensionManager)
+	ml := metrics.MakeListener(metricsConfig, extensionManager)
 	return []wrapper.HandlerListener{
 		&tl, &ml,
 	}

internal/extension/extension.go

@@ -15,10 +15,10 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"strconv"
 	"time"
 
 	"github.com/DataDog/datadog-lambda-go/internal/logger"
+
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 )
 
@@ -50,27 +50,29 @@ const (
 )
 
 type ExtensionManager struct {
-	helloRoute         string
-	flushRoute         string
-	extensionPath      string
-	startInvocationUrl string
-	endInvocationUrl   string
-	httpClient         HTTPClient
-	isExtensionRunning bool
+	helloRoute                 string
+	flushRoute                 string
+	extensionPath              string
+	startInvocationUrl         string
+	endInvocationUrl           string
+	httpClient                 HTTPClient
+	isExtensionRunning         bool
+	isUniversalInstrumentation bool
 }
 
 type HTTPClient interface {
 	Do(req *http.Request) (*http.Response, error)
 }
 
-func BuildExtensionManager() *ExtensionManager {
+func BuildExtensionManager(isUniversalInstrumentation bool) *ExtensionManager {
 	em := &ExtensionManager{
-		helloRoute:         helloUrl,
-		flushRoute:         flushUrl,
-		startInvocationUrl: startInvocationUrl,
-		endInvocationUrl:   endInvocationUrl,
-		extensionPath:      extensionPath,
-		httpClient:         &http.Client{Timeout: timeout},
+		helloRoute:                 helloUrl,
+		flushRoute:                 flushUrl,
+		startInvocationUrl:         startInvocationUrl,
+		endInvocationUrl:           endInvocationUrl,
+		extensionPath:              extensionPath,
+		httpClient:                 &http.Client{Timeout: timeout},
+		isUniversalInstrumentation: isUniversalInstrumentation,
 	}
 	em.checkAgentRunning()
 	return em
@@ -85,8 +87,7 @@ func (em *ExtensionManager) checkAgentRunning() {
 		em.isExtensionRunning = true
 
 		// Tell the extension not to create an execution span if universal instrumentation is disabled
-		isUniversalInstrumentation, _ := strconv.ParseBool(os.Getenv("DD_UNIVERSAL_INSTRUMENTATION"))
-		if !isUniversalInstrumentation {
+		if !em.isUniversalInstrumentation {
 			req, _ := http.NewRequest(http.MethodGet, em.helloRoute, nil)
 			if response, err := em.httpClient.Do(req); err == nil && response.StatusCode == 200 {
 				logger.Debug("Hit the extension /hello route")

internal/extension/extension_test.go

@@ -72,12 +72,13 @@ func captureLog(f func()) string {
 }
 
 func TestBuildExtensionManager(t *testing.T) {
-	em := BuildExtensionManager()
+	em := BuildExtensionManager(false)
 	assert.Equal(t, "http://localhost:8124/lambda/hello", em.helloRoute)
 	assert.Equal(t, "http://localhost:8124/lambda/flush", em.flushRoute)
 	assert.Equal(t, "http://localhost:8124/lambda/start-invocation", em.startInvocationUrl)
 	assert.Equal(t, "http://localhost:8124/lambda/end-invocation", em.endInvocationUrl)
 	assert.Equal(t, "/opt/extensions/datadog-agent", em.extensionPath)
+	assert.Equal(t, false, em.isUniversalInstrumentation)
 	assert.NotNil(t, em.httpClient)
 }