tracker/udp: Improve context handling

go.mod

@@ -45,7 +45,8 @@ require (
 	github.com/pion/webrtc/v3 v3.1.42
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.12.2
-	github.com/stretchr/testify v1.8.1
+	github.com/protolambda/ctxlock v0.1.0
+	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/btree v1.6.0
 	go.etcd.io/bbolt v1.3.6
 	go.opentelemetry.io/otel v1.11.1

go.sum

@@ -485,6 +485,8 @@ github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O
 github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
 github.com/prometheus/procfs v0.7.3 h1:4jVXhlkAyzOScmCkXBTOLRLTz8EeU+eyjrwB/EPq0VU=
 github.com/prometheus/procfs v0.7.3/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
+github.com/protolambda/ctxlock v0.1.0 h1:rCUY3+vRdcdZXqT07iXgyr744J2DU2LCBIXowYAjBCE=
+github.com/protolambda/ctxlock v0.1.0/go.mod h1:vefhX6rIZH8rsg5ZpOJfEDYQOppZi19SfPiGOFrNnwM=
 github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
@@ -529,8 +531,9 @@ github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/tidwall/btree v1.6.0 h1:LDZfKfQIBHGHWSwckhXI0RPSXzlo+KYdjK7FWSqOzzg=
 github.com/tidwall/btree v1.6.0/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY=
 github.com/tinylib/msgp v1.0.2/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE=

tracker/udp/client.go

@@ -5,9 +5,9 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
+	"github.com/protolambda/ctxlock"
 	"io"
 	"net"
-	"sync"
 	"time"
 
 	"github.com/anacrolix/dht/v2/krpc"
@@ -16,7 +16,7 @@ import (
 // Client interacts with UDP trackers via its Writer and Dispatcher. It has no knowledge of
 // connection specifics.
 type Client struct {
-	mu           sync.Mutex
+	mu           ctxlock.Lock
 	connId       ConnectionId
 	connIdIssued time.Time
 
@@ -145,7 +145,10 @@ func (cl *Client) writeRequest(
 		// written before allowing the connection ID to change again. This is to ensure the server
 		// doesn't assign us another ID before we've sent this request. Note that this doesn't allow
 		// for us to return if the context is cancelled while we wait to obtain a new ID.
-		cl.mu.Lock()
+		err = cl.mu.LockCtx(ctx)
+		if err != nil {
+			return fmt.Errorf("locking connection id: %w", err)
+		}
 		defer cl.mu.Unlock()
 		connId, err = cl.connIdForRequest(ctx, action)
 		if err != nil {
@@ -228,7 +231,7 @@ func (cl *Client) request(
 	case err = <-writeErr:
 		err = fmt.Errorf("write error: %w", err)
 	case <-ctx.Done():
-		err = ctx.Err()
+		err = context.Cause(ctx)
 	}
 	return
 }