enable self-insrumentation tracing

Author: 1pkg

internal/beatcmd/beat_test.go

@@ -179,7 +179,7 @@ func TestRunManager(t *testing.T) {
 			},
 		},
 		"instrumentation": map[string]interface{}{
-			"enabled":     false,
+			"enabled":     true,
 			"environment": "testenv",
 		},
 	}, m)

internal/beatcmd/reloader.go

@@ -195,8 +195,7 @@ func (r *Reloader) reload(inputConfig, outputConfig, apmTracingConfig *config.C)
 			return fmt.Errorf("APM tracing config for elastic not found")
 		}
 		// set enabled manually as APMConfig doesn't contain it.
-		// TODO set "enable" to true after the issue https://github.com/elastic/elastic-agent/issues/5211 gets resolved.
-		c.SetBool("enabled", -1, false)
+		c.SetBool("enabled", -1, true)
 		wrappedApmTracingConfig = config.MustNewConfigFrom(map[string]interface{}{
 			"instrumentation": c,
 		})

internal/beatcmd/reloader_test.go

@@ -182,7 +182,7 @@ func TestReloaderNewRunnerParams(t *testing.T) {
 	args := <-calls
 	assert.NotNil(t, args.Logger)
 	assert.Equal(t, info, args.Info)
-	assert.Equal(t, config.MustNewConfigFrom(`{"revision": 1, "input": 123, "output.console.enabled": true, "instrumentation.enabled":false, "instrumentation.environment":"test"}`), args.Config)
+	assert.Equal(t, config.MustNewConfigFrom(`{"revision": 1, "input": 123, "output.console.enabled": true, "instrumentation.enabled":true, "instrumentation.environment":"test"}`), args.Config)
 }
 
 func expectNoEvent(t testing.TB, ch <-chan struct{}, message string) {

internal/beater/beater.go

@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"os"
 	"runtime"
+	"strconv"
 	"time"
 
 	"github.com/dustin/go-humanize"
@@ -525,6 +526,7 @@ func newInstrumentation(rawConfig *agentconfig.C) (instrumentation.Instrumentati
 			ServerCertificate string `config:"servercert"`
 			ServerCA          string `config:"serverca"`
 		} `config:"tls"`
+		SamplingRate *float32 `config:"samplingrate"`
 	}
 	cfg, err := rawConfig.Child("instrumentation", -1)
 	if err != nil || !cfg.Enabled() {
@@ -541,6 +543,7 @@ func newInstrumentation(rawConfig *agentconfig.C) (instrumentation.Instrumentati
 		envServerCert       = "ELASTIC_APM_SERVER_CERT"
 		envCACert           = "ELASTIC_APM_SERVER_CA_CERT_FILE"
 		envGlobalLabels     = "ELASTIC_APM_GLOBAL_LABELS"
+		envSamplingRate     = "ELASTIC_APM_TRANSACTION_SAMPLE_RATE"
 	)
 	if apmCfg.APIKey != "" {
 		os.Setenv(envAPIKey, apmCfg.APIKey)
@@ -566,6 +569,11 @@ func newInstrumentation(rawConfig *agentconfig.C) (instrumentation.Instrumentati
 		os.Setenv(envGlobalLabels, apmCfg.GlobalLabels)
 		defer os.Unsetenv(envGlobalLabels)
 	}
+	if apmCfg.SamplingRate != nil {
+		r := max(min(*apmCfg.SamplingRate, 1.0), 0.0)
+		os.Setenv(envSamplingRate, strconv.FormatFloat(float64(r), 'f', -1, 32))
+		defer os.Unsetenv(envSamplingRate)
+	}
 	return instrumentation.New(rawConfig, "apm-server", version.Version)
 }
 

internal/beater/beater_test.go

@@ -24,10 +24,12 @@ import (
 	"encoding/pem"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -285,6 +287,46 @@ func TestNewInstrumentation(t *testing.T) {
 	assert.Equal(t, "Bearer secret", auth)
 }
 
+func TestNewInstrumentationWithSampling(t *testing.T) {
+	runSampled := func(rate float32) {
+		var events int
+		s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/intake/v2/events" {
+				zr, _ := zlib.NewReader(r.Body)
+				b, _ := io.ReadAll(zr)
+				// Skip metadata and transaction keys, only count span.
+				events = strings.Count(string(b), "\n") - 2
+			}
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer s.Close()
+		cfg := agentconfig.MustNewConfigFrom(map[string]interface{}{
+			"instrumentation": map[string]interface{}{
+				"enabled": true,
+				"hosts":   []string{s.URL},
+				"tls": map[string]interface{}{
+					"skipverify": true,
+				},
+				"samplingrate": fmt.Sprintf("%f", rate),
+			},
+		})
+		i, err := newInstrumentation(cfg)
+		require.NoError(t, err)
+		tracer := i.Tracer()
+		tr := tracer.StartTransaction("name", "type")
+		tr.StartSpan("span", "type", nil).End()
+		tr.End()
+		tracer.Flush(nil)
+		assert.Equal(t, int(rate), events)
+	}
+	t.Run("100% sampling", func(t *testing.T) {
+		runSampled(1.0)
+	})
+	t.Run("0% sampling", func(t *testing.T) {
+		runSampled(0.0)
+	})
+}
+
 func TestProcessMemoryLimit(t *testing.T) {
 	l := logp.NewLogger("test")
 	const gb = 1 << 30