Skip to content

Commit dbdd056

Browse files
authored
Support sampling rate in APM configuration (#4037)
* Support sampling rate in APM configuration
1 parent 100e434 commit dbdd056

File tree

2 files changed

+96
-4
lines changed

2 files changed

+96
-4
lines changed

internal/pkg/server/agent.go

+7
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"errors"
1010
"fmt"
1111
"io"
12+
"strconv"
1213
"strings"
1314
"sync"
1415
"time"
@@ -556,6 +557,12 @@ func apmConfigToInstrumentation(src *proto.APMConfig) (config.Instrumentation, e
556557
Hosts: apmest.GetHosts(),
557558
GlobalLabels: apmest.GetGlobalLabels(),
558559
}
560+
561+
if apmest.SamplingRate != nil {
562+
// set the sampling rate in config
563+
cfg.TransactionSampleRate = strconv.FormatFloat(float64(*apmest.SamplingRate), 'f', -1, 32)
564+
}
565+
559566
return cfg, nil
560567
}
561568
return config.Instrumentation{}, fmt.Errorf("unable to transform APMConfig to instrumentation")

internal/pkg/server/agent_test.go

+89-4
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,78 @@ func Test_Agent_configFromUnits(t *testing.T) {
239239
assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
240240
require.Len(t, cfg.Output.Elasticsearch.Hosts, 2)
241241
})
242+
t.Run("Minimal APM config is specified", func(t *testing.T) {
243+
outStruct, err := structpb.NewStruct(map[string]interface{}{
244+
"service_token": "test-token",
245+
})
246+
require.NoError(t, err)
247+
mockOutClient := &mockClientUnit{}
248+
mockOutClient.On("Expected").Return(
249+
client.Expected{
250+
State: client.UnitStateHealthy,
251+
LogLevel: client.UnitLogLevelInfo,
252+
Config: &proto.UnitExpectedConfig{Source: outStruct},
253+
})
254+
255+
inStruct, err := structpb.NewStruct(map[string]interface{}{
256+
"type": "fleet-server",
257+
"server": map[string]interface{}{
258+
"host": "0.0.0.0",
259+
},
260+
"policy": map[string]interface{}{
261+
"id": "test-policy",
262+
},
263+
})
264+
require.NoError(t, err)
265+
266+
mockInClient := &mockClientUnit{}
267+
mockInClient.On("Expected").Return(
268+
client.Expected{
269+
State: client.UnitStateHealthy,
270+
LogLevel: client.UnitLogLevelInfo,
271+
Config: &proto.UnitExpectedConfig{Source: inStruct},
272+
APMConfig: &proto.APMConfig{
273+
Elastic: &proto.ElasticAPM{
274+
Environment: "test",
275+
SecretToken: "secretToken",
276+
Hosts: []string{"testhost:8080"},
277+
},
278+
},
279+
})
280+
281+
cliCfg, err := ucfg.NewFrom(map[string]interface{}{
282+
"inputs": []interface{}{
283+
map[string]interface{}{
284+
"policy": map[string]interface{}{
285+
"id": "test-policy",
286+
},
287+
},
288+
},
289+
})
290+
require.NoError(t, err)
291+
a := &Agent{
292+
cliCfg: cliCfg,
293+
agent: mockAgent,
294+
inputUnit: mockInClient,
295+
outputUnit: mockOutClient,
296+
}
297+
298+
cfg, err := a.configFromUnits(context.Background())
299+
require.NoError(t, err)
300+
require.Len(t, cfg.Inputs, 1)
301+
assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
302+
assert.Equal(t, "test-policy", cfg.Inputs[0].Policy.ID)
303+
assert.Equal(t, "0.0.0.0", cfg.Inputs[0].Server.Host)
304+
assert.True(t, cfg.Inputs[0].Server.Instrumentation.Enabled)
305+
assert.False(t, cfg.Inputs[0].Server.Instrumentation.TLS.SkipVerify)
306+
assert.Empty(t, cfg.Inputs[0].Server.Instrumentation.TLS.ServerCA)
307+
assert.Equal(t, "test", cfg.Inputs[0].Server.Instrumentation.Environment)
308+
assert.Empty(t, cfg.Inputs[0].Server.Instrumentation.APIKey)
309+
assert.Equal(t, "secretToken", cfg.Inputs[0].Server.Instrumentation.SecretToken)
310+
assert.Equal(t, []string{"testhost:8080"}, cfg.Inputs[0].Server.Instrumentation.Hosts)
311+
assert.Empty(t, cfg.Inputs[0].Server.Instrumentation.GlobalLabels)
312+
assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
313+
})
242314
t.Run("APM config is specified", func(t *testing.T) {
243315
outStruct, err := structpb.NewStruct(map[string]interface{}{
244316
"service_token": "test-token",
@@ -262,6 +334,8 @@ func Test_Agent_configFromUnits(t *testing.T) {
262334
},
263335
})
264336
require.NoError(t, err)
337+
338+
samplingRate := float32(0.5)
265339
mockInClient := &mockClientUnit{}
266340
mockInClient.On("Expected").Return(
267341
client.Expected{
@@ -279,6 +353,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
279353
SecretToken: "secretToken",
280354
Hosts: []string{"testhost:8080"},
281355
GlobalLabels: "test",
356+
SamplingRate: &samplingRate,
282357
},
283358
},
284359
})
@@ -315,6 +390,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
315390
assert.Equal(t, []string{"testhost:8080"}, cfg.Inputs[0].Server.Instrumentation.Hosts)
316391
assert.Equal(t, "test", cfg.Inputs[0].Server.Instrumentation.GlobalLabels)
317392
assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
393+
assert.Equal(t, "0.5", cfg.Inputs[0].Server.Instrumentation.TransactionSampleRate)
318394
})
319395
t.Run("APM config no tls", func(t *testing.T) {
320396
outStruct, err := structpb.NewStruct(map[string]interface{}{
@@ -336,6 +412,8 @@ func Test_Agent_configFromUnits(t *testing.T) {
336412
},
337413
})
338414
require.NoError(t, err)
415+
416+
samplingRate := float32(0.01)
339417
mockInClient := &mockClientUnit{}
340418
mockInClient.On("Expected").Return(
341419
client.Expected{
@@ -349,6 +427,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
349427
SecretToken: "secretToken",
350428
Hosts: []string{"testhost:8080"},
351429
GlobalLabels: "test",
430+
SamplingRate: &samplingRate,
352431
},
353432
},
354433
})
@@ -374,6 +453,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
374453
assert.Equal(t, []string{"testhost:8080"}, cfg.Inputs[0].Server.Instrumentation.Hosts)
375454
assert.Equal(t, "test", cfg.Inputs[0].Server.Instrumentation.GlobalLabels)
376455
assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
456+
assert.Equal(t, "0.01", cfg.Inputs[0].Server.Instrumentation.TransactionSampleRate)
377457
})
378458
t.Run("APM config and instrumentation is specified", func(t *testing.T) {
379459
outStruct, err := structpb.NewStruct(map[string]interface{}{
@@ -398,14 +478,17 @@ func Test_Agent_configFromUnits(t *testing.T) {
398478
"skip_verify": true,
399479
"server_certificate": "/path/to/cert.crt",
400480
},
401-
"environment": "replace",
402-
"api_key": "replace",
403-
"secret_token": "replace",
404-
"hosts": []interface{}{"replace"},
481+
"environment": "replace",
482+
"api_key": "replace",
483+
"secret_token": "replace",
484+
"hosts": []interface{}{"replace"},
485+
"transaction_sample_rate": "0.75",
405486
},
406487
},
407488
})
408489
require.NoError(t, err)
490+
491+
samplingRate := float32(0.01)
409492
mockInClient := &mockClientUnit{}
410493
mockInClient.On("Expected").Return(
411494
client.Expected{
@@ -423,6 +506,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
423506
SecretToken: "secretToken",
424507
Hosts: []string{"testhost:8080"},
425508
GlobalLabels: "test",
509+
SamplingRate: &samplingRate,
426510
},
427511
},
428512
})
@@ -449,6 +533,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
449533
assert.Equal(t, []string{"testhost:8080"}, cfg.Inputs[0].Server.Instrumentation.Hosts)
450534
assert.Equal(t, "test", cfg.Inputs[0].Server.Instrumentation.GlobalLabels)
451535
assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
536+
assert.Equal(t, "0.01", cfg.Inputs[0].Server.Instrumentation.TransactionSampleRate)
452537
})
453538
t.Run("APM config error", func(t *testing.T) {
454539
outStruct, err := structpb.NewStruct(map[string]interface{}{

0 commit comments

Comments
 (0)