Add retries to BBE DNS probe

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
 )
 

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,
 		},

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,

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) {

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

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 {

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())
+}