fix latency

dmachard · Jan 5, 2024 · 05e9178 · 05e9178
1 parent 8cd4d0f
commit 05e9178
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 2 deletions.
diff --git a/dnsutils/message.go b/dnsutils/message.go
@@ -25,7 +25,9 @@ import (
 
 var (
 	DNSQuery                  = "QUERY"
+	DNSQueryQuiet             = "Q"
 	DNSReply                  = "REPLY"
+	DNSReplyQuiet             = "R"
 	PdnsDirectives            = regexp.MustCompile(`^powerdns-*`)
 	GeoIPDirectives           = regexp.MustCompile(`^geoip-*`)
 	SuspiciousDirectives      = regexp.MustCompile(`^suspicious-*`)

diff --git a/transformers/latency.go b/transformers/latency.go
@@ -149,7 +149,7 @@ func (s *LatencyProcessor) MeasureLatency(dm *dnsutils.DNSMessage) {
 		hashfnv := fnv.New64a()
 		hashfnv.Write([]byte(strings.Join(hashData, "+")))
 
-		if dm.DNS.Type == dnsutils.DNSQuery {
+		if dm.DNS.Type == dnsutils.DNSQuery || dm.DNS.Type == dnsutils.DNSQueryQuiet {
 			s.hashQueries.Set(hashfnv.Sum64(), dm.DNSTap.Timestamp)
 		} else {
 			key := hashfnv.Sum64()
@@ -174,7 +174,7 @@ func (s *LatencyProcessor) DetectEvictedTimeout(dm *dnsutils.DNSMessage) {
 		hashfnv.Write([]byte(strings.Join(hashData, "+")))
 		key := hashfnv.Sum64()
 
-		if dm.DNS.Type == dnsutils.DNSQuery {
+		if dm.DNS.Type == dnsutils.DNSQuery || dm.DNS.Type == dnsutils.DNSQueryQuiet {
 			s.mapQueries.Set(key, *dm)
 		} else if s.mapQueries.Exists(key) {
 			s.mapQueries.Delete(key)

diff --git a/transformers/latency_test.go b/transformers/latency_test.go
@@ -5,8 +5,113 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/dmachard/go-dnscollector/dnsutils"
+	"github.com/dmachard/go-dnscollector/pkgconfig"
+	"github.com/dmachard/go-logger"
 )
 
+func TestLatency_MeasureLatency(t *testing.T) {
+	// enable feature
+	config := pkgconfig.GetFakeConfigTransformers()
+	log := logger.New(false)
+	outChannels := []chan dnsutils.DNSMessage{}
+
+	// init transformer
+	latencyProcessor := NewLatencySubprocessor(config, logger.New(true), "test", 0, outChannels, log.Info, log.Error)
+
+	testcases := []struct {
+		name string
+		cq   string
+		cr   string
+	}{
+		{
+			name: "standard_mode",
+			cq:   dnsutils.DNSQuery,
+			cr:   dnsutils.DNSReply,
+		},
+		{
+			name: "quiet_mode",
+			cq:   dnsutils.DNSQueryQuiet,
+			cr:   dnsutils.DNSReplyQuiet,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Register Query
+			CQ := dnsutils.GetFakeDNSMessage()
+			CQ.DNS.Type = tc.cq
+			CQ.DNSTap.Timestamp = 1704486841216166066
+
+			// Measure latency
+			latencyProcessor.MeasureLatency(&CQ)
+
+			// Register Query
+			CR := dnsutils.GetFakeDNSMessage()
+			CR.DNS.Type = tc.cr
+			CR.DNSTap.Timestamp = 1704486841227961611
+
+			// Measure latency
+			latencyProcessor.MeasureLatency(&CR)
+
+			if CR.DNSTap.Latency == 0.0 {
+				t.Errorf("incorrect latency, got 0.0")
+			}
+		})
+	}
+}
+
+func TestLatency_DetectEvictedTimeout(t *testing.T) {
+	// enable feature
+	config := pkgconfig.GetFakeConfigTransformers()
+	config.Latency.Enable = true
+	config.Latency.QueriesTimeout = 1
+
+	log := logger.New(false)
+	outChannels := []chan dnsutils.DNSMessage{}
+	outChannels = append(outChannels, make(chan dnsutils.DNSMessage, 1))
+
+	// init transformer
+	latencyProcessor := NewLatencySubprocessor(config, logger.New(true), "test", 0, outChannels, log.Info, log.Error)
+
+	testcases := []struct {
+		name string
+		cq   string
+		cr   string
+	}{
+		{
+			name: "standard_mode",
+			cq:   dnsutils.DNSQuery,
+			cr:   dnsutils.DNSReply,
+		},
+		{
+			name: "quiet_mode",
+			cq:   dnsutils.DNSQueryQuiet,
+			cr:   dnsutils.DNSReplyQuiet,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Register Query
+			CQ := dnsutils.GetFakeDNSMessage()
+			CQ.DNS.Type = tc.cq
+			CQ.DNSTap.Timestamp = 1704486841216166066
+
+			// Measure latency
+			latencyProcessor.DetectEvictedTimeout(&CQ)
+
+			time.Sleep(2 * time.Second)
+
+			dmTimeout := <-outChannels[0]
+			if dmTimeout.DNS.Rcode != "TIMEOUT" {
+				t.Errorf("incorrect rcode, expected=TIMEOUT, got=%s", dmTimeout.DNS.Rcode)
+			}
+		})
+	}
+}
+
 func Test_HashQueries(t *testing.T) {
 	// init map
 	mapttl := NewHashQueries(2 * time.Second)