Add retries to BBE DNS probe

Copy the BBE DNS probe code to modify it in order to add retry logic. We can try to upstream this, that's why this is not adding any metrics yet. This will retry 3 times if there's enough time (at least 15 seconds). No DNS check should actually take more than 5 seconds, so when we see a long timeout, inject the retry logic in there. This experimental DNS prober is behind a feature flag called "experimental-dns-prober". The existing behavior is the default. Signed-off-by: Marcelo E. Magallon <[email protected]>
grafana · Jul 29, 2024 · 0d818c1 · 0d818c1
1 parent b2a24c4
commit 0d818c1
Show file tree

Hide file tree

Showing 13 changed files with 1,214 additions and 27 deletions.
diff --git a/go.mod b/go.mod
@@ -28,6 +28,7 @@ require (
 
 require (
 	github.com/KimMachineGun/automemlimit v0.6.1
+	github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9
 	github.com/felixge/httpsnoop v1.0.4
 	github.com/go-kit/log v0.2.1
 	github.com/go-ping/ping v1.1.0
@@ -38,11 +39,11 @@ require (
 	github.com/quasilyte/go-ruleguard/dsl v0.3.22
 	github.com/spf13/afero v1.11.0
 	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56
+	gopkg.in/yaml.v3 v3.0.1
 	kernel.org/pub/linux/libs/security/libcap/cap v1.2.70
 )
 
 require (
-	github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9 // indirect
 	github.com/andybalholm/brotli v1.1.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/buger/goterm v1.0.4 // indirect
@@ -73,7 +74,6 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240528184218-531527333157 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 	kernel.org/pub/linux/libs/security/libcap/psx v1.2.70 // indirect
 )
 

diff --git a/internal/adhoc/adhoc.go b/internal/adhoc/adhoc.go
@@ -139,7 +139,7 @@ func NewHandler(opts HandlerOpts) (*Handler, error) {
 		tenantCh:                     opts.TenantCh,
 		runnerFactory:                opts.runnerFactory,
 		grpcAdhocChecksClientFactory: opts.grpcAdhocChecksClientFactory,
-		proberFactory:                prober.NewProberFactory(opts.K6Runner, 0),
+		proberFactory:                prober.NewProberFactory(opts.K6Runner, 0, opts.Features),
 		api: apiInfo{
 			conn: opts.Conn,
 		},

diff --git a/internal/checks/checks.go b/internal/checks/checks.go
@@ -888,7 +888,9 @@ func (c *Updater) addAndStartScraperWithLock(ctx context.Context, check model.Ch
 	)
 
 	scraper, err := c.scraperFactory(
-		ctx, check, c.publisher, *c.probe, c.logger,
+		ctx, check, c.publisher, *c.probe,
+		c.features,
+		c.logger,
 		metrics,
 		c.k6Runner,
 		c.tenantLimits, c.telemeter,

diff --git a/internal/checks/checks_test.go b/internal/checks/checks_test.go
@@ -388,9 +388,13 @@ func (l testLabelsLimiter) LogLabels(ctx context.Context, tenantID model.GlobalI
 	return l.logLabelsLimit, nil
 }
 
-func testScraperFactory(ctx context.Context, check model.Check, publisher pusher.Publisher,
-	_ sm.Probe, logger zerolog.Logger, metrics scraper.Metrics,
-	k6Runner k6runner.Runner, labelsLimiter scraper.LabelsLimiter, telemeter *telemetry.Telemeter,
+func testScraperFactory(ctx context.Context, check model.Check, publisher pusher.Publisher, _ sm.Probe,
+	_ feature.Collection,
+	logger zerolog.Logger,
+	metrics scraper.Metrics,
+	k6Runner k6runner.Runner,
+	labelsLimiter scraper.LabelsLimiter,
+	telemeter *telemetry.Telemeter,
 ) (*scraper.Scraper, error) {
 	return scraper.NewWithOpts(
 		ctx,
@@ -406,6 +410,8 @@ func testScraperFactory(ctx context.Context, check model.Check, publisher pusher
 	)
 }
 
+var _ scraper.Factory = testScraperFactory
+
 type channelPublisher chan pusher.Payload
 
 func (c channelPublisher) Publish(payload pusher.Payload) {

diff --git a/internal/feature/feature.go b/internal/feature/feature.go
@@ -10,8 +10,9 @@ import (
 
 // TODO: this doesn't seem like the right place for this
 const (
-	Traceroute = "traceroute"
-	K6         = "k6"
+	Traceroute            = "traceroute"
+	K6                    = "k6"
+	ExperimentalDnsProber = "experimental-dns-prober"
 )
 
 // ErrInvalidCollection is returned when you try to set a flag in an

diff --git a/internal/prober/dns/dns.go b/internal/prober/dns/dns.go
@@ -6,18 +6,19 @@ import (
 	"strings"
 	"time"
 
+	"github.com/grafana/synthetic-monitoring-agent/internal/prober/dns/internal/bbe/config"
+	bbeprober "github.com/grafana/synthetic-monitoring-agent/internal/prober/dns/internal/bbe/prober"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/logger"
 	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
-	"github.com/prometheus/blackbox_exporter/config"
-	bbeprober "github.com/prometheus/blackbox_exporter/prober"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
 var errUnsupportedCheck = errors.New("unsupported check")
 
 type Prober struct {
-	target string
-	config config.Module
+	target       string
+	config       config.Module
+	experimental bool
 }
 
 func NewProber(check sm.Check) (Prober, error) {
@@ -34,16 +35,50 @@ func NewProber(check sm.Check) (Prober, error) {
 	}, nil
 }
 
+func NewExperimentalProber(check sm.Check) (Prober, error) {
+	p, err := NewProber(check)
+	if err != nil {
+		return p, err
+	}
+
+	p.experimental = true
+
+	return p, nil
+}
+
 func (p Prober) Name() string {
 	return "dns"
 }
 
 func (p Prober) Probe(ctx context.Context, target string, registry *prometheus.Registry, logger logger.Logger) bool {
+	cfg := p.config
+
+	if p.experimental {
+		const (
+			cutoff  = 15 * time.Second
+			retries = 3
+		)
+
+		if deadline, found := ctx.Deadline(); found {
+			budget := time.Until(deadline)
+			if budget >= cutoff {
+				cfg.DNS.Retries = retries
+				// Split 99% of the budget between three retries. For a
+				// budget of 15s, this allows for 150 ms per retry for
+				// other operations.
+				cfg.DNS.RetryTimeout = budget * 99 / (retries * 100)
+			}
+		}
+
+		_ = logger.Log("msg", "probing DNS", "target", target, "retries", cfg.DNS.Retries, "retry_timeout", cfg.DNS.RetryTimeout)
+	}
+
 	// The target of the BBE DNS check is the _DNS server_, while
 	// the target of the SM DNS check is the _query_, so we need
 	// pass the server as the target parameter, and ignore the
 	// _target_ paramater that is passed to this function.
-	return bbeprober.ProbeDNS(ctx, p.target, p.config, registry, logger)
+
+	return bbeprober.ProbeDNS(ctx, p.target, cfg, registry, logger)
 }
 
 func settingsToModule(settings *sm.DnsSettings, target string) config.Module {

diff --git a/internal/prober/dns/dns_test.go b/internal/prober/dns/dns_test.go
@@ -1,10 +1,22 @@
 package dns
 
 import (
+	"bytes"
+	"context"
+	"net"
+	"os"
+	"slices"
+	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/go-kit/log"
+	"github.com/grafana/synthetic-monitoring-agent/internal/prober/dns/internal/bbe/config"
 	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
-	"github.com/prometheus/blackbox_exporter/config"
+	"github.com/miekg/dns"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/common/expfmt"
 	"github.com/stretchr/testify/require"
 )
 
@@ -157,3 +169,99 @@ func TestSettingsToModule(t *testing.T) {
 		})
 	}
 }
+
+func TestProberRetries(t *testing.T) {
+	if !slices.Contains(strings.Split(os.Getenv("SM_TEST_RUN"), ","), "TestProberRetries") {
+		t.Skip("Skipping long test TestProberRetries")
+	}
+
+	mux := dns.NewServeMux()
+	var counter atomic.Int32
+	mux.Handle(".", dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) {
+		answer := dns.Msg{
+			MsgHdr: dns.MsgHdr{
+				Id:       r.Id,
+				Response: true,
+				Rcode:    dns.RcodeSuccess,
+			},
+			Question: r.Question,
+			Answer: []dns.RR{
+				&dns.A{
+					Hdr: dns.RR_Header{
+						Name:   r.Question[0].Name,
+						Rrtype: dns.TypeA,
+						Class:  dns.ClassINET,
+						Ttl:    60,
+					},
+					A: net.ParseIP("1.2.3.4"),
+				},
+			},
+		}
+
+		counter.Add(1)
+		t.Logf("Received request %d: %v", counter.Load(), r)
+		delay := 10 * time.Second
+		if counter.Load()%3 == 0 {
+			delay = 0
+		}
+		t.Log("Answer...")
+		time.Sleep(delay)
+		_ = w.WriteMsg(&answer)
+	}))
+	addr, err := net.ResolveUDPAddr("udp4", "127.0.0.1:0")
+	require.NoError(t, err)
+	l, err := net.ListenUDP("udp4", addr)
+	require.NoError(t, err)
+	t.Log(l.LocalAddr().String())
+	server := &dns.Server{Addr: ":0", PacketConn: l, Net: "udp", Handler: mux}
+	go func() {
+		err := server.ActivateAndServe()
+		if err != nil {
+			panic(err)
+		}
+	}()
+
+	p, err := NewExperimentalProber(sm.Check{
+		Target:  "www.grafana.com",
+		Timeout: 20000,
+		Settings: sm.CheckSettings{
+			Dns: &sm.DnsSettings{
+				Server:     l.LocalAddr().String(),
+				RecordType: sm.DnsRecordType_A,
+				Protocol:   sm.DnsProtocol_UDP,
+				IpVersion:  sm.IpVersion_V4,
+			},
+		},
+	})
+
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), p.config.Timeout)
+	t.Cleanup(cancel)
+
+	defer func() {
+		err := server.ShutdownContext(ctx)
+		if err != nil {
+			panic(err)
+		}
+	}()
+
+	registry := prometheus.NewPedanticRegistry()
+
+	var buf bytes.Buffer
+	logger := log.NewLogfmtLogger(&buf)
+
+	t0 := time.Now()
+	success := p.Probe(ctx, p.target, registry, logger)
+	t.Log(success, time.Since(t0))
+	require.True(t, success)
+
+	mfs, err := registry.Gather()
+	require.NoError(t, err)
+	enc := expfmt.NewEncoder(&buf, expfmt.NewFormat(expfmt.TypeTextPlain))
+	for _, mf := range mfs {
+		require.NoError(t, enc.Encode(mf))
+	}
+
+	t.Log(buf.String())
+}