tracker/udp: Improve context handling

diff --git a/go.mod b/go.mod
index bcbd7caccb6fe63a144beb2c80337bc4c7805678..1e1929b6908eba67a332a09eb6f1c32b92dd44cf 100644 (file)
--- a/go.mod
+++ b/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

diff --git a/go.sum b/go.sum
index ae5edd04677b7523257f58df1ed8dedfc079025a..375d9e5a86344c9dbad0c06fc4fb6b4fca9494eb 100644 (file)
--- a/go.sum
+++ b/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=

diff --git a/tracker/udp/client.go b/tracker/udp/client.go
index 7f6d88a8ee7de42a4e9af11e05fd7626505bffaa..077cc7a40bac83d9259e4d0c65fb417d054d1437 100644 (file)
--- a/tracker/udp/client.go
+++ b/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
 }