From 8daf162e782f9eb4e2c259fcd7cb34cc76457b1f Mon Sep 17 00:00:00 2001
From: Fedor Setrakov <f.setrakov@adguard.com>
Date: Wed, 30 Sep 2026 15:06:35 +0300
Subject: [PATCH] AGH-15 exchange-ctx (18)

---
 dnsproxytest/interface.go                 |  6 +-
 fastip/fastest.go                         |  8 +-
 fastip/fastest_test.go                    | 21 +++--
 fastip/ping.go                            | 48 +++++++----
 fastip/ping_internal_test.go              | 17 ++--
 proxy/cache_internal_test.go              |  9 ++-
 proxy/dns64.go                            | 10 ++-
 proxy/dns64_internal_test.go              | 15 ++--
 proxy/exchange.go                         | 19 +++--
 proxy/exchange_internal_test.go           | 31 +++++---
 proxy/optimisticresolver.go               | 19 +++--
 proxy/optimisticresolver_internal_test.go | 39 ++++++---
 proxy/pending_test.go                     |  2 +-
 proxy/proxy.go                            | 26 +++---
 proxy/proxy_internal_test.go              | 57 ++++++++-----
 proxy/proxy_test.go                       | 17 ++--
 proxy/proxycache.go                       |  8 +-
 proxy/servertcp_test.go                   |  9 ++-
 proxy/stats.go                            |  5 +-
 proxy/stats_test.go                       |  5 +-
 upstream/dnscrypt.go                      | 15 +++-
 upstream/dnscrypt_internal_test.go        | 16 ++--
 upstream/doh.go                           | 97 ++++++++++++++---------
 upstream/doh_internal_test.go             | 21 ++---
 upstream/doq.go                           | 56 +++++++------
 upstream/doq_internal_test.go             | 31 ++++----
 upstream/dot.go                           | 61 ++++++++++----
 upstream/dot_internal_test.go             | 29 ++++---
 upstream/parallel.go                      | 28 ++++---
 upstream/parallel_internal_test.go        | 12 ++-
 upstream/plain.go                         | 47 ++++++++---
 upstream/plain_internal_test.go           |  8 +-
 upstream/resolver.go                      | 26 +++---
 upstream/resolver_test.go                 |  4 +-
 upstream/upstream.go                      | 10 +--
 upstream/upstream_internal_test.go        | 28 ++++---
 36 files changed, 551 insertions(+), 309 deletions(-)

diff --git a/dnsproxytest/interface.go b/dnsproxytest/interface.go
index 5f700abb..6210fb58 100644
--- a/dnsproxytest/interface.go
+++ b/dnsproxytest/interface.go
@@ -14,7 +14,7 @@ import (
 // Upstream is a mock [upstream.Upstream] implementation for tests.
 type Upstream struct {
 	OnAddress  func() (addr string)
-	OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
+	OnExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
 	OnClose    func() (err error)
 }
 
@@ -27,8 +27,8 @@ func (u *Upstream) Address() (addr string) {
 }
 
 // Exchange implements the [upstream.Upstream] interface for *Upstream.
-func (u *Upstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
-	return u.OnExchange(req)
+func (u *Upstream) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
+	return u.OnExchange(ctx, req)
 }
 
 // Close implements the [upstream.Upstream] interface for *Upstream.
diff --git a/fastip/fastest.go b/fastip/fastest.go
index bc35d477..b65b5a48 100644
--- a/fastip/fastest.go
+++ b/fastip/fastest.go
@@ -4,6 +4,7 @@
 package fastip
 
 import (
+	"context"
 	"log/slog"
 	"net"
 	"net/netip"
@@ -92,10 +93,11 @@ func New(c *Config) (f *FastestAddr) {
 // the fastest IP address.  The fastest IP address is considered to be the first
 // one successfully dialed and other addresses are removed from the answer.
 func (f *FastestAddr) ExchangeFastest(
+	ctx context.Context,
 	req *dns.Msg,
 	ups []upstream.Upstream,
 ) (resp *dns.Msg, u upstream.Upstream, err error) {
-	replies, err := upstream.ExchangeAll(ups, req)
+	replies, err := upstream.ExchangeAll(ctx, ups, req)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -112,11 +114,11 @@ func (f *FastestAddr) ExchangeFastest(
 
 	ips := ipSet.Values()
 	host := strings.ToLower(req.Question[0].Name)
-	if pingRes := f.pingAll(host, ips); pingRes != nil {
+	if pingRes := f.pingAll(ctx, host, ips); pingRes != nil {
 		return f.prepareReply(pingRes, replies)
 	}
 
-	f.logger.Debug("no fastest ip found, using the first response", "host", host)
+	f.logger.DebugContext(ctx, "no fastest ip found, using the first response", "host", host)
 
 	return replies[0].Resp, replies[0].Upstream, nil
 }
diff --git a/fastip/fastest_test.go b/fastip/fastest_test.go
index 03c3b3da..dcd89cdd 100644
--- a/fastip/fastest_test.go
+++ b/fastip/fastest_test.go
@@ -1,6 +1,7 @@
 package fastip_test
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"strings"
@@ -8,6 +9,7 @@ import (
 
 	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/fastip"
+	proxytest "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/internal/nettest"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
@@ -30,16 +32,19 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
 		t.Parallel()
 
 		u := &dnsproxytest.Upstream{
-			OnAddress:  func() (addr string) { return "bad_upstream" },
-			OnExchange: func(_ *dns.Msg) (resp *dns.Msg, err error) { return nil, assert.AnError },
-			OnClose:    func() (err error) { return nil },
+			OnAddress: func() (addr string) { return "bad_upstream" },
+			OnExchange: func(_ context.Context, _ *dns.Msg) (resp *dns.Msg, err error) {
+				return nil, assert.AnError
+			},
+			OnClose: func() (err error) { return nil },
 		}
 		f := fastip.New(&fastip.Config{
 			Logger:          l,
 			PingWaitTimeout: fastip.DefaultPingWaitTimeout,
 		})
 
-		resp, up, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{u})
+		ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
+		resp, up, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{u})
 		require.Error(t, err)
 
 		assert.ErrorIs(t, err, assert.AnError)
@@ -67,7 +72,8 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
 		alive := newTestAUpstream(t, []*dns.A{newTestRec(t, aliveAddr)})
 		dead := newTestAUpstream(t, []*dns.A{newTestRec(t, netip.MustParseAddr("192.0.2.1"))})
 
-		rep, ups, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{dead, alive})
+		ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
+		rep, ups, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{dead, alive})
 		require.NoError(t, err)
 
 		assert.Equal(t, ups, alive)
@@ -96,7 +102,8 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
 			newTestRec(t, netip.MustParseAddr("127.0.0.3")),
 		})
 
-		resp, _, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{ups})
+		ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
+		resp, _, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{ups})
 		require.NoError(t, err)
 
 		require.NotNil(t, resp)
@@ -112,7 +119,7 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
 func newTestAUpstream(tb testing.TB, recs []*dns.A) (ups *dnsproxytest.Upstream) {
 	tb.Helper()
 
-	onExchange := func(m *dns.Msg) (resp *dns.Msg, err error) {
+	onExchange := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 		resp = &dns.Msg{}
 		resp.SetReply(m)
 
diff --git a/fastip/ping.go b/fastip/ping.go
index 041531d4..de980fd9 100644
--- a/fastip/ping.go
+++ b/fastip/ping.go
@@ -1,6 +1,7 @@
 package fastip
 
 import (
+	"context"
 	"net/netip"
 	"time"
 
@@ -29,6 +30,7 @@ type pingResult struct {
 // Returns scheduled flag which indicates that some goroutines have been
 // scheduled.
 func (f *FastestAddr) schedulePings(
+	ctx context.Context,
 	resCh chan *pingResult,
 	ips []netip.Addr,
 	host string,
@@ -38,7 +40,7 @@ func (f *FastestAddr) schedulePings(
 		if cached == nil {
 			scheduled = true
 			for _, port := range f.pingPorts {
-				go f.pingDoTCP(host, netip.AddrPortFrom(ip, uint16(port)), resCh)
+				go f.pingDoTCP(ctx, host, netip.AddrPortFrom(ip, uint16(port)), resCh)
 			}
 
 			continue
@@ -58,7 +60,7 @@ func (f *FastestAddr) schedulePings(
 
 // pingAll pings all ips concurrently and returns as soon as the fastest one is
 // found or the timeout is exceeded.
-func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
+func (f *FastestAddr) pingAll(ctx context.Context, host string, ips []netip.Addr) (pr *pingResult) {
 	ipN := len(ips)
 	switch ipN {
 	case 0:
@@ -71,22 +73,23 @@ func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
 	}
 
 	resCh := make(chan *pingResult, ipN*len(f.pingPorts))
-	pr, scheduled := f.schedulePings(resCh, ips, host)
+	pr, scheduled := f.schedulePings(ctx, resCh, ips, host)
 	if !scheduled {
 		if pr != nil {
-			f.logger.Debug(
+			f.logger.DebugContext(
+				ctx,
 				"pinging all returns cached response",
 				"host", host,
 				"addr", pr.addrPort,
 			)
 		} else {
-			f.logger.Debug("pinging all returns nothing", "host", host)
+			f.logger.DebugContext(ctx, "pinging all returns nothing", "host", host)
 		}
 
 		return pr
 	}
 
-	res := f.firstSuccessRes(resCh, host)
+	res := f.firstSuccessRes(ctx, resCh, host)
 	if res == nil {
 		// In case of timeout return cached or nil.
 		return pr
@@ -103,12 +106,17 @@ func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
 
 // firstSuccessRes waits and returns the first successful ping result or nil in
 // case of timeout.
-func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res *pingResult) {
+func (f *FastestAddr) firstSuccessRes(
+	ctx context.Context,
+	resCh chan *pingResult,
+	host string,
+) (res *pingResult) {
 	after := time.After(f.pingWaitTimeout)
 	for {
 		select {
 		case res = <-resCh:
-			f.logger.Debug(
+			f.logger.DebugContext(
+				ctx,
 				"pinging all got result",
 				"host", host,
 				"addr", res.addrPort,
@@ -121,7 +129,7 @@ func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res
 
 			return res
 		case <-after:
-			f.logger.Debug("pinging all timed out", "host", host)
+			f.logger.DebugContext(ctx, "pinging all timed out", "host", host)
 
 			return nil
 		}
@@ -129,18 +137,25 @@ func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res
 }
 
 // pingDoTCP sends the result of dialing the specified address into resCh.
-func (f *FastestAddr) pingDoTCP(host string, addrPort netip.AddrPort, resCh chan *pingResult) {
+func (f *FastestAddr) pingDoTCP(
+	ctx context.Context,
+	host string,
+	addrPort netip.AddrPort,
+	resCh chan *pingResult,
+) {
 	l := f.logger.With("host", host, "addr", addrPort)
-	l.Debug("open tcp connection")
+	l.DebugContext(ctx, "open tcp connection")
 
 	start := time.Now()
+
+	// TODO(f.setrakov): Respect context timeout.
 	conn, err := f.pinger.Dial(bootstrap.NetworkTCP, addrPort.String())
 	elapsed := time.Since(start)
 
 	success := err == nil
 	if success {
 		if cErr := conn.Close(); cErr != nil {
-			l.Debug("closing tcp connection", slogutil.KeyError, cErr)
+			l.DebugContext(ctx, "closing tcp connection", slogutil.KeyError, cErr)
 		}
 	}
 
@@ -154,10 +169,15 @@ func (f *FastestAddr) pingDoTCP(host string, addrPort netip.AddrPort, resCh chan
 
 	addr := addrPort.Addr().Unmap()
 	if success {
-		l.Debug("tcp ping success", "elapsed", elapsed)
+		l.DebugContext(ctx, "tcp ping success", "elapsed", elapsed)
 		f.cacheAddSuccessful(addr, latency)
 	} else {
-		l.Debug("tcp ping failed to connect", "elapsed", elapsed, slogutil.KeyError, err)
+		l.DebugContext(
+			ctx,
+			"tcp ping failed to connect",
+			"elapsed", elapsed,
+			slogutil.KeyError, err,
+		)
 		f.cacheAddFailure(addr)
 	}
 }
diff --git a/fastip/ping_internal_test.go b/fastip/ping_internal_test.go
index df69ac3a..c0658409 100644
--- a/fastip/ping_internal_test.go
+++ b/fastip/ping_internal_test.go
@@ -7,6 +7,7 @@ import (
 	"syscall"
 	"testing"
 
+	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
 	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/testutil"
@@ -30,7 +31,9 @@ func TestFastestAddr_PingAll_timeout(t *testing.T) {
 		}
 
 		ip := netutil.IPv4Localhost()
-		res := f.pingAll("", []netip.Addr{ip, ip})
+
+		ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
+		res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
 		require.Nil(t, res)
 
 		waitCh <- unit{}
@@ -52,7 +55,8 @@ func TestFastestAddr_PingAll_timeout(t *testing.T) {
 			return nil
 		}
 
-		res := f.pingAll("", []netip.Addr{ip1, ip2})
+		ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
+		res := f.pingAll(ctx, "", []netip.Addr{ip1, ip2})
 		require.NotNil(t, res)
 
 		assert.True(t, res.success)
@@ -83,7 +87,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
 		f := New(&Config{Logger: slogutil.NewDiscardLogger()})
 		f.cacheAddFailure(ip)
 
-		res := f.pingAll("", []netip.Addr{ip, ip})
+		ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
+		res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
 		require.Nil(t, res)
 	})
 
@@ -93,7 +98,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
 		f := New(&Config{Logger: slogutil.NewDiscardLogger()})
 		f.cacheAddSuccessful(ip, lat)
 
-		res := f.pingAll("", []netip.Addr{ip, ip})
+		ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
+		res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
 		require.NotNil(t, res)
 		assert.True(t, res.success)
 		assert.Equal(t, lat, res.latency)
@@ -125,7 +131,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
 			return nil
 		}
 
-		res := f.pingAll("", ips)
+		ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
+		res := f.pingAll(ctx, "", ips)
 		require.NotNil(t, res)
 
 		assert.True(t, res.success)
diff --git a/proxy/cache_internal_test.go b/proxy/cache_internal_test.go
index 6089ed1c..2494ddeb 100644
--- a/proxy/cache_internal_test.go
+++ b/proxy/cache_internal_test.go
@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"cmp"
+	"context"
 	"net"
 	"net/netip"
 	"strings"
@@ -26,9 +27,11 @@ const testUpsAddr = "https://upstream.address"
 // upstreamWithAddr is a [dnsproxytest.Upstream] that is only expected to be
 // used to get its address.
 var upstreamWithAddr = &testUpstream{
-	OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
-	OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
-	OnAddress:  func() (addr string) { return testUpsAddr },
+	OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
+		panic(testutil.UnexpectedCall(ctx, m))
+	},
+	OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
+	OnAddress: func() (addr string) { return testUpsAddr },
 }
 
 // newTestCache is a helper that returns new cache and fills its config with
diff --git a/proxy/dns64.go b/proxy/dns64.go
index a35b18e4..312f64bb 100644
--- a/proxy/dns64.go
+++ b/proxy/dns64.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/netip"
@@ -278,6 +279,7 @@ func (p *Proxy) synthRR(rr dns.RR, soaTTL uint32) (result dns.RR) {
 // nil, if the request was not performed.  If origResp is not nil, origReq must
 // also not be nil.
 func (p *Proxy) performDNS64(
+	ctx context.Context,
 	origReq *dns.Msg,
 	origResp *dns.Msg,
 	upstreams []upstream.Upstream,
@@ -292,17 +294,17 @@ func (p *Proxy) performDNS64(
 	}
 
 	host := origReq.Question[0].Name
-	p.logger.Debug("received an empty aaaa response, checking dns64", "host", host)
+	p.logger.DebugContext(ctx, "received an empty aaaa response, checking dns64", "host", host)
 
-	dns64Resp, u, err := p.exchangeUpstreams(dns64Req, upstreams)
+	dns64Resp, u, err := p.exchangeUpstreams(ctx, dns64Req, upstreams)
 	if err != nil {
-		p.logger.Error("dns64 request failed", slogutil.KeyError, err)
+		p.logger.ErrorContext(ctx, "dns64 request failed", slogutil.KeyError, err)
 
 		return nil
 	}
 
 	if dns64Resp != nil && p.synthDNS64(origReq, origResp, dns64Resp) {
-		p.logger.Debug("synthesized aaaa response", "host", host)
+		p.logger.DebugContext(ctx, "synthesized aaaa response", "host", host)
 
 		return u
 	}
diff --git a/proxy/dns64_internal_test.go b/proxy/dns64_internal_test.go
index 66084005..11b4b475 100644
--- a/proxy/dns64_internal_test.go
+++ b/proxy/dns64_internal_test.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"sync"
@@ -21,7 +22,7 @@ const ipv4OnlyFqdn = "ipv4.only."
 func TestProxy_HandleDNSRequest_dns64Race(t *testing.T) {
 	ans := newRR(t, ipv4OnlyFqdn, dns.TypeA, 3600, dnsproxytest.IPv4)
 	ups := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			resp = (&dns.Msg{}).SetReply(req)
 			if req.Question[0].Qtype == dns.TypeA {
 				resp.Answer = []dns.RR{dns.Copy(ans)}
@@ -33,9 +34,11 @@ func TestProxy_HandleDNSRequest_dns64Race(t *testing.T) {
 		OnClose:   func() (err error) { return nil },
 	}
 	localUps := &testUpstream{
-		OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
-		OnAddress:  func() (addr string) { return "fake.address" },
-		OnClose:    func() (err error) { return nil },
+		OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
+			panic(testutil.UnexpectedCall(ctx, m))
+		},
+		OnAddress: func() (addr string) { return "fake.address" },
+		OnClose:   func() (err error) { return nil },
 	}
 
 	dnsProxy := mustNew(t, &Config{
@@ -193,7 +196,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 	pt := testutil.NewPanicT(t)
 	newUps := func(answers answerMap) (u upstream.Upstream) {
 		return &testUpstream{
-			OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+			OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 				q := req.Question[0]
 				require.Contains(pt, answers, q.Qtype)
 
@@ -213,7 +216,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 
 	localRR := newRR(t, ptr64Domain, dns.TypePTR, 3600, domainPointed)
 	localUps := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			require.Equal(pt, req.Question[0].Name, ptr64Domain)
 			resp = (&dns.Msg{}).SetReply(req)
 			resp.Answer = []dns.RR{localRR}
diff --git a/proxy/exchange.go b/proxy/exchange.go
index 99693b11..36ee433e 100644
--- a/proxy/exchange.go
+++ b/proxy/exchange.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -15,16 +16,17 @@ import (
 // response, the upstream that successfully resolved the request, and the error
 // if any.
 func (p *Proxy) exchangeUpstreams(
+	ctx context.Context,
 	req *dns.Msg,
 	ups []upstream.Upstream,
 ) (resp *dns.Msg, u upstream.Upstream, err error) {
 	switch p.upstreamMode {
 	case UpstreamModeParallel:
-		return upstream.ExchangeParallel(ups, req)
+		return upstream.ExchangeParallel(ctx, ups, req)
 	case UpstreamModeFastestAddr:
 		switch req.Question[0].Qtype {
 		case dns.TypeA, dns.TypeAAAA:
-			return p.fastestAddr.ExchangeFastest(req, ups)
+			return p.fastestAddr.ExchangeFastest(ctx, req, ups)
 		default:
 			// Go on to the load-balancing mode.
 		}
@@ -34,7 +36,7 @@ func (p *Proxy) exchangeUpstreams(
 
 	if len(ups) == 1 {
 		u = ups[0]
-		resp, _, err = p.exchange(u, req)
+		resp, _, err = p.exchange(ctx, u, req)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -50,7 +52,7 @@ func (p *Proxy) exchangeUpstreams(
 		u = ups[i]
 
 		var elapsed time.Duration
-		resp, elapsed, err = p.exchange(u, req)
+		resp, elapsed, err = p.exchange(ctx, u, req)
 		if err == nil {
 			p.updateRTT(u.Address(), elapsed)
 
@@ -72,11 +74,12 @@ func (p *Proxy) exchangeUpstreams(
 // exchange returns the result of the DNS request exchange with the given
 // upstream and the elapsed time.  u and req must not be nil.
 func (p *Proxy) exchange(
+	ctx context.Context,
 	u upstream.Upstream,
 	req *dns.Msg,
 ) (resp *dns.Msg, dur time.Duration, err error) {
 	startTime := p.time.Now()
-	resp, err = u.Exchange(req)
+	resp, err = u.Exchange(ctx, req)
 
 	// Don't use [time.Since] because it uses [time.Now].
 	dur = p.time.Now().Sub(startTime)
@@ -84,7 +87,8 @@ func (p *Proxy) exchange(
 	addr := u.Address()
 	q := &req.Question[0]
 	if err != nil {
-		p.logger.Error(
+		p.logger.ErrorContext(
+			ctx,
 			"exchange failed",
 			"upstream", addr,
 			"question", q,
@@ -92,7 +96,8 @@ func (p *Proxy) exchange(
 			slogutil.KeyError, err,
 		)
 	} else {
-		p.logger.Debug(
+		p.logger.DebugContext(
+			ctx,
 			"exchange successfully finished",
 			"upstream", addr,
 			"question", q,
diff --git a/proxy/exchange_internal_test.go b/proxy/exchange_internal_test.go
index 75813fbc..de10ef77 100644
--- a/proxy/exchange_internal_test.go
+++ b/proxy/exchange_internal_test.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"math/rand/v2"
 	"net"
 	"net/netip"
@@ -25,7 +26,7 @@ func newUpstreamWithErrorRate(rate uint, name string) (u upstream.Upstream) {
 	var n uint
 
 	return &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			n++
 			if n%rate == 0 {
 				return nil, assert.AnError
@@ -52,10 +53,10 @@ type measuredUpstream struct {
 var _ upstream.Upstream = measuredUpstream{}
 
 // Exchange implements the [upstream.Upstream] interface for measuredUpstream.
-func (u measuredUpstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (u measuredUpstream) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	u.stats[u.Address()]++
 
-	return u.Upstream.Exchange(req)
+	return u.Upstream.Exchange(ctx, req)
 }
 
 func TestProxy_Exchange_loadBalance(t *testing.T) {
@@ -88,7 +89,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 	}
 
 	fastUps := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 100)
 
 			return (&dns.Msg{}).SetReply(req), nil
@@ -97,7 +98,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 	slowerUps := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 10)
 
 			return (&dns.Msg{}).SetReply(req), nil
@@ -106,7 +107,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 	slowestUps := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 2)
 
 			return (&dns.Msg{}).SetReply(req), nil
@@ -116,20 +117,24 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 	}
 
 	err1Ups := &testUpstream{
-		OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
-		OnAddress:  func() (addr string) { return "error1" },
-		OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
+		OnExchange: func(_ context.Context, _ *dns.Msg) (r *dns.Msg, err error) {
+			return nil, assert.AnError
+		},
+		OnAddress: func() (addr string) { return "error1" },
+		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 	err2Ups := &testUpstream{
-		OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
-		OnAddress:  func() (addr string) { return "error2" },
-		OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
+		OnExchange: func(_ context.Context, _ *dns.Msg) (r *dns.Msg, err error) {
+			return nil, assert.AnError
+		},
+		OnAddress: func() (addr string) { return "error2" },
+		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 
 	singleError := &sync.Once{}
 	// fastestUps responds with an error on the first request.
 	fastestUps := &testUpstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			singleError.Do(func() { err = assert.AnError })
 			currentNow = zeroTime.Add(testRTT / 200)
 
diff --git a/proxy/optimisticresolver.go b/proxy/optimisticresolver.go
index 011ade64..f7e587a1 100644
--- a/proxy/optimisticresolver.go
+++ b/proxy/optimisticresolver.go
@@ -15,7 +15,7 @@ type cachingResolver interface {
 	// resolved and the response may be cached.
 	//
 	// TODO(e.burkov):  Find out when ok can be false with nil err.
-	replyFromUpstream(dctx *DNSContext) (ok bool, err error)
+	replyFromUpstream(ctx context.Context, dctx *DNSContext) (ok bool, err error)
 
 	// cacheResp caches the response from dctx.
 	cacheResp(dctx *DNSContext)
@@ -46,10 +46,13 @@ func newOptimisticResolver(cr cachingResolver) (s *optimisticResolver) {
 // with the same key at the same period of time.  It runs in a separate
 // goroutine.  Do not pass the *DNSContext which is used elsewhere since it
 // isn't intended to be used concurrently.
-//
-// TODO(e.burkov):  Pass the context.
-func (s *optimisticResolver) resolveOnce(dctx *DNSContext, key []byte, l *slog.Logger) {
-	defer slogutil.RecoverAndLog(context.TODO(), l)
+func (s *optimisticResolver) resolveOnce(
+	ctx context.Context,
+	dctx *DNSContext,
+	key []byte,
+	l *slog.Logger,
+) {
+	defer slogutil.RecoverAndLog(ctx, l)
 
 	keyHexed := hex.EncodeToString(key)
 	if _, ok := s.reqs.LoadOrStore(keyHexed, unit{}); ok {
@@ -57,9 +60,11 @@ func (s *optimisticResolver) resolveOnce(dctx *DNSContext, key []byte, l *slog.L
 	}
 	defer s.reqs.Delete(keyHexed)
 
-	ok, err := s.cr.replyFromUpstream(dctx)
+	// NOTE: Discard the context deadline to prevent cancellation during
+	// refresh.
+	ok, err := s.cr.replyFromUpstream(context.WithoutCancel(ctx), dctx)
 	if err != nil {
-		l.Debug("resolving request for optimistic cache", slogutil.KeyError, err)
+		l.DebugContext(ctx, "resolving request for optimistic cache", slogutil.KeyError, err)
 	}
 
 	if ok {
diff --git a/proxy/optimisticresolver_internal_test.go b/proxy/optimisticresolver_internal_test.go
index e9c63c2f..50b2f680 100644
--- a/proxy/optimisticresolver_internal_test.go
+++ b/proxy/optimisticresolver_internal_test.go
@@ -2,25 +2,30 @@ package proxy
 
 import (
 	"bytes"
+	"context"
 	"log/slog"
 	"sync"
 	"testing"
 
 	"github.com/AdguardTeam/golibs/errors"
+	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/stretchr/testify/assert"
 )
 
 // testCachingResolver is a stub implementation of the cachingResolver interface
 // to simplify testing.
 type testCachingResolver struct {
-	onReplyFromUpstream func(dctx *DNSContext) (ok bool, err error)
+	onReplyFromUpstream func(ctx context.Context, dctx *DNSContext) (ok bool, err error)
 	onCacheResp         func(dctx *DNSContext)
 }
 
 // replyFromUpstream implements the cachingResolver interface for
 // *testCachingResolver.
-func (tcr *testCachingResolver) replyFromUpstream(dctx *DNSContext) (ok bool, err error) {
-	return tcr.onReplyFromUpstream(dctx)
+func (tcr *testCachingResolver) replyFromUpstream(
+	ctx context.Context,
+	dctx *DNSContext,
+) (ok bool, err error) {
+	return tcr.onReplyFromUpstream(ctx, dctx)
 }
 
 // cacheResp implements the cachingResolver interface for *testCachingResolver.
@@ -33,7 +38,7 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
 	var timesResolved, timesSet int
 
 	tcr := &testCachingResolver{
-		onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) {
+		onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
 			timesResolved++
 
 			return true, nil
@@ -52,7 +57,8 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
 	sameKey := []byte{1, 2, 3}
 
 	// Start the primary goroutine.
-	go s.resolveOnce(nil, sameKey, testLogger)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	go s.resolveOnce(ctx, nil, sameKey, testLogger)
 	// Block until the primary goroutine reaches the resolve function.
 	<-out
 
@@ -61,7 +67,8 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
 	const secondaryNum = 10
 	for range secondaryNum {
 		wg.Go(func() {
-			s.resolveOnce(nil, sameKey, testLogger)
+			resolveCtx := testutil.ContextWithTimeout(t, testTimeout)
+			s.resolveOnce(resolveCtx, nil, sameKey, testLogger)
 		})
 	}
 
@@ -90,10 +97,14 @@ func TestOptimisticResolver_ResolveOnce_unsuccessful(t *testing.T) {
 
 		cached := false
 		s := newOptimisticResolver(&testCachingResolver{
-			onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) { return true, rErr },
-			onCacheResp:         func(_ *DNSContext) { cached = true },
+			onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
+				return true, rErr
+			},
+			onCacheResp: func(_ *DNSContext) { cached = true },
 		})
-		s.resolveOnce(nil, key, l)
+
+		ctx := testutil.ContextWithTimeout(t, testTimeout)
+		s.resolveOnce(ctx, nil, key, l)
 
 		assert.True(t, cached)
 		assert.Contains(t, logOutput.String(), rErr.Error())
@@ -102,10 +113,14 @@ func TestOptimisticResolver_ResolveOnce_unsuccessful(t *testing.T) {
 	t.Run("not_ok", func(t *testing.T) {
 		cached := false
 		s := newOptimisticResolver(&testCachingResolver{
-			onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) { return false, nil },
-			onCacheResp:         func(_ *DNSContext) { cached = true },
+			onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
+				return false, nil
+			},
+			onCacheResp: func(_ *DNSContext) { cached = true },
 		})
-		s.resolveOnce(nil, key, testLogger)
+
+		ctx := testutil.ContextWithTimeout(t, testTimeout)
+		s.resolveOnce(ctx, nil, key, testLogger)
 
 		assert.False(t, cached)
 	})
diff --git a/proxy/pending_test.go b/proxy/pending_test.go
index 58ce8755..710f2026 100644
--- a/proxy/pending_test.go
+++ b/proxy/pending_test.go
@@ -64,7 +64,7 @@ func TestPendingRequests(t *testing.T) {
 
 	once := &sync.Once{}
 	u := &dnsproxytest.Upstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			once.Do(func() {
 				resp = (&dns.Msg{}).SetReply(req)
 			})
diff --git a/proxy/proxy.go b/proxy/proxy.go
index f6c67822..d072b2b7 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -791,7 +791,7 @@ func (p *Proxy) selectUpstreams(d *DNSContext) (upstreams []upstream.Upstream, i
 // replyFromUpstream tries to resolve the request via configured upstream
 // servers.  It returns true if the response actually came from an upstream.  d
 // must not be nil.
-func (p *Proxy) replyFromUpstream(d *DNSContext) (ok bool, err error) {
+func (p *Proxy) replyFromUpstream(ctx context.Context, d *DNSContext) (ok bool, err error) {
 	req := d.Req
 
 	upstreams, isPrivate := p.selectUpstreams(d)
@@ -809,40 +809,42 @@ func (p *Proxy) replyFromUpstream(d *DNSContext) (ok bool, err error) {
 	wrapped := upstreamsWithStats(upstreams)
 
 	// Perform the DNS request.
-	resp, u, err := p.exchangeUpstreams(req, wrapped)
-	if dns64Ups := p.performDNS64(req, resp, wrapped); dns64Ups != nil {
+	resp, u, err := p.exchangeUpstreams(ctx, req, wrapped)
+	if dns64Ups := p.performDNS64(ctx, req, resp, wrapped); dns64Ups != nil {
 		u = dns64Ups
 	} else if p.isBogusNXDomain(resp) {
-		p.logger.Debug("response contains bogus-nxdomain ip")
+		p.logger.DebugContext(ctx, "response contains bogus-nxdomain ip")
 		resp = p.messages.NewMsgNXDOMAIN(req)
 	}
 
 	var wrappedFallbacks []upstream.Upstream
 	if err != nil && !isPrivate && p.fallbacks != nil {
-		p.logger.Debug("using fallback", slogutil.KeyError, err)
+		p.logger.DebugContext(ctx, "using fallback", slogutil.KeyError, err)
 
 		src = "fallback"
 
-		// upstreams mustn't appear empty since they have been validated when
+		// Upstreams mustn't appear empty since they have been validated when
 		// creating proxy.
 		upstreams = p.fallbacks.getUpstreamsForDomain(req.Question[0].Name)
 
 		wrappedFallbacks = upstreamsWithStats(upstreams)
-		resp, u, err = upstream.ExchangeParallel(wrappedFallbacks, req)
+
+		// NOTE: Use [context.WithoutCancel] to have a separate timeout on
+		// retry.
+		resp, u, err = upstream.ExchangeParallel(context.WithoutCancel(ctx), wrappedFallbacks, req)
 	}
 
 	if err != nil {
-		p.logger.Debug("resolving err", "src", src, slogutil.KeyError, err)
+		p.logger.DebugContext(ctx, "resolving err", "src", src, slogutil.KeyError, err)
 	}
 
 	if resp != nil {
-		p.logger.Debug("resolved", "upstream", u.Address(), "src", src)
+		p.logger.DebugContext(ctx, "resolved", "upstream", u.Address(), "src", src)
 	}
 
 	unwrapped, stats := collectQueryStats(p.upstreamMode, u, wrapped, wrappedFallbacks)
 	d.queryStatistics = stats
 
-	ctx := context.TODO()
 	p.handleExchangeResult(ctx, d, req, resp, unwrapped)
 
 	return resp != nil, err
@@ -927,7 +929,7 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
 		}
 		defer func() { p.pendingRequests.done(ctx, dctx, err) }()
 
-		if p.replyFromCache(dctx) {
+		if p.replyFromCache(ctx, dctx) {
 			// Complete the response from cache.
 			filterMsg(dctx.Res, dctx.Res, dctx.adBit, dctx.doBit, 0)
 			dctx.scrub()
@@ -937,7 +939,7 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
 	}
 
 	var ok bool
-	ok, err = p.replyFromUpstream(dctx)
+	ok, err = p.replyFromUpstream(ctx, dctx)
 
 	// Don't cache the responses having CD flag, just like Dnsmasq does.  It
 	// prevents the cache from being poisoned with unvalidated answers which may
diff --git a/proxy/proxy_internal_test.go b/proxy/proxy_internal_test.go
index bc259ba0..9c056d9c 100644
--- a/proxy/proxy_internal_test.go
+++ b/proxy/proxy_internal_test.go
@@ -82,7 +82,7 @@ func newTestUpstreamConfig(tb testing.TB, ups ...upstream.Upstream) (u *Upstream
 func newTestUpstream(tb testing.TB) (uc upstream.Upstream) {
 	tb.Helper()
 
-	onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
+	onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		return dnsproxytest.NewTestResponse(req), nil
 	}
 
@@ -94,7 +94,7 @@ func newTestUpstream(tb testing.TB) (uc upstream.Upstream) {
 // be replaced with stub implementations.
 func newTestUpstreamWithExchange(
 	tb testing.TB,
-	onExchange func(req *dns.Msg) (resp *dns.Msg, err error),
+	onExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error),
 ) (uc upstream.Upstream) {
 	tb.Helper()
 
@@ -220,7 +220,7 @@ func TestProxy_Resolve_dnssecCache(t *testing.T) {
 	}
 
 	u := &testUpstream{
-		OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 			resp = (&dns.Msg{}).SetReply(m)
 
 			q := m.Question[0]
@@ -348,9 +348,12 @@ func TestProxy_Resolve_dnssecCache(t *testing.T) {
 func TestProxy_HandleDNSRequest_exchangeWithReservedDomains(t *testing.T) {
 	t.Parallel()
 
-	emptyUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
-		return (&dns.Msg{}).SetReply(req), nil
-	})
+	emptyUpstream := newTestUpstreamWithExchange(
+		t,
+		func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
+			return (&dns.Msg{}).SetReply(req), nil
+		},
+	)
 
 	const (
 		host1 = "example-1.test"
@@ -422,9 +425,12 @@ func TestProxy_HandleDNSRequest_exchangeWithReservedDomains(t *testing.T) {
 func TestProxy_HandleDNSRequest_oneByOneUpstreamsExchange(t *testing.T) {
 	t.Parallel()
 
-	errUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
-		return nil, assert.AnError
-	})
+	errUpstream := newTestUpstreamWithExchange(
+		t,
+		func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
+			return nil, assert.AnError
+		},
+	)
 
 	dnsProxy := mustNew(t, &Config{
 		Logger:        testLogger,
@@ -477,18 +483,21 @@ func TestProxy_HandleDNSRequest_fallback(t *testing.T) {
 
 	pt := testutil.NewPanicT(t)
 
-	successExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
+	successExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		testutil.RequireSend(pt, responseCh, req.Id, dnsproxytest.Timeout)
 
 		return dnsproxytest.NewTestResponse(req), nil
 	}
 	successUpstream := newTestUpstreamWithExchange(t, successExchange)
 
-	failUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
-		testutil.RequireSend(pt, failCh, req.Id, dnsproxytest.Timeout)
+	failUpstream := newTestUpstreamWithExchange(
+		t,
+		func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
+			testutil.RequireSend(pt, failCh, req.Id, dnsproxytest.Timeout)
 
-		return nil, assert.AnError
-	})
+			return nil, assert.AnError
+		},
+	)
 
 	upsConf := &UpstreamConfig{
 		DomainReservedUpstreams: map[string][]upstream.Upstream{
@@ -636,7 +645,7 @@ func TestProxy_Resolve_customUpstreamConfigCache(t *testing.T) {
 
 	var count int
 
-	exchangeFunc := func(m *dns.Msg) (resp *dns.Msg, err error) {
+	exchangeFunc := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 		resp = &dns.Msg{}
 		resp.SetReply(m)
 		resp.Answer = append(resp.Answer, &dns.A{
@@ -992,7 +1001,7 @@ func TestProxy_Resolve_withOptimisticResolver(t *testing.T) {
 	p.initCache()
 	out, in := make(chan unit), make(chan unit)
 	p.shortFlighter.cr = &testCachingResolver{
-		onReplyFromUpstream: func(dctx *DNSContext) (ok bool, err error) {
+		onReplyFromUpstream: func(_ context.Context, dctx *DNSContext) (ok bool, err error) {
 			dctx.Res = buildResp(dctx.Req, nonOptimisticTTL)
 
 			return true, nil
@@ -1077,7 +1086,7 @@ func TestProxy_ValidateRequest(t *testing.T) {
 	}
 
 	ups := &testUpstream{
-		OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 			resp = &dns.Msg{}
 			resp.SetReply(m)
 
@@ -1206,7 +1215,7 @@ func (h *testHandler) ServeDNS(ctx context.Context, p *Proxy, dctx *DNSContext)
 // TODO(m.kazantsev):  Use [dnsproxytest.Upstream].
 type testUpstream struct {
 	OnAddress  func() (addr string)
-	OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
+	OnExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
 	OnClose    func() (err error)
 }
 
@@ -1216,8 +1225,12 @@ type testUpstream struct {
 // The handler dereferences ans, ecsIP, and ecsReqIP at call time, so tests
 // can reassign those variables between sub-tests and the handler will see
 // the current values.
-func newECSReplyHandler(ans *[]dns.RR, ecsIP, ecsReqIP *net.IP) func(*dns.Msg) (*dns.Msg, error) {
-	return func(m *dns.Msg) (resp *dns.Msg, err error) {
+func newECSReplyHandler(
+	ans *[]dns.RR,
+	ecsIP *net.IP,
+	ecsReqIP *net.IP,
+) func(context.Context, *dns.Msg) (*dns.Msg, error) {
+	return func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 		resp = (&dns.Msg{}).SetReply(m)
 
 		if ans != nil && *ans != nil {
@@ -1245,8 +1258,8 @@ func newECSReplyHandler(ans *[]dns.RR, ecsIP, ecsReqIP *net.IP) func(*dns.Msg) (
 var _ upstream.Upstream = (*testUpstream)(nil)
 
 // Exchange implements the upstream.Upstream interface for *testUpstream.
-func (u *testUpstream) Exchange(m *dns.Msg) (resp *dns.Msg, err error) {
-	return u.OnExchange(m)
+func (u *testUpstream) Exchange(ctx context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
+	return u.OnExchange(ctx, m)
 }
 
 // Address implements the upstream.Upstream interface for *testUpstream.
diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go
index 2c0d0cf3..cd8661f7 100644
--- a/proxy/proxy_test.go
+++ b/proxy/proxy_test.go
@@ -1,6 +1,7 @@
 package proxy_test
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"os"
@@ -44,7 +45,7 @@ var (
 //
 // TODO(f.setrakov): DRY with internal version.
 func newTestUpstreamConfig(tb testing.TB) (uc *proxy.UpstreamConfig) {
-	onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
+	onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		pt := testutil.NewPanicT(tb)
 
 		require.Len(pt, req.Question, 1)
@@ -103,7 +104,7 @@ func sendTestMessages(tb testing.TB, conn *dns.Conn) {
 func TestProxy_Resolve_badResponse(t *testing.T) {
 	dnsProxy := mustStartDefaultProxy(t)
 
-	onExchange := func(m *dns.Msg) (resp *dns.Msg, err error) {
+	onExchange := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
 		resp = (&dns.Msg{}).SetReply(m)
 		resp.Answer = append(resp.Answer, &dns.A{
 			Hdr: dns.RR_Header{
@@ -257,7 +258,7 @@ func TestProxy_Resolve_cache(t *testing.T) {
 		OnAddress: func() (addr string) { return "stub" },
 		OnClose:   func() (err error) { return nil },
 	}
-	ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
+	ups.OnExchange = func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		resp = (&dns.Msg{}).SetReply(req)
 		resp.Answer = append(resp.Answer, &dns.A{
 			Hdr: dns.RR_Header{
@@ -349,9 +350,11 @@ func TestProxy_Start_closeOnFail(t *testing.T) {
 	tcpAddr := testutil.RequireTypeAssert[*net.TCPAddr](t, l.Addr())
 
 	ups := &dnsproxytest.Upstream{
-		OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
-		OnAddress:  func() (_ string) { panic(testutil.UnexpectedCall()) },
-		OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
+		OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
+			panic(testutil.UnexpectedCall(ctx, m))
+		},
+		OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
+		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 
 	p, err := proxy.New(&proxy.Config{
@@ -389,7 +392,7 @@ func TestProxy_ServeDNS_formatError(t *testing.T) {
 		OnAddress: func() (addr string) { return testIPv4.String() },
 		OnClose:   func() (err error) { return nil },
 	}
-	ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
+	ups.OnExchange = func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		panic(testutil.UnexpectedCall(req))
 	}
 
diff --git a/proxy/proxycache.go b/proxy/proxycache.go
index 98a350b5..cb3c56ce 100644
--- a/proxy/proxycache.go
+++ b/proxy/proxycache.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"net"
 	"slices"
 )
@@ -17,7 +18,7 @@ func (p *Proxy) cacheForContext(d *DNSContext) (c *cache) {
 // replyFromCache tries to get the response from general or subnet cache.  In
 // case the cache is present in d, it's used first.  Returns true on success.  d
 // must not be nil.
-func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
+func (p *Proxy) replyFromCache(ctx context.Context, d *DNSContext) (hit bool) {
 	dctxCache := p.cacheForContext(d)
 
 	var ci *cacheItem
@@ -41,7 +42,8 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
 	d.Res = ci.m
 	d.queryStatistics = cachedQueryStatistics(ci.u)
 
-	p.logger.Debug(
+	p.logger.DebugContext(
+		ctx,
 		"replying from cache",
 		"source", cacheSource,
 		"ecs_enabled", p.enableEDNSClientSubnet,
@@ -59,7 +61,7 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
 			minCtxClone.Req = d.Req.Copy()
 		}
 
-		go p.shortFlighter.resolveOnce(minCtxClone, key, p.logger)
+		go p.shortFlighter.resolveOnce(ctx, minCtxClone, key, p.logger)
 	}
 
 	return hit
diff --git a/proxy/servertcp_test.go b/proxy/servertcp_test.go
index 1467f3e5..59095196 100644
--- a/proxy/servertcp_test.go
+++ b/proxy/servertcp_test.go
@@ -1,6 +1,7 @@
 package proxy_test
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/binary"
@@ -100,9 +101,11 @@ func TestProxy_handleDNSRequest_emptyTCPMessage(t *testing.T) {
 	t.Parallel()
 
 	u := &dnsproxytest.Upstream{
-		OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
-		OnAddress:  func() (_ string) { panic(testutil.UnexpectedCall()) },
-		OnClose:    func() (err error) { return nil },
+		OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
+			panic(testutil.UnexpectedCall(ctx, m))
+		},
+		OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
+		OnClose:   func() (err error) { return nil },
 	}
 	upsConf := &proxy.UpstreamConfig{
 		Upstreams: []upstream.Upstream{u},
diff --git a/proxy/stats.go b/proxy/stats.go
index 4cb0830e..ce03b8be 100644
--- a/proxy/stats.go
+++ b/proxy/stats.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -25,9 +26,9 @@ type upstreamWithStats struct {
 var _ upstream.Upstream = (*upstreamWithStats)(nil)
 
 // Exchange implements the [upstream.Upstream] for *upstreamWithStats.
-func (u *upstreamWithStats) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (u *upstreamWithStats) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	start := time.Now()
-	resp, err = u.upstream.Exchange(req)
+	resp, err = u.upstream.Exchange(ctx, req)
 	u.err = err
 	u.queryDuration = time.Since(start)
 
diff --git a/proxy/stats_test.go b/proxy/stats_test.go
index f9f74166..dc8b2807 100644
--- a/proxy/stats_test.go
+++ b/proxy/stats_test.go
@@ -1,6 +1,7 @@
 package proxy_test
 
 import (
+	"context"
 	"net"
 	"testing"
 
@@ -27,7 +28,7 @@ func TestDNSContext_QueryStatistics(t *testing.T) {
 	}
 
 	ups := &dnsproxytest.Upstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			return (&dns.Msg{}).SetReply(req), nil
 		},
 		OnAddress: func() (addr string) { return "upstream" },
@@ -35,7 +36,7 @@ func TestDNSContext_QueryStatistics(t *testing.T) {
 	}
 
 	failUps := &dnsproxytest.Upstream{
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			return nil, errors.Error("exchange error")
 		},
 		OnAddress: func() (addr string) { return "fail.upstream" },
diff --git a/upstream/dnscrypt.go b/upstream/dnscrypt.go
index 040d2770..3ee1680b 100644
--- a/upstream/dnscrypt.go
+++ b/upstream/dnscrypt.go
@@ -54,8 +54,7 @@ var _ Upstream = (*dnsCrypt)(nil)
 func (p *dnsCrypt) Address() string { return p.addr.String() }
 
 // Exchange implements the [Upstream] interface for *dnsCrypt.
-func (p *dnsCrypt) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
-	ctx := context.Background()
+func (p *dnsCrypt) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	if p.timeout > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, p.timeout)
@@ -103,7 +102,8 @@ func (p *dnsCrypt) exchangeDNSCrypt(ctx context.Context, req *dns.Msg) (resp *dn
 	resp, err = client.ExchangeContext(ctx, req, resolverInfo)
 	if resp != nil && resp.Truncated {
 		q := &req.Question[0]
-		p.logger.Debug(
+		p.logger.DebugContext(
+			ctx,
 			"dnscrypt received truncated, falling back to tcp",
 			"addr", p.addr,
 			"question", q,
@@ -114,6 +114,15 @@ func (p *dnsCrypt) exchangeDNSCrypt(ctx context.Context, req *dns.Msg) (resp *dn
 			Proto:  dnscrypt.ProtoTCP,
 		})
 
+		// NOTE: Give the TCP retry its own timeout.
+		if p.timeout > 0 {
+			ctx = context.WithoutCancel(ctx)
+
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithTimeout(ctx, p.timeout)
+			defer cancel()
+		}
+
 		resp, err = tcpClient.ExchangeContext(ctx, req, resolverInfo)
 	}
 	if err == nil && resp != nil && resp.Id != req.Id {
diff --git a/upstream/dnscrypt_internal_test.go b/upstream/dnscrypt_internal_test.go
index 8d12a6a9..828403fa 100644
--- a/upstream/dnscrypt_internal_test.go
+++ b/upstream/dnscrypt_internal_test.go
@@ -111,7 +111,7 @@ func TestUpstreamDNSCrypt(t *testing.T) {
 
 	// Test that it responds properly
 	for range 10 {
-		checkUpstream(t, u, address)
+		checkUpstream(t, u, address, testTimeout)
 	}
 }
 
@@ -161,8 +161,9 @@ func TestDNSCrypt_Exchange_truncated(t *testing.T) {
 
 	req := (&dns.Msg{}).SetQuestion("unit-test2.dns.adguard.com.", dns.TypeTXT)
 
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
 	// Check that response is not truncated (even though it's huge).
-	res, err := u.Exchange(req)
+	res, err := u.Exchange(ctx, req)
 	require.NoError(t, err)
 
 	assert.False(t, res.Truncated)
@@ -191,7 +192,8 @@ func TestDNSCrypt_Exchange_deadline(t *testing.T) {
 
 	req := (&dns.Msg{}).SetQuestion("unit-test2.dns.adguard.com.", dns.TypeTXT)
 
-	res, err := u.Exchange(req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	res, err := u.Exchange(ctx, req)
 	require.ErrorIs(t, err, context.DeadlineExceeded)
 
 	assert.Nil(t, res)
@@ -219,8 +221,10 @@ func TestDNSCrypt_Exchange_dialFail(t *testing.T) {
 	require.True(t, t.Run("dial_fail", func(t *testing.T) {
 		testutil.CleanupAndRequireSuccess(t, u.Close)
 
+		ctx := testutil.ContextWithTimeout(t, testTimeout)
+
 		var res *dns.Msg
-		res, err = u.Exchange(req)
+		res, err = u.Exchange(ctx, req)
 		require.Error(t, err)
 
 		assert.Nil(t, res)
@@ -241,8 +245,10 @@ func TestDNSCrypt_Exchange_dialFail(t *testing.T) {
 		})
 		require.NoError(t, err)
 
+		ctx := testutil.ContextWithTimeout(t, testTimeout)
+
 		var res *dns.Msg
-		res, err = u.Exchange(req)
+		res, err = u.Exchange(ctx, req)
 		require.ErrorIs(t, err, validationErr)
 
 		assert.Nil(t, res)
diff --git a/upstream/doh.go b/upstream/doh.go
index ac44e189..b4793241 100644
--- a/upstream/doh.go
+++ b/upstream/doh.go
@@ -153,16 +153,16 @@ var _ Upstream = (*dnsOverHTTPS)(nil)
 func (p *dnsOverHTTPS) Address() string { return p.addrRedacted }
 
 // Exchange implements the [Upstream] interface for *dnsOverHTTPS.
-func (p *dnsOverHTTPS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (p *dnsOverHTTPS) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	// Check if there was already an active client before sending the request.
 	// We'll only attempt to re-connect if there was one.
-	client, isCached, err := p.getClient()
+	client, isCached, err := p.getClient(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to init http client: %w", err)
 	}
 
 	// Make the first attempt to send the DNS query.
-	resp, err = p.exchangeHTTPS(client, req)
+	resp, err = p.exchangeHTTPS(ctx, client, req)
 
 	// Make up to 2 attempts to re-create the HTTP client and send the request
 	// again.  There are several cases (mostly, with QUIC) where this workaround
@@ -170,17 +170,17 @@ func (p *dnsOverHTTPS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
 	// the case when the connection was closed (due to inactivity for example)
 	// AND the server refuses to open a 0-RTT connection.
 	for i := 0; isCached && p.shouldRetry(err) && i < 2; i++ {
-		client, err = p.resetClient(err)
+		client, err = p.resetClient(ctx, err)
 		if err != nil {
 			return nil, fmt.Errorf("failed to reset http client: %w", err)
 		}
 
-		resp, err = p.exchangeHTTPS(client, req)
+		resp, err = p.exchangeHTTPS(ctx, client, req)
 	}
 
 	if err != nil {
 		// If the request failed anyway, make sure we don't use this client.
-		_, resErr := p.resetClient(err)
+		_, resErr := p.resetClient(ctx, err)
 
 		return nil, errors.WithDeferred(err, resErr)
 	}
@@ -217,14 +217,18 @@ func (p *dnsOverHTTPS) closeClient(client *http.Client) (err error) {
 
 // exchangeHTTPS logs the request and its result and calls exchangeHTTPSClient.
 // client and req must not be nil.
-func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *dns.Msg, err error) {
+func (p *dnsOverHTTPS) exchangeHTTPS(
+	ctx context.Context,
+	client *http.Client,
+	req *dns.Msg,
+) (resp *dns.Msg, err error) {
 	n := networkTCP
 	if isHTTP3(client) {
 		n = networkUDP
 	}
 
-	logBegin(p.logger, p.addrRedacted, n, req)
-	defer func() { logFinish(p.logger, p.addrRedacted, n, err) }()
+	logBegin(ctx, p.logger, p.addrRedacted, n, req)
+	defer func() { logFinish(ctx, p.logger, p.addrRedacted, n, err) }()
 
 	buf, err := req.Pack()
 	if err != nil {
@@ -238,7 +242,7 @@ func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *d
 	// See https://www.rfc-editor.org/rfc/rfc8484.html.
 	binary.BigEndian.PutUint16(buf, 0)
 
-	resp, err = p.exchangeHTTPSClient(client, buf)
+	resp, err = p.exchangeHTTPSClient(ctx, client, buf)
 	if err != nil {
 		return nil, fmt.Errorf("exchanging: %w", err)
 	}
@@ -264,6 +268,7 @@ func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *d
 // http.Client instance.  buf is the packed DNS message that will be sent to the
 // resolver.  client must not be nil.
 func (p *dnsOverHTTPS) exchangeHTTPSClient(
+	ctx context.Context,
 	client *http.Client,
 	buf []byte,
 ) (resp *dns.Msg, err error) {
@@ -287,7 +292,7 @@ func (p *dnsOverHTTPS) exchangeHTTPSClient(
 		RawQuery: q.Encode(),
 	}
 
-	httpReq, err := http.NewRequest(method, u.String(), nil)
+	httpReq, err := http.NewRequestWithContext(ctx, method, u.String(), nil)
 	if err != nil {
 		return nil, fmt.Errorf("creating http request to %s: %w", p.addrRedacted, err)
 	}
@@ -301,7 +306,7 @@ func (p *dnsOverHTTPS) exchangeHTTPSClient(
 	if err != nil {
 		return nil, fmt.Errorf("requesting %s: %w", p.addrRedacted, err)
 	}
-	defer slogutil.CloseAndLog(httpReq.Context(), p.logger, httpResp.Body, slog.LevelDebug)
+	defer slogutil.CloseAndLog(ctx, p.logger, httpResp.Body, slog.LevelDebug)
 
 	limitBody := ioutil.LimitReader(httpResp.Body, dns.MaxMsgSize)
 	body, err := io.ReadAll(limitBody)
@@ -359,7 +364,10 @@ func (p *dnsOverHTTPS) shouldRetry(err error) (ok bool) {
 // resetClient triggers re-creation of the *http.Client that is used by this
 // upstream.  This method accepts the error that caused resetting client as
 // depending on the error we may also reset the QUIC config.
-func (p *dnsOverHTTPS) resetClient(resetErr error) (client *http.Client, err error) {
+func (p *dnsOverHTTPS) resetClient(
+	ctx context.Context,
+	resetErr error,
+) (client *http.Client, err error) {
 	p.clientMu.Lock()
 	defer p.clientMu.Unlock()
 
@@ -372,12 +380,16 @@ func (p *dnsOverHTTPS) resetClient(resetErr error) (client *http.Client, err err
 	if oldClient != nil {
 		closeErr := p.closeClient(oldClient)
 		if closeErr != nil {
-			p.logger.Warn("failed to close the old http client", slogutil.KeyError, closeErr)
+			p.logger.WarnContext(
+				ctx,
+				"failed to close the old http client",
+				slogutil.KeyError, closeErr,
+			)
 		}
 	}
 
-	p.logger.Debug("recreating the http client", slogutil.KeyError, resetErr)
-	p.client, err = p.createClient()
+	p.logger.DebugContext(ctx, "recreating the http client", slogutil.KeyError, resetErr)
+	p.client, err = p.createClient(ctx)
 
 	return p.client, err
 }
@@ -403,7 +415,7 @@ func (p *dnsOverHTTPS) resetQUICConfig() {
 
 // getClient gets or lazily initializes an HTTP client (and transport) that will
 // be used for this DoH resolver.
-func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
+func (p *dnsOverHTTPS) getClient(ctx context.Context) (c *http.Client, isCached bool, err error) {
 	startTime := time.Now()
 
 	p.clientMu.Lock()
@@ -420,8 +432,8 @@ func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
 		return nil, false, fmt.Errorf("timeout exceeded: %s", elapsed)
 	}
 
-	p.logger.Debug("creating a new http client")
-	p.client, err = p.createClient()
+	p.logger.DebugContext(ctx, "creating a new http client")
+	p.client, err = p.createClient(ctx)
 
 	return p.client, false, err
 }
@@ -430,8 +442,8 @@ func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
 // will depend on whether HTTP3 is allowed and provided by this upstream.  Note,
 // that we'll attempt to establish a QUIC connection when creating the client in
 // order to check whether HTTP3 is supported.
-func (p *dnsOverHTTPS) createClient() (*http.Client, error) {
-	transport, err := p.createTransport()
+func (p *dnsOverHTTPS) createClient(ctx context.Context) (*http.Client, error) {
+	transport, err := p.createTransport(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("initializing http transport: %w", err)
 	}
@@ -455,7 +467,7 @@ func (p *dnsOverHTTPS) createClient() (*http.Client, error) {
 // that this function will first attempt to establish a QUIC connection (if
 // HTTP3 is enabled in the upstream options).  If this attempt is successful,
 // it returns an HTTP3 transport, otherwise it returns the H1/H2 transport.
-func (p *dnsOverHTTPS) createTransport() (t http.RoundTripper, err error) {
+func (p *dnsOverHTTPS) createTransport(ctx context.Context) (t http.RoundTripper, err error) {
 	dialContext, err := p.getDialer()
 	if err != nil {
 		return nil, fmt.Errorf("bootstrapping %s: %w", p.addrRedacted, err)
@@ -465,14 +477,14 @@ func (p *dnsOverHTTPS) createTransport() (t http.RoundTripper, err error) {
 	// connection is established successfully, we'll be using HTTP3 for this
 	// upstream.
 	tlsConf := p.tlsConf.Clone()
-	transportH3, err := p.createTransportH3(tlsConf, dialContext)
+	transportH3, err := p.createTransportH3(ctx, tlsConf, dialContext)
 	if err == nil {
-		p.logger.Debug("using http/3 for this upstream, quic was faster")
+		p.logger.DebugContext(ctx, "using http/3 for this upstream, quic was faster")
 
 		return transportH3, nil
 	}
 
-	p.logger.Debug("got error, switching to http/2 for this upstream", slogutil.KeyError, err)
+	p.logger.DebugContext(ctx, "got error, switching to http/2 for this upstream", slogutil.KeyError, err)
 
 	if !p.supportsHTTP() {
 		return nil, errors.Error("HTTP1/1 and HTTP2 are not supported by this upstream")
@@ -558,6 +570,7 @@ func (h *http3Transport) Close() (err error) {
 // parallel (one for TLS, the other one for QUIC) and if QUIC is faster it will
 // create the [*http3.Transport] instance.
 func (p *dnsOverHTTPS) createTransportH3(
+	ctx context.Context,
 	tlsConfig *tls.Config,
 	dialContext bootstrap.DialHandler,
 ) (roundTripper http.RoundTripper, err error) {
@@ -565,7 +578,7 @@ func (p *dnsOverHTTPS) createTransportH3(
 		return nil, errors.Error("HTTP3 support is not enabled")
 	}
 
-	addr, err := p.probeH3(tlsConfig, dialContext)
+	addr, err := p.probeH3(ctx, tlsConfig, dialContext)
 	if err != nil {
 		return nil, err
 	}
@@ -594,13 +607,14 @@ func (p *dnsOverHTTPS) createTransportH3(
 // upstream.  If the test is successful it will return the address that we
 // should use to establish the QUIC connections.
 func (p *dnsOverHTTPS) probeH3(
+	ctx context.Context,
 	tlsConfig *tls.Config,
 	dialContext bootstrap.DialHandler,
 ) (addr string, err error) {
 	// We're using bootstrapped address instead of what's passed to the function
 	// it does not create an actual connection, but it helps us determine
 	// what IP is actually reachable (when there are v4/v6 addresses).
-	rawConn, err := dialContext(context.Background(), "udp", "")
+	rawConn, err := dialContext(ctx, "udp", "")
 	if err != nil {
 		return "", fmt.Errorf("failed to dial: %w", err)
 	}
@@ -635,8 +649,8 @@ func (p *dnsOverHTTPS) probeH3(
 	// Run probeQUIC and probeTLS in parallel and see which one is faster.
 	chQUIC := make(chan error, 1)
 	chTLS := make(chan error, 1)
-	go p.probeQUIC(addr, probeTLSCfg, chQUIC)
-	go p.probeTLS(dialContext, probeTLSCfg, chTLS)
+	go p.probeQUIC(ctx, addr, probeTLSCfg, chQUIC)
+	go p.probeTLS(ctx, dialContext, probeTLSCfg, chTLS)
 
 	select {
 	case quicErr := <-chQUIC:
@@ -650,7 +664,7 @@ func (p *dnsOverHTTPS) probeH3(
 	case tlsErr := <-chTLS:
 		if tlsErr != nil {
 			// Return immediately, TLS failed.
-			p.logger.Debug("probing tls", slogutil.KeyError, tlsErr)
+			p.logger.DebugContext(ctx, "probing tls", slogutil.KeyError, tlsErr)
 
 			return addr, nil
 		}
@@ -661,14 +675,20 @@ func (p *dnsOverHTTPS) probeH3(
 
 // probeQUIC attempts to establish a QUIC connection to the specified address.
 // We run probeQUIC and probeTLS in parallel and see which one is faster.
-func (p *dnsOverHTTPS) probeQUIC(addr string, tlsConfig *tls.Config, ch chan error) {
+func (p *dnsOverHTTPS) probeQUIC(
+	ctx context.Context,
+	addr string,
+	tlsConfig *tls.Config,
+	ch chan error,
+) {
 	startTime := time.Now()
 
 	t := p.timeout
 	if t == 0 {
 		t = dialTimeout
 	}
-	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(t))
+
+	ctx, cancel := context.WithTimeout(ctx, t)
 	defer cancel()
 
 	conn, err := quic.DialAddrEarly(ctx, addr, tlsConfig, p.getQUICConfig())
@@ -683,15 +703,20 @@ func (p *dnsOverHTTPS) probeQUIC(addr string, tlsConfig *tls.Config, ch chan err
 	ch <- nil
 
 	elapsed := time.Since(startTime)
-	p.logger.Debug("quic connection established", "elapsed", elapsed)
+	p.logger.DebugContext(ctx, "quic connection established", "elapsed", elapsed)
 }
 
 // probeTLS attempts to establish a TLS connection to the specified address. We
 // run probeQUIC and probeTLS in parallel and see which one is faster.
-func (p *dnsOverHTTPS) probeTLS(dialContext bootstrap.DialHandler, tlsConfig *tls.Config, ch chan error) {
+func (p *dnsOverHTTPS) probeTLS(
+	ctx context.Context,
+	dialContext bootstrap.DialHandler,
+	tlsConfig *tls.Config,
+	ch chan error,
+) {
 	startTime := time.Now()
 
-	conn, err := tlsDial(dialContext, tlsConfig)
+	conn, err := tlsDial(ctx, dialContext, tlsConfig)
 	if err != nil {
 		ch <- fmt.Errorf("opening TLS connection: %w", err)
 		return
@@ -703,7 +728,7 @@ func (p *dnsOverHTTPS) probeTLS(dialContext bootstrap.DialHandler, tlsConfig *tl
 	ch <- nil
 
 	elapsed := time.Since(startTime)
-	p.logger.Debug("tls connection established", "elapsed", elapsed)
+	p.logger.DebugContext(ctx, "tls connection established", "elapsed", elapsed)
 }
 
 // supportsH3 returns true if HTTP/3 is supported by this upstream.
diff --git a/upstream/doh_internal_test.go b/upstream/doh_internal_test.go
index a358b861..19722d7d 100644
--- a/upstream/doh_internal_test.go
+++ b/upstream/doh_internal_test.go
@@ -97,7 +97,7 @@ func TestUpstreamDoH(t *testing.T) {
 
 			// Test that it responds properly.
 			for range 10 {
-				checkUpstream(t, u, address)
+				checkUpstream(t, u, address, testTimeout)
 			}
 
 			doh := testutil.RequireTypeAssert[*dnsOverHTTPS](t, u)
@@ -106,7 +106,7 @@ func TestUpstreamDoH(t *testing.T) {
 			doh.client = nil
 
 			// Force it to establish the connection again.
-			checkUpstream(t, u, address)
+			checkUpstream(t, u, address, testTimeout)
 
 			// Check that TLS session was resumed properly.
 			require.True(t, lastState.DidResume)
@@ -194,7 +194,7 @@ func TestUpstreamDoH_raceReconnect(t *testing.T) {
 			require.NoError(t, err)
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
-			checkRaceCondition(u)
+			checkRaceCondition(t, u)
 		})
 	}
 }
@@ -238,7 +238,7 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
 				})
 				require.NoError(t, err)
 
-				checkUpstream(t, u, upsAddr)
+				checkUpstream(t, u, upsAddr, testTimeout)
 			})
 			require.False(t, t.Failed())
 			testutil.CleanupAndRequireSuccess(t, u.Close)
@@ -249,12 +249,13 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
 					port:         int(addr.Port()),
 				})
 
-				checkUpstream(t, u, upsAddr)
+				checkUpstream(t, u, upsAddr, testTimeout)
 			})
 			require.False(t, t.Failed())
 
 			t.Run("retry", func(t *testing.T) {
-				_, err := u.Exchange(createTestMessage())
+				ctx := testutil.ContextWithTimeout(t, testTimeout)
+				_, err := u.Exchange(ctx, createTestMessage())
 				require.Error(t, err)
 
 				_ = startDoHServer(t, testDoHServerOptions{
@@ -262,7 +263,7 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
 					port:         int(addr.Port()),
 				})
 
-				checkUpstream(t, u, upsAddr)
+				checkUpstream(t, u, upsAddr, testTimeout)
 			})
 		})
 	}
@@ -291,7 +292,8 @@ func TestUpstreamDoH_0RTT(t *testing.T) {
 	req := createTestMessage()
 
 	// Trigger connection to a DoH3 server.
-	resp, err := uh.Exchange(req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	resp, err := uh.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, resp)
 
@@ -307,7 +309,8 @@ func TestUpstreamDoH_0RTT(t *testing.T) {
 	}()
 
 	// Trigger second connection.
-	resp, err = uh.Exchange(req)
+	ctx = testutil.ContextWithTimeout(t, testTimeout)
+	resp, err = uh.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, resp)
 
diff --git a/upstream/doq.go b/upstream/doq.go
index 604b6630..98873c7b 100644
--- a/upstream/doq.go
+++ b/upstream/doq.go
@@ -167,7 +167,7 @@ var _ Upstream = (*dnsOverQUIC)(nil)
 func (p *dnsOverQUIC) Address() string { return p.addr.String() }
 
 // Exchange implements the [Upstream] interface for *dnsOverQUIC.
-func (p *dnsOverQUIC) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (p *dnsOverQUIC) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	// When sending queries over a QUIC connection, the DNS Message ID MUST be
 	// set to 0.  The stream mapping for DoQ allows for unambiguous correlation
 	// of queries and responses, so the Message ID field is not required.
@@ -184,39 +184,43 @@ func (p *dnsOverQUIC) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
 	}()
 
 	// Gets or opens a QUIC connection to use for this query.
-	conn, cached, err := p.getConnection()
+	conn, cached, err := p.getConnection(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("getting conn: %w", err)
 	}
 
 	// Make the first attempt to send the DNS query.
-	resp, err = p.exchangeQUIC(req, conn)
+	resp, err = p.exchangeQUIC(ctx, req, conn)
 
 	// Failure to use a cached connection should be handled gracefully as this
 	// connection could have been closed by the server or simply be broken due
 	// to how UDP NAT works.  In this case the connection should be re-created.
 	if cached && err != nil {
-		p.logger.Debug("recreating the quic connection and retrying", slogutil.KeyError, err)
+		p.logger.DebugContext(
+			ctx,
+			"recreating the quic connection and retrying",
+			slogutil.KeyError, err,
+		)
 
 		// Close the active connection to make sure the cached connection is
 		// cleaned up.
-		p.closeConnWithError(conn, err)
+		p.closeConnWithError(ctx, conn, err)
 
 		// Get or re-create the QUIC connection in order to make the second
 		// attempt.
-		conn, _, err = p.getConnection()
+		conn, _, err = p.getConnection(ctx)
 		if err != nil {
 			return nil, fmt.Errorf("getting new conn: %w", err)
 		}
 
 		// Retry sending the request through the new connection.
-		resp, err = p.exchangeQUIC(req, conn)
+		resp, err = p.exchangeQUIC(ctx, req, conn)
 	}
 
 	if err != nil {
 		// If we're unable to exchange messages, make sure the connection is
 		// closed and signal about an internal error.
-		p.closeConnWithError(conn, err)
+		p.closeConnWithError(ctx, conn, err)
 	}
 
 	return resp, err
@@ -238,18 +242,22 @@ func (p *dnsOverQUIC) Close() (err error) {
 
 // exchangeQUIC attempts to open a new QUIC stream, send the DNS message
 // through it and return the response it got from the server.
-func (p *dnsOverQUIC) exchangeQUIC(req *dns.Msg, conn *quic.Conn) (resp *dns.Msg, err error) {
+func (p *dnsOverQUIC) exchangeQUIC(
+	ctx context.Context,
+	req *dns.Msg,
+	conn *quic.Conn,
+) (resp *dns.Msg, err error) {
 	addr := p.Address()
 
-	logBegin(p.logger, addr, networkUDP, req)
-	defer func() { logFinish(p.logger, addr, networkUDP, err) }()
+	logBegin(ctx, p.logger, addr, networkUDP, req)
+	defer func() { logFinish(ctx, p.logger, addr, networkUDP, err) }()
 
 	buf, err := req.Pack()
 	if err != nil {
 		return nil, fmt.Errorf("failed to pack DNS message for DoQ: %w", err)
 	}
 
-	stream, err := p.openStream(conn)
+	stream, err := p.openStream(ctx, conn)
 	if err != nil {
 		return nil, fmt.Errorf("opening stream: %w", err)
 	}
@@ -272,7 +280,7 @@ func (p *dnsOverQUIC) exchangeQUIC(req *dns.Msg, conn *quic.Conn) (resp *dns.Msg
 	// of the stream, but does not prevent reading from it.
 	err = stream.Close()
 	if err != nil {
-		p.logger.Debug("closing quic stream", slogutil.KeyError, err)
+		p.logger.DebugContext(ctx, "closing quic stream", slogutil.KeyError, err)
 	}
 
 	return p.readMsg(stream)
@@ -298,7 +306,7 @@ func (p *dnsOverQUIC) getBytesPool() (pool *sync.Pool) {
 
 // getConnection opens or returns an existing *quic.Conn and indicates whether
 // it opened a new connection or used an existing cached one.
-func (p *dnsOverQUIC) getConnection() (conn *quic.Conn, cached bool, err error) {
+func (p *dnsOverQUIC) getConnection(ctx context.Context) (conn *quic.Conn, cached bool, err error) {
 	p.connMu.Lock()
 	defer p.connMu.Unlock()
 
@@ -307,7 +315,7 @@ func (p *dnsOverQUIC) getConnection() (conn *quic.Conn, cached bool, err error)
 		return conn, true, nil
 	}
 
-	conn, err = p.openConnection()
+	conn, err = p.openConnection(ctx)
 	if err != nil {
 		return nil, false, err
 	}
@@ -337,8 +345,8 @@ func (p *dnsOverQUIC) resetQUICConfig() {
 }
 
 // openStream opens a new QUIC stream for the specified connection.
-func (p *dnsOverQUIC) openStream(conn *quic.Conn) (*quic.Stream, error) {
-	ctx, cancel := p.withDeadline(context.Background())
+func (p *dnsOverQUIC) openStream(ctx context.Context, conn *quic.Conn) (*quic.Stream, error) {
+	ctx, cancel := p.withDeadline(ctx)
 	defer cancel()
 
 	stream, err := conn.OpenStreamSync(ctx)
@@ -350,16 +358,16 @@ func (p *dnsOverQUIC) openStream(conn *quic.Conn) (*quic.Stream, error) {
 }
 
 // openConnection dials a new QUIC connection.
-func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
+func (p *dnsOverQUIC) openConnection(ctx context.Context) (conn *quic.Conn, err error) {
 	dialContext, err := p.getDialer()
 	if err != nil {
 		return nil, fmt.Errorf("bootstrapping %s: %w", p.addr, err)
 	}
 
-	// we're using bootstrapped address instead of what's passed to the function
+	// We're using bootstrapped address instead of what's passed to the function
 	// it does not create an actual connection, but it helps us determine
 	// what IP is actually reachable (when there're v4/v6 addresses).
-	rawConn, err := dialContext(context.Background(), "udp", "")
+	rawConn, err := dialContext(ctx, "udp", "")
 	if err != nil {
 		return nil, fmt.Errorf("dialing raw connection to %s: %w", p.addr, err)
 	}
@@ -367,7 +375,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
 	// It's never actually used.
 	err = rawConn.Close()
 	if err != nil {
-		p.logger.Debug("closing raw connection", "addr", p.addr, slogutil.KeyError, err)
+		p.logger.DebugContext(ctx, "closing raw connection", "addr", p.addr, slogutil.KeyError, err)
 	}
 
 	udpConn, ok := rawConn.(*net.UDPConn)
@@ -377,7 +385,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
 
 	addr := udpConn.RemoteAddr().String()
 
-	ctx, cancel := p.withDeadline(context.Background())
+	ctx, cancel := p.withDeadline(ctx)
 	defer cancel()
 
 	conn, err = quic.DialAddrEarly(ctx, addr, p.tlsConf.Clone(), p.getQUICConfig())
@@ -391,7 +399,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
 // closeConnWithError closes the active connection with error to make sure that
 // new queries were processed in another connection.  We can do that in the case
 // of a fatal error.
-func (p *dnsOverQUIC) closeConnWithError(conn *quic.Conn, err error) {
+func (p *dnsOverQUIC) closeConnWithError(ctx context.Context, conn *quic.Conn, err error) {
 	p.connMu.Lock()
 	defer p.connMu.Unlock()
 
@@ -407,7 +415,7 @@ func (p *dnsOverQUIC) closeConnWithError(conn *quic.Conn, err error) {
 
 	err = conn.CloseWithError(code, "")
 	if err != nil {
-		p.logger.Error("failed to close the conn", slogutil.KeyError, err)
+		p.logger.ErrorContext(ctx, "failed to close the conn", slogutil.KeyError, err)
 	}
 
 	// If the connection that's being closed is cached, reset the cache.
diff --git a/upstream/doq_internal_test.go b/upstream/doq_internal_test.go
index 2add53b2..8338f96f 100644
--- a/upstream/doq_internal_test.go
+++ b/upstream/doq_internal_test.go
@@ -49,7 +49,7 @@ func TestDNSOverQUIC(t *testing.T) {
 
 	// Test that it responds properly
 	for range 10 {
-		checkUpstream(t, u, address)
+		checkUpstream(t, u, address, testTimeout)
 
 		if conn == nil {
 			conn = uq.conn
@@ -63,7 +63,7 @@ func TestDNSOverQUIC(t *testing.T) {
 	_ = conn.CloseWithError(quic.ApplicationErrorCode(0), "")
 
 	// Try to establish it again.
-	checkUpstream(t, u, address)
+	checkUpstream(t, u, address, testTimeout)
 
 	// Make sure that the session has been resumed.
 	require.True(t, lastState.DidResume)
@@ -74,7 +74,7 @@ func TestDNSOverQUIC(t *testing.T) {
 	require.NoError(t, err)
 	testutil.CleanupAndRequireSuccess(t, u.Close)
 
-	checkRaceCondition(u)
+	checkRaceCondition(t, u)
 }
 
 func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
@@ -96,7 +96,7 @@ func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
 	testutil.CleanupAndRequireSuccess(t, u.Close)
 
 	// Test that the upstream works properly.
-	checkUpstream(t, u, address)
+	checkUpstream(t, u, address, testTimeout)
 
 	// Close all active connections.
 	err = srv.closeConns()
@@ -113,17 +113,17 @@ func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
 	wg := sync.WaitGroup{}
 	wg.Add(parallelQueries)
 
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
 	for range 10 {
 		pt := testutil.PanicT{}
-
-		go func(t assert.TestingT) {
+		go func(ctx context.Context, t assert.TestingT) {
 			defer wg.Done()
 
 			req := createTestMessage()
-			_, errExch := u.Exchange(req)
+			_, errExch := u.Exchange(ctx, req)
 
 			assert.NoError(t, errExch)
-		}(pt)
+		}(ctx, pt)
 	}
 
 	wg.Wait()
@@ -160,7 +160,7 @@ func TestDNSOverQUIC_serverRestart(t *testing.T) {
 		)
 		require.NoError(t, err)
 
-		checkUpstream(t, u, upsStr)
+		checkUpstream(t, u, upsStr, testTimeout)
 	})
 	require.False(t, t.Failed())
 	testutil.CleanupAndRequireSuccess(t, u.Close)
@@ -168,17 +168,18 @@ func TestDNSOverQUIC_serverRestart(t *testing.T) {
 	t.Run("second_try", func(t *testing.T) {
 		_ = startDoQServer(t, tlsConf, int(addr.Port()))
 
-		checkUpstream(t, u, upsStr)
+		checkUpstream(t, u, upsStr, testTimeout)
 	})
 	require.False(t, t.Failed())
 
 	t.Run("retry", func(t *testing.T) {
-		_, err := u.Exchange(createTestMessage())
+		ctx := testutil.ContextWithTimeout(t, testTimeout)
+		_, err := u.Exchange(ctx, createTestMessage())
 		require.Error(t, err)
 
 		_ = startDoQServer(t, tlsConf, int(addr.Port()))
 
-		checkUpstream(t, u, upsStr)
+		checkUpstream(t, u, upsStr, testTimeout)
 	})
 }
 
@@ -199,9 +200,10 @@ func TestDNSOverQUIC_0RTT(t *testing.T) {
 
 	uq := testutil.RequireTypeAssert[*dnsOverQUIC](t, u)
 	req := createTestMessage()
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
 
 	// Trigger connection to a QUIC server.
-	resp, err := uq.Exchange(req)
+	resp, err := uq.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, resp)
 
@@ -216,8 +218,9 @@ func TestDNSOverQUIC_0RTT(t *testing.T) {
 		uq.conn = nil
 	}()
 
+	ctx = testutil.ContextWithTimeout(t, testTimeout)
 	// Trigger second connection.
-	resp, err = uq.Exchange(req)
+	resp, err = uq.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, resp)
 
diff --git a/upstream/dot.go b/upstream/dot.go
index a4d7c29e..72e55c08 100644
--- a/upstream/dot.go
+++ b/upstream/dot.go
@@ -49,6 +49,9 @@ type dnsOverTLS struct {
 	// This leads to weak performance for all exchanges coming across such
 	// connections.
 	conns []net.Conn
+
+	// timeout is the timeout for DNS request reads and writes.
+	timeout time.Duration
 }
 
 // newDoT returns the DNS-over-TLS Upstream.
@@ -81,6 +84,7 @@ func newDoT(addr *url.URL, opts *Options) (ups Upstream, err error) {
 		},
 		connsMu: &sync.Mutex{},
 		logger:  opts.Logger,
+		timeout: opts.Timeout,
 	}
 
 	runtime.SetFinalizer(tlsUps, (*dnsOverTLS).Close)
@@ -95,28 +99,33 @@ var _ Upstream = (*dnsOverTLS)(nil)
 func (p *dnsOverTLS) Address() string { return p.addr.String() }
 
 // Exchange implements the [Upstream] interface for *dnsOverTLS.
-func (p *dnsOverTLS) Exchange(req *dns.Msg) (reply *dns.Msg, err error) {
+func (p *dnsOverTLS) Exchange(ctx context.Context, req *dns.Msg) (reply *dns.Msg, err error) {
 	h, err := p.getDialer()
 	if err != nil {
 		return nil, fmt.Errorf("getting conn to %s: %w", p.addr, err)
 	}
 
-	conn, err := p.conn(h)
+	conn, err := p.conn(ctx, h)
 	if err != nil {
 		return nil, fmt.Errorf("getting conn to %s: %w", p.addr, err)
 	}
 
-	reply, err = p.exchangeWithConn(conn, req)
+	reply, err = p.exchangeWithConn(ctx, conn, req)
 	if err != nil {
 		// The pooled connection might have been closed already, see
 		// https://github.com/AdguardTeam/dnsproxy/issues/3.  The following
 		// connection from pool may also be malformed, so dial a new one.
 
 		err = errors.WithDeferred(err, conn.Close())
-		p.logger.Debug("dot got bad conn from pool", "addr", p.addr, slogutil.KeyError, err)
+		p.logger.DebugContext(
+			ctx,
+			"dot got bad conn from pool",
+			"addr", p.addr,
+			slogutil.KeyError, err,
+		)
 
 		// Retry.
-		conn, err = tlsDial(h, p.tlsConf.Clone())
+		conn, err = tlsDial(ctx, h, p.tlsConf.Clone())
 		if err != nil {
 			return nil, fmt.Errorf(
 				"dialing %s: connecting to %s: %w",
@@ -126,7 +135,7 @@ func (p *dnsOverTLS) Exchange(req *dns.Msg) (reply *dns.Msg, err error) {
 			)
 		}
 
-		reply, err = p.exchangeWithConn(conn, req)
+		reply, err = p.exchangeWithConn(ctx, conn, req)
 		if err != nil {
 			return reply, errors.WithDeferred(err, conn.Close())
 		}
@@ -157,11 +166,11 @@ func (p *dnsOverTLS) Close() (err error) {
 
 // conn returns the first available connection from the pool if there is any, or
 // dials a new one otherwise.
-func (p *dnsOverTLS) conn(h bootstrap.DialHandler) (conn net.Conn, err error) {
+func (p *dnsOverTLS) conn(ctx context.Context, h bootstrap.DialHandler) (conn net.Conn, err error) {
 	// Dial a new connection outside the lock, if needed.
 	defer func() {
 		if conn == nil {
-			conn, err = tlsDial(h, p.tlsConf.Clone())
+			conn, err = tlsDial(ctx, h, p.tlsConf.Clone())
 			err = errors.Annotate(err, "connecting to %s: %w", p.tlsConf.ServerName)
 		}
 	}()
@@ -178,14 +187,18 @@ func (p *dnsOverTLS) conn(h bootstrap.DialHandler) (conn net.Conn, err error) {
 
 	err = conn.SetDeadline(time.Now().Add(dialTimeout))
 	if err != nil {
-		p.logger.Debug("dot upstream setting deadline to conn from pool", slogutil.KeyError, err)
+		p.logger.DebugContext(
+			ctx,
+			"dot upstream setting deadline to conn from pool",
+			slogutil.KeyError, err,
+		)
 
 		// If deadLine can't be updated it means that connection was already
 		// closed.
 		return nil, nil
 	}
 
-	p.logger.Debug("dot upstream using existing conn", "raddr", conn.RemoteAddr())
+	p.logger.DebugContext(ctx, "dot upstream using existing conn", "raddr", conn.RemoteAddr())
 
 	return conn, nil
 }
@@ -198,14 +211,26 @@ func (p *dnsOverTLS) putBack(conn net.Conn) {
 }
 
 // exchangeWithConn tries to exchange the query using conn.
-func (p *dnsOverTLS) exchangeWithConn(conn net.Conn, req *dns.Msg) (reply *dns.Msg, err error) {
+func (p *dnsOverTLS) exchangeWithConn(
+	ctx context.Context,
+	conn net.Conn,
+	req *dns.Msg,
+) (reply *dns.Msg, err error) {
 	addr := p.Address()
 
-	logBegin(p.logger, addr, networkTCP, req)
-	defer func() { logFinish(p.logger, addr, networkTCP, err) }()
+	logBegin(ctx, p.logger, addr, networkTCP, req)
+	defer func() { logFinish(ctx, p.logger, addr, networkTCP, err) }()
 
 	dnsConn := dns.Conn{Conn: conn}
 
+	if p.timeout > 0 {
+		deadline := time.Now().Add(p.timeout)
+		err = dnsConn.SetDeadline(deadline)
+		if err != nil {
+			return nil, fmt.Errorf("setting connection deadline: %w", err)
+		}
+	}
+
 	err = dnsConn.WriteMsg(req)
 	if err != nil {
 		return nil, fmt.Errorf("sending request to %s: %w", addr, err)
@@ -223,10 +248,14 @@ func (p *dnsOverTLS) exchangeWithConn(conn net.Conn, req *dns.Msg) (reply *dns.M
 
 // tlsDial is basically the same as tls.DialWithDialer, but we will call our own
 // dialContext function to get connection.
-func tlsDial(dialContext bootstrap.DialHandler, conf *tls.Config) (c *tls.Conn, err error) {
+func tlsDial(
+	ctx context.Context,
+	dialContext bootstrap.DialHandler,
+	conf *tls.Config,
+) (c *tls.Conn, err error) {
 	// We're using bootstrapped address instead of what's passed to the
 	// function.
-	rawConn, err := dialContext(context.Background(), networkTCP, "")
+	rawConn, err := dialContext(ctx, networkTCP, "")
 	if err != nil {
 		return nil, err
 	}
@@ -234,6 +263,8 @@ func tlsDial(dialContext bootstrap.DialHandler, conf *tls.Config) (c *tls.Conn,
 	// We want the timeout to cover the whole process: TCP connection and TLS
 	// handshake dialTimeout will be used as connection deadLine.
 	conn := tls.Client(rawConn, conf)
+
+	// TODO(f.setrakov): Make it configurable.
 	err = conn.SetDeadline(time.Now().Add(dialTimeout))
 	if err != nil {
 		// Must not happen in normal circumstances.
diff --git a/upstream/dot_internal_test.go b/upstream/dot_internal_test.go
index 6bc26c5d..d0872389 100644
--- a/upstream/dot_internal_test.go
+++ b/upstream/dot_internal_test.go
@@ -1,6 +1,7 @@
 package upstream
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -38,7 +39,7 @@ func TestUpstream_dnsOverTLS(t *testing.T) {
 
 	// Test that it responds properly.
 	for range 10 {
-		checkUpstream(t, u, addr)
+		checkUpstream(t, u, addr, testTimeout)
 	}
 }
 
@@ -69,8 +70,9 @@ func TestUpstream_dnsOverTLS_race(t *testing.T) {
 		wg.Go(func() {
 			pt := testutil.PanicT{}
 
+			ctx := testutil.ContextWithTimeout(t, testTimeout)
 			req := createTestMessage()
-			resp, uErr := u.Exchange(req)
+			resp, uErr := u.Exchange(ctx, req)
 			require.NoError(pt, uErr)
 			requireResponse(pt, req, resp)
 		})
@@ -112,7 +114,8 @@ func TestUpstream_dnsOverTLS_poolReconnect(t *testing.T) {
 
 	// Send the first test message.
 	req := createTestMessage()
-	reply, err := u.Exchange(req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	reply, err := u.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, reply)
 
@@ -123,7 +126,8 @@ func TestUpstream_dnsOverTLS_poolReconnect(t *testing.T) {
 
 	// Send the second test message.
 	req = createTestMessage()
-	reply, err = u.Exchange(req)
+	ctx = testutil.ContextWithTimeout(t, testTimeout)
+	reply, err = u.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, reply)
 
@@ -154,7 +158,8 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
 
 	// Send the first test message.
 	req := createTestMessage()
-	response, err := u.Exchange(req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	response, err := u.Exchange(ctx, req)
 	require.NoError(t, err)
 	requireResponse(t, req, response)
 
@@ -167,11 +172,12 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
 	dialHandler, err := p.getDialer()
 	require.NoError(t, err)
 
-	usedConn, err := p.conn(dialHandler)
+	ctx = testutil.ContextWithTimeout(t, testTimeout)
+	usedConn, err := p.conn(ctx, dialHandler)
 	require.NoError(t, err)
 	require.Same(t, usedConn, conn)
 
-	response, err = p.exchangeWithConn(conn, req)
+	response, err = p.exchangeWithConn(ctx, conn, req)
 	require.NoError(t, err)
 	requireResponse(t, req, response)
 
@@ -185,11 +191,12 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
 	require.Len(t, p.conns, 1)
 	conn = p.conns[0]
 
-	usedConn, err = p.conn(dialHandler)
+	ctx = testutil.ContextWithTimeout(t, testTimeout)
+	usedConn, err = p.conn(ctx, dialHandler)
 	require.NoError(t, err)
 	require.Same(t, usedConn, conn)
 
-	response, err = p.exchangeWithConn(usedConn, req)
+	response, err = p.exchangeWithConn(ctx, usedConn, req)
 	require.NoError(t, err)
 	requireResponse(t, req, response)
 
@@ -198,7 +205,7 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
 	require.NoError(t, err)
 
 	// Connection with expired deadLine can't be used.
-	response, err = p.exchangeWithConn(usedConn, req)
+	response, err = p.exchangeWithConn(ctx, usedConn, req)
 	require.Error(t, err)
 	require.Nil(t, response)
 }
@@ -297,7 +304,7 @@ func BenchmarkDoTUpstream(b *testing.B) {
 
 		b.RunParallel(func(p *testing.PB) {
 			for p.Next() {
-				_, _ = u.Exchange(<-reqChan)
+				_, _ = u.Exchange(context.Background(), <-reqChan)
 			}
 		})
 	})
diff --git a/upstream/parallel.go b/upstream/parallel.go
index a2c3c495..367f5a93 100644
--- a/upstream/parallel.go
+++ b/upstream/parallel.go
@@ -1,6 +1,7 @@
 package upstream
 
 import (
+	"context"
 	"fmt"
 	"slices"
 
@@ -21,13 +22,17 @@ const (
 
 // ExchangeParallel returns the first successful response from one of u.  It
 // returns an error if all upstreams failed to exchange the request.
-func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Upstream, err error) {
+func ExchangeParallel(
+	ctx context.Context,
+	ups []Upstream,
+	req *dns.Msg,
+) (reply *dns.Msg, resolved Upstream, err error) {
 	upsNum := len(ups)
 	switch upsNum {
 	case 0:
 		return nil, nil, ErrNoUpstreams
 	case 1:
-		return exchangeSingle(ups[0], req)
+		return exchangeSingle(ctx, ups[0], req)
 	default:
 		// Go on.
 	}
@@ -39,7 +44,7 @@ func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Up
 		//
 		// TODO(s.chzhen):  Consider using buffer pool.
 		copyReq := req.Copy()
-		go exchangeAsync(f, copyReq, resCh)
+		go exchangeAsync(ctx, f, copyReq, resCh)
 	}
 
 	errs := []error{}
@@ -68,10 +73,11 @@ func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Up
 // exchangeSingle returns a successful response and resolver if a DNS lookup was
 // successful.
 func exchangeSingle(
+	ctx context.Context,
 	ups Upstream,
 	req *dns.Msg,
 ) (resp *dns.Msg, resolved Upstream, err error) {
-	resp, err = ups.Exchange(req)
+	resp, err = ups.Exchange(ctx, req)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -91,14 +97,18 @@ type ExchangeAllResult struct {
 
 // ExchangeAll returns the responses from all of u.  It returns an error only if
 // all upstreams failed to exchange the request.
-func ExchangeAll(ups []Upstream, req *dns.Msg) (res []ExchangeAllResult, err error) {
+func ExchangeAll(
+	ctx context.Context,
+	ups []Upstream,
+	req *dns.Msg,
+) (res []ExchangeAllResult, err error) {
 	upsNum := len(ups)
 	switch upsNum {
 	case 0:
 		return nil, ErrNoUpstreams
 	case 1:
 		var reply *dns.Msg
-		reply, err = ups[0].Exchange(req)
+		reply, err = ups[0].Exchange(ctx, req)
 		if err != nil {
 			return nil, err
 		} else if reply == nil {
@@ -122,7 +132,7 @@ func ExchangeAll(ups []Upstream, req *dns.Msg) (res []ExchangeAllResult, err err
 		//
 		// TODO(s.chzhen):  Consider using buffer pool.
 		copyReq := req.Copy()
-		go exchangeAsync(u, copyReq, resCh)
+		go exchangeAsync(ctx, u, copyReq, resCh)
 	}
 
 	// Wait for all exchanges to finish.
@@ -162,8 +172,8 @@ func receiveAsyncResult(resCh chan any) (res *ExchangeAllResult, err error) {
 
 // exchangeAsync tries to resolve DNS request with one upstream and sends the
 // result to respCh.
-func exchangeAsync(u Upstream, req *dns.Msg, resCh chan any) {
-	reply, err := u.Exchange(req)
+func exchangeAsync(ctx context.Context, u Upstream, req *dns.Msg, resCh chan any) {
+	reply, err := u.Exchange(ctx, req)
 	if err != nil {
 		resCh <- err
 	} else {
diff --git a/upstream/parallel_internal_test.go b/upstream/parallel_internal_test.go
index ccd64254..90cb6c1d 100644
--- a/upstream/parallel_internal_test.go
+++ b/upstream/parallel_internal_test.go
@@ -1,6 +1,7 @@
 package upstream
 
 import (
+	"context"
 	"fmt"
 	"net/netip"
 	"testing"
@@ -30,7 +31,8 @@ func TestExchangeParallel(t *testing.T) {
 
 	req := createTestMessage()
 	start := time.Now()
-	resp, u, err := ExchangeParallel(upstreams, req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	resp, u, err := ExchangeParallel(ctx, upstreams, req)
 	if err != nil {
 		t.Fatalf("no response from test upstreams: %s", err)
 	}
@@ -53,7 +55,8 @@ func TestExchangeParallelEmpty(t *testing.T) {
 	}
 
 	req := createTestMessage()
-	resp, up, err := ExchangeParallel(ups, req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	resp, up, err := ExchangeParallel(ctx, ups, req)
 	require.Error(t, err)
 
 	assert.Nil(t, resp)
@@ -79,7 +82,7 @@ type testUpstream struct {
 var _ Upstream = (*testUpstream)(nil)
 
 // Exchange implements the [Upstream] interface for *testUpstream.
-func (u *testUpstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (u *testUpstream) Exchange(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	if u.sleep != 0 {
 		time.Sleep(u.sleep)
 	}
@@ -130,7 +133,8 @@ func TestExchangeAll(t *testing.T) {
 	}}
 
 	req := createHostTestMessage("test.org")
-	res, err := ExchangeAll(ups, req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	res, err := ExchangeAll(ctx, ups, req)
 	require.NoError(t, err)
 	require.Len(t, res, 2)
 
diff --git a/upstream/plain.go b/upstream/plain.go
index a7036e0f..cf11544e 100644
--- a/upstream/plain.go
+++ b/upstream/plain.go
@@ -86,11 +86,21 @@ func (p *plainDNS) Address() string {
 // dialExchange performs a DNS exchange with the specified dial handler.
 // network must be either [networkUDP] or [networkTCP].
 func (p *plainDNS) dialExchange(
+	ctx context.Context,
 	network network,
 	dial bootstrap.DialHandler,
 	req *dns.Msg,
 ) (resp *dns.Msg, err error) {
+	var cancel context.CancelFunc
+	if p.timeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, p.timeout)
+		defer cancel()
+	}
+
 	addr := p.Address()
+
+	// NOTE: It is important to set a client timeout to prevent the use of the
+	// default timeout settings.
 	client := &dns.Client{Timeout: p.timeout}
 
 	conn := &dns.Conn{}
@@ -101,25 +111,34 @@ func (p *plainDNS) dialExchange(
 		}
 	}()
 
-	logBegin(p.logger, addr, network, upstreamReq)
-	defer func() { logFinish(p.logger, addr, network, err) }()
+	logBegin(ctx, p.logger, addr, network, upstreamReq)
+	defer func() { logFinish(ctx, p.logger, addr, network, err) }()
 
-	ctx := context.Background()
-	conn.Conn, err = dial(ctx, network, "")
+	// NOTE: Discard context timeout to prevent deadline sharing between dial
+	// calls.
+	//
+	// TODO(f.setrakov): Consider supporting context timeout here.
+	conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
 	if err != nil {
 		return nil, fmt.Errorf("dialing %s over %s: %w", p.addr.Host, network, err)
 	}
 	defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)
 
-	resp, _, err = client.ExchangeWithConn(upstreamReq, conn)
+	resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
 	if isExpectedConnErr(err) {
-		conn.Conn, err = dial(ctx, network, "")
+		if p.timeout > 0 {
+			ctx = context.WithoutCancel(ctx)
+			ctx, cancel = context.WithTimeout(ctx, p.timeout)
+			defer cancel()
+		}
+
+		conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
 		if err != nil {
 			return nil, fmt.Errorf("dialing %s over %s again: %w", p.addr.Host, network, err)
 		}
 		defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)
 
-		resp, _, err = client.ExchangeWithConn(upstreamReq, conn)
+		resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
 	}
 
 	if err != nil {
@@ -159,7 +178,7 @@ func isExpectedConnErr(err error) (is bool) {
 }
 
 // Exchange implements the [Upstream] interface for *plainDNS.
-func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
+func (p *plainDNS) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 	dial, err := p.getDialer()
 	if err != nil {
 		// Don't wrap the error since it's informative enough as is.
@@ -168,7 +187,7 @@ func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
 
 	addr := p.Address()
 
-	resp, err = p.dialExchange(p.net, dial, req)
+	resp, err = p.dialExchange(ctx, p.net, dial, req)
 	if p.net != networkUDP {
 		// The network is already TCP.
 		return resp, err
@@ -181,22 +200,24 @@ func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
 
 	if errors.Is(err, errQuestion) {
 		// The upstream responds with malformed messages, so try TCP.
-		p.logger.Debug(
+		p.logger.DebugContext(
+			ctx,
 			"plain response is malformed, using tcp",
 			"addr", addr,
 			slogutil.KeyError, err,
 		)
 
-		return p.dialExchange(networkTCP, dial, req)
+		return p.dialExchange(ctx, networkTCP, dial, req)
 	} else if resp.Truncated {
 		// Fallback to TCP on truncated responses.
-		p.logger.Debug(
+		p.logger.DebugContext(
+			ctx,
 			"plain response is truncated, using tcp",
 			"question", &req.Question[0],
 			"addr", addr,
 		)
 
-		return p.dialExchange(networkTCP, dial, req)
+		return p.dialExchange(ctx, networkTCP, dial, req)
 	}
 
 	// There is either no error or the error isn't related to the received
diff --git a/upstream/plain_internal_test.go b/upstream/plain_internal_test.go
index 35967a40..aca2d24b 100644
--- a/upstream/plain_internal_test.go
+++ b/upstream/plain_internal_test.go
@@ -34,7 +34,7 @@ func TestUpstream_plainDNS(t *testing.T) {
 	testutil.CleanupAndRequireSuccess(t, u.Close)
 
 	for range 10 {
-		checkUpstream(t, u, addr)
+		checkUpstream(t, u, addr, testTimeout)
 	}
 }
 
@@ -57,7 +57,8 @@ func TestUpstream_plainDNS_badID(t *testing.T) {
 	require.NoError(t, err)
 	testutil.CleanupAndRequireSuccess(t, u.Close)
 
-	resp, err := u.Exchange(req)
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+	resp, err := u.Exchange(ctx, req)
 
 	var netErr net.Error
 	require.ErrorAs(t, err, &netErr)
@@ -132,7 +133,8 @@ func TestUpstream_plainDNS_fallbackToTCP(t *testing.T) {
 			require.NoError(t, err)
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
-			resp, err := u.Exchange(req)
+			ctx := testutil.ContextWithTimeout(t, testTimeout)
+			resp, err := u.Exchange(ctx, req)
 			require.NoError(t, err)
 			requireResponse(t, req, resp)
 
diff --git a/upstream/resolver.go b/upstream/resolver.go
index 2dbe8f88..d9f9f9ed 100644
--- a/upstream/resolver.go
+++ b/upstream/resolver.go
@@ -154,16 +154,14 @@ type ipResult struct {
 // lookupNetIP performs a DNS lookup of host and returns the result.  network
 // must be either [bootstrap.NetworkIP4], [bootstrap.NetworkIP6], or
 // [bootstrap.NetworkIP].  host must be in a lower-case FQDN form.
-//
-// TODO(e.burkov):  Use context.
 func (r *UpstreamResolver) lookupNetIP(
-	_ context.Context,
+	ctx context.Context,
 	network bootstrap.Network,
 	host string,
 ) (result *ipResult, err error) {
 	switch network {
 	case bootstrap.NetworkIP4, bootstrap.NetworkIP6:
-		return r.request(host, network)
+		return r.request(ctx, host, network)
 	case bootstrap.NetworkIP:
 		// Go on.
 	default:
@@ -171,8 +169,8 @@ func (r *UpstreamResolver) lookupNetIP(
 	}
 
 	resCh := make(chan any, 2)
-	go r.resolveAsync(resCh, host, bootstrap.NetworkIP4)
-	go r.resolveAsync(resCh, host, bootstrap.NetworkIP6)
+	go r.resolveAsync(ctx, resCh, host, bootstrap.NetworkIP4)
+	go r.resolveAsync(ctx, resCh, host, bootstrap.NetworkIP6)
 
 	var errs []error
 	result = &ipResult{}
@@ -199,7 +197,11 @@ func (r *UpstreamResolver) lookupNetIP(
 //
 // TODO(e.burkov):  Consider NS and Extra sections when setting TTL.  Check out
 // what RFCs say about it.
-func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipResult, err error) {
+func (r *UpstreamResolver) request(
+	ctx context.Context,
+	host string,
+	n bootstrap.Network,
+) (res *ipResult, err error) {
 	var qtype uint16
 	switch n {
 	case bootstrap.NetworkIP4:
@@ -224,7 +226,7 @@ func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipRes
 
 	// As per [Upstream.Exchange] documentation, the response is always returned
 	// if no error occurred.
-	resp, err := r.Exchange(req)
+	resp, err := r.Exchange(ctx, req)
 	if err != nil {
 		return res, err
 	}
@@ -251,8 +253,12 @@ func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipRes
 
 // resolveAsync performs a single DNS lookup and sends the result to ch.  It's
 // intended to be used as a goroutine.
-func (r *UpstreamResolver) resolveAsync(resCh chan<- any, host, network string) {
-	res, err := r.request(host, network)
+func (r *UpstreamResolver) resolveAsync(
+	ctx context.Context,
+	resCh chan<- any,
+	host, network string,
+) {
+	res, err := r.request(ctx, host, network)
 	if err != nil {
 		resCh <- err
 	} else {
diff --git a/upstream/resolver_test.go b/upstream/resolver_test.go
index 4d13139a..c485a6c5 100644
--- a/upstream/resolver_test.go
+++ b/upstream/resolver_test.go
@@ -20,7 +20,7 @@ func TestNewUpstreamResolver(t *testing.T) {
 	ups := &dnsproxytest.Upstream{
 		OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
-		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+		OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 			resp = (&dns.Msg{}).SetReply(req)
 			resp.Answer = []dns.RR{&dns.A{
 				Hdr: dns.RR_Header{
@@ -133,7 +133,7 @@ func TestCachingResolver_staleness(t *testing.T) {
 		fqdn = "test.fully.qualified.name."
 	)
 
-	onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
+	onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
 		resp = (&dns.Msg{}).SetReply(req)
 
 		hdr := dns.RR_Header{
diff --git a/upstream/upstream.go b/upstream/upstream.go
index 97e6b6db..a473058e 100644
--- a/upstream/upstream.go
+++ b/upstream/upstream.go
@@ -36,7 +36,7 @@ type Upstream interface {
 	// must not modify req as well as the caller must not modify it until the
 	// method returns.  It shouldn't be called after closing.  req must not be
 	// nil.
-	Exchange(req *dns.Msg) (resp *dns.Msg, err error)
+	Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
 
 	// Address returns the human-readable address of the upstream DNS resolver.
 	// It may differ from what was passed to [AddressToUpstream].
@@ -330,7 +330,7 @@ func addPort(u *url.URL, port uint16) {
 // logBegin logs the start of DNS request resolution.  It should be called right
 // before dialing the connection to the upstream.  n is the [network] that will
 // be used to send the request.
-func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
+func logBegin(ctx context.Context, l *slog.Logger, addr string, n network, req *dns.Msg) {
 	var qtype dns.Type
 	var qname string
 	if len(req.Question) != 0 {
@@ -339,7 +339,7 @@ func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
 	}
 
 	l.DebugContext(
-		context.TODO(),
+		ctx,
 		"sending request",
 		"addr", addr,
 		"proto", n,
@@ -352,7 +352,7 @@ func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
 // logFinish logs the end of DNS request resolution.  It should be called right
 // after receiving the response from the upstream or the failing action.  n is
 // the [network] that was used to send the request.
-func logFinish(l *slog.Logger, addr string, n network, err error) {
+func logFinish(ctx context.Context, l *slog.Logger, addr string, n network, err error) {
 	lvl := slog.LevelDebug
 	status := "ok"
 
@@ -364,7 +364,7 @@ func logFinish(l *slog.Logger, addr string, n network, err error) {
 		}
 	}
 
-	l.Log(context.TODO(), lvl, "response received", "addr", addr, "proto", n, "status", status)
+	l.Log(ctx, lvl, "response received", "addr", addr, "proto", n, "status", status)
 }
 
 // isTimeout returns true if err is a timeout error.
diff --git a/upstream/upstream_internal_test.go b/upstream/upstream_internal_test.go
index 38d5fb56..1d2a2b9e 100644
--- a/upstream/upstream_internal_test.go
+++ b/upstream/upstream_internal_test.go
@@ -90,7 +90,9 @@ func TestUpstream_bootstrapTimeout(t *testing.T) {
 			req := createTestMessage()
 
 			start := time.Now()
-			_, rErr := u.Exchange(req)
+
+			ctx := testutil.ContextWithTimeout(t, testTimeout)
+			_, rErr := u.Exchange(ctx, req)
 			elapsed := time.Since(start)
 
 			// Require an error, since the bootstrap server cannot work.
@@ -229,7 +231,7 @@ func TestUpstreams(t *testing.T) {
 			require.NoErrorf(t, upsErr, "failed to generate upstream from address %s", test.address)
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
-			checkUpstream(t, u, test.address)
+			checkUpstream(t, u, test.address, testTimeout)
 		})
 	}
 }
@@ -399,7 +401,7 @@ func TestUpstreamDoTBootstrap(t *testing.T) {
 			require.NoErrorf(t, err, "failed to generate upstream from address %s", tc.address)
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
-			checkUpstream(t, u, tc.address)
+			checkUpstream(t, u, tc.address, testTimeout)
 		})
 	}
 }
@@ -431,6 +433,7 @@ func TestUpstreamsInvalidBootstrap(t *testing.T) {
 	}}
 
 	l := testLogger
+	const upstreamTimeout = 5 * time.Second
 
 	for _, tc := range upstreams {
 		t.Run(tc.address, func(t *testing.T) {
@@ -450,12 +453,12 @@ func TestUpstreamsInvalidBootstrap(t *testing.T) {
 			u, err := AddressToUpstream(tc.address, &Options{
 				Logger:    l,
 				Bootstrap: rslv,
-				Timeout:   testTimeout,
+				Timeout:   upstreamTimeout,
 			})
 			require.NoErrorf(t, err, "failed to generate upstream from address %s", tc.address)
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
-			checkUpstream(t, u, tc.address)
+			checkUpstream(t, u, tc.address, upstreamTimeout)
 		})
 	}
 
@@ -527,7 +530,7 @@ func TestAddressToUpstream_StaticResolver(t *testing.T) {
 			testutil.CleanupAndRequireSuccess(t, u.Close)
 
 			assert.NotPanics(t, func() {
-				checkUpstream(t, u, tc.address)
+				checkUpstream(t, u, tc.address, testTimeout)
 			})
 		})
 	}
@@ -594,11 +597,13 @@ func TestAddPort(t *testing.T) {
 }
 
 // checkUpstream sends a test message to the upstream and checks the result.
-func checkUpstream(tb testing.TB, u Upstream, addr string) {
+func checkUpstream(tb testing.TB, u Upstream, addr string, timeout time.Duration) {
 	tb.Helper()
 
 	req := createTestMessage()
-	reply, err := u.Exchange(req)
+
+	ctx := testutil.ContextWithTimeout(tb, timeout)
+	reply, err := u.Exchange(ctx, req)
 	require.NoErrorf(tb, err, "couldn't talk to upstream %s", addr)
 
 	requireResponse(tb, req, reply)
@@ -606,7 +611,9 @@ func checkUpstream(tb testing.TB, u Upstream, addr string) {
 
 // checkRaceCondition runs several goroutines in parallel and each of them calls
 // checkUpstream several times.
-func checkRaceCondition(u Upstream) {
+func checkRaceCondition(tb testing.TB, u Upstream) {
+	tb.Helper()
+
 	wg := sync.WaitGroup{}
 
 	// The number of requests to run in every goroutine.
@@ -619,7 +626,8 @@ func checkRaceCondition(u Upstream) {
 		for range reqCount {
 			req := createTestMessage()
 			// Ignore exchange errors here, the point is to check for races.
-			_, _ = u.Exchange(req)
+			ctx := testutil.ContextWithTimeout(tb, testTimeout)
+			_, _ = u.Exchange(ctx, req)
 		}
 	}