From 8167de5805395eb1de1b052bf0ecf7b39f0eb660 Mon Sep 17 00:00:00 2001
From: Simone Basso <bassosimone@gmail.com>
Date: Thu, 8 Sep 2022 11:05:38 +0200
Subject: [PATCH] fix(measurexlite): emit resolve_start and resolve_done (#944)

Part of https://github.com/ooni/probe/issues/2238
---
 internal/measurexlite/dns.go      | 26 ++++++++++++++++++++++++++
 internal/measurexlite/dns_test.go | 27 ++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/internal/measurexlite/dns.go b/internal/measurexlite/dns.go
index 0e9e241..5d37485 100644
--- a/internal/measurexlite/dns.go
+++ b/internal/measurexlite/dns.go
@@ -49,18 +49,44 @@ func (r *resolverTrace) CloseIdleConnections() {
 	r.r.CloseIdleConnections()
 }
 
+// emits the resolve_start event
+func (r *resolverTrace) emitResolveStart() {
+	select {
+	case r.tx.networkEvent <- NewAnnotationArchivalNetworkEvent(
+		r.tx.Index, r.tx.TimeSince(r.tx.ZeroTime), "resolve_start",
+	):
+	default: // buffer is full
+	}
+}
+
+// emits the resolve_done event
+func (r *resolverTrace) emiteResolveDone() {
+	select {
+	case r.tx.networkEvent <- NewAnnotationArchivalNetworkEvent(
+		r.tx.Index, r.tx.TimeSince(r.tx.ZeroTime), "resolve_done",
+	):
+	default: // buffer is full
+	}
+}
+
 // LookupHost implements model.Resolver.LookupHost
 func (r *resolverTrace) LookupHost(ctx context.Context, hostname string) ([]string, error) {
+	defer r.emiteResolveDone()
+	r.emitResolveStart()
 	return r.r.LookupHost(netxlite.ContextWithTrace(ctx, r.tx), hostname)
 }
 
 // LookupHTTPS implements model.Resolver.LookupHTTPS
 func (r *resolverTrace) LookupHTTPS(ctx context.Context, domain string) (*model.HTTPSSvc, error) {
+	defer r.emiteResolveDone()
+	r.emitResolveStart()
 	return r.r.LookupHTTPS(netxlite.ContextWithTrace(ctx, r.tx), domain)
 }
 
 // LookupNS implements model.Resolver.LookupNS
 func (r *resolverTrace) LookupNS(ctx context.Context, domain string) ([]*net.NS, error) {
+	defer r.emiteResolveDone()
+	r.emitResolveStart()
 	return r.r.LookupNS(netxlite.ContextWithTrace(ctx, r.tx), domain)
 }
 
diff --git a/internal/measurexlite/dns_test.go b/internal/measurexlite/dns_test.go
index f6e5e22..998b3f4 100644
--- a/internal/measurexlite/dns_test.go
+++ b/internal/measurexlite/dns_test.go
@@ -177,7 +177,7 @@ func TestNewResolver(t *testing.T) {
 		t.Run("DNSLookup events", func(t *testing.T) {
 			events := trace.DNSLookupsFromRoundTrip()
 			if len(events) != 2 {
-				t.Fatal("unexpected DNS events")
+				t.Fatal("unexpected DNS events length")
 			}
 			for _, ev := range events {
 				if ev.ResolverAddress != "dns.google" {
@@ -200,6 +200,23 @@ func TestNewResolver(t *testing.T) {
 				}
 			}
 		})
+
+		t.Run("Network events", func(t *testing.T) {
+			events := trace.NetworkEvents()
+			if len(events) != 2 {
+				t.Fatal("unexpected network events length")
+			}
+			foundNames := map[string]int{}
+			for _, ev := range events {
+				foundNames[ev.Operation]++
+			}
+			if foundNames["resolve_start"] != 1 {
+				t.Fatal("missing resolve_start")
+			}
+			if foundNames["resolve_done"] != 1 {
+				t.Fatal("missing resolve_done")
+			}
+		})
 	})
 
 	t.Run("LookupHost discards events when buffers are full", func(t *testing.T) {
@@ -207,6 +224,7 @@ func TestNewResolver(t *testing.T) {
 		td := testingx.NewTimeDeterministic(zeroTime)
 		trace := NewTrace(0, zeroTime)
 		trace.dnsLookup = make(chan *model.ArchivalDNSLookupResult) // no buffer
+		trace.networkEvent = make(chan *model.ArchivalNetworkEvent) // ditto
 		trace.TimeNowFn = td.Now
 		txp := &mocks.DNSTransport{
 			MockRoundTrip: func(ctx context.Context, query model.DNSQuery) (model.DNSResponse, error) {
@@ -262,6 +280,13 @@ func TestNewResolver(t *testing.T) {
 				t.Fatal("expected to see no DNSLookup events")
 			}
 		})
+
+		t.Run("Network events", func(t *testing.T) {
+			events := trace.NetworkEvents()
+			if len(events) != 0 {
+				t.Fatal("unexpected to see no network events")
+			}
+		})
 	})
 }