Repository: decred/dcrd
Pull request #1729 (2019-04-30 17:10:24)

connmgr: start v3, add context, use context


Patch:
diff --git a/config.go b/config.go
index fa662461..c9368381 100644
--- a/config.go
+++ b/config.go
@@ -6,6 +6,7 @@
 package main
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/base64"
 	"errors"
@@ -173,8 +174,8 @@ type config struct {
 	AltDNSNames          []string      `long:"altdnsnames" description:"Specify additional DNS names to use when generating the RPC server certificate" env:"DCRD_ALT_DNSNAMES" env-delim:","`
 	onionlookup          func(string) ([]net.IP, error)
 	lookup               func(string) ([]net.IP, error)
-	oniondial            func(string, string) (net.Conn, error)
-	dial                 func(string, string) (net.Conn, error)
+	oniondial            func(context.Context, string, string) (net.Conn, error)
+	dial                 func(context.Context, string, string) (net.Conn, error)
 	miningAddrs          []dcrutil.Address
 	minRelayTxFee        dcrutil.Amount
 	whitelists           []*net.IPNet
@@ -1075,7 +1076,8 @@ func loadConfig() (*config, []string, error) {
 	// specified, the dial function is set to the proxy specific dial
 	// function and the lookup is set to use tor (unless --noonion is
 	// specified in which case the system DNS resolver is used).
-	cfg.dial = net.Dial
+	var d net.Dialer
+	cfg.dial = d.DialContext
 	cfg.lookup = net.LookupIP
 	if cfg.Proxy != "" {
 		_, _, err := net.SplitHostPort(cfg.Proxy)
@@ -1099,7 +1101,7 @@ func loadConfig() (*config, []string, error) {
 			Password:     cfg.ProxyPass,
 			TorIsolation: cfg.TorIsolation,
 		}
-		cfg.dial = proxy.Dial
+		cfg.dial = proxy.DialContext
 		if !cfg.NoOnion {
 			cfg.lookup = func(host string) ([]net.IP, error) {
 				return connmgr.TorLookupIP(host, cfg.Proxy)
@@ -1132,14 +1134,14 @@ func loadConfig() (*config, []string, error) {
 				"credentials ")
 		}
 
-		cfg.oniondial = func(a, b string) (net.Conn, error) {
+		cfg.oniondial = func(ctx context.Context, a, b string) (net.Conn, error) {
 			proxy := &socks.Proxy{
 				Addr:         cfg.OnionProxy,
 				Username:     cfg.OnionProxyUser,
 				Password:     cfg.OnionProxyPass,
 				TorIsolation: cfg.TorIsolation,
 			}
-			return proxy.Dial(a, b)
+			return proxy.DialContext(ctx, a, b)
 		}
 		cfg.onionlookup = func(host string) ([]net.IP, error) {
 			return connmgr.TorLookupIP(host, cfg.OnionProxy)
@@ -1152,7 +1154,7 @@ func loadConfig() (*config, []string, error) {
 	// Specifying --noonion means the onion address dial and DNS resolution
 	// (lookup) functions result in an error.
 	if cfg.NoOnion {
-		cfg.oniondial = func(a, b string) (net.Conn, error) {
+		cfg.oniondial = func(ctx context.Context, a, b string) (net.Conn, error) {
 			return nil, errors.New("tor has been disabled")
 		}
 		cfg.onionlookup = func(a string) ([]net.IP, error) {
@@ -1190,11 +1192,11 @@ func loadConfig() (*config, []string, error) {
 // example, .onion addresses will be dialed using the onion specific proxy if
 // one was specified, but will otherwise use the normal dial function (which
 // could itself use a proxy or not).
-func dcrdDial(network, addr string) (net.Conn, error) {
+func dcrdDial(ctx context.Context, network, addr string) (net.Conn, error) {
 	if strings.Contains(addr, ".onion:") {
-		return cfg.oniondial(network, addr)
+		return cfg.oniondial(ctx, network, addr)
 	}
-	return cfg.dial(network, addr)
+	return cfg.dial(ctx, network, addr)
 }
 
 // dcrdLookup returns the correct DNS lookup function to use depending on the
diff --git a/connmgr/connmanager.go b/connmgr/connmanager.go
index 7ce81a90..b9084ec4 100644
--- a/connmgr/connmanager.go
+++ b/connmgr/connmanager.go
@@ -6,6 +6,7 @@
 package connmgr
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -140,12 +141,12 @@ type Config struct {
 
 	// Dial connects to the address on the named network. Either Dial or
 	// DialAddr need to be specified (but not both).
-	Dial func(network, addr string) (net.Conn, error)
+	Dial func(ctx context.Context, network, addr string) (net.Conn, error)
 
 	// DialAddr is an alternative to Dial which receives a full net.Addr instead
 	// of just the protocol family and address. Either DialAddr or Dial need
 	// to be specified (but not both).
-	DialAddr func(net.Addr) (net.Conn, error)
+	DialAddr func(context.Context, net.Addr) (net.Conn, error)
 }
 
 // registerPending is used to register a pending connection attempt. By
@@ -206,7 +207,7 @@ func (cm *ConnManager) handleFailedConn(c *ConnReq) {
 		}
 		log.Debugf("Retrying connection to %v in %v", c, d)
 		time.AfterFunc(d, func() {
-			cm.Connect(c)
+			cm.Connect(context.Background(), c)
 		})
 	} else if cm.cfg.GetNewAddress != nil {
 		cm.failedAttempts++
@@ -215,10 +216,10 @@ func (cm *ConnManager) handleFailedConn(c *ConnReq) {
 				"-- retrying connection in: %v", maxFailedAttempts,
 				cm.cfg.RetryDuration)
 			time.AfterFunc(cm.cfg.RetryDuration, func() {
-				cm.NewConnReq()
+				cm.newConnReq()
 			})
 		} else {
-			go cm.NewConnReq()
+			go cm.newConnReq()
 		}
 	}
 }
@@ -357,9 +358,9 @@ out:
 	log.Trace("Connection handler done")
 }
 
-// NewConnReq creates a new connection request and connects to the
+// newConnReq creates a new connection request and connects to the
 // corresponding address.
-func (cm *ConnManager) NewConnReq() {
+func (cm *ConnManager) newConnReq() {
 	if atomic.LoadInt32(&cm.stop) != 0 {
 		return
 	}
@@ -400,15 +401,18 @@ func (cm *ConnManager) NewConnReq() {
 
 	c.Addr = addr
 
-	cm.Connect(c)
+	cm.Connect(context.Background(), c)
 }
 
 // Connect assigns an id and dials a connection to the address of the
 // connection request.
-func (cm *ConnManager) Connect(c *ConnReq) {
+func (cm *ConnManager) Connect(ctx context.Context, c *ConnReq) {
 	if atomic.LoadInt32(&cm.stop) != 0 {
 		return
 	}
+	if ctx.Err() != nil {
+		return
+	}
 
 	// During the time we wait for retry there is a chance that this
 	// connection was already cancelled.
@@ -445,9 +449,9 @@ func (cm *ConnManager) Connect(c *ConnReq) {
 	var conn net.Conn
 	var err error
 	if cm.cfg.Dial != nil {
-		conn, err = cm.cfg.Dial(c.Addr.Network(), c.Addr.String())
+		conn, err = cm.cfg.Dial(ctx, c.Addr.Network(), c.Addr.String())
 	} else {
-		conn, err = cm.cfg.DialAddr(c.Addr)
+		conn, err = cm.cfg.DialAddr(ctx, c.Addr)
 	}
 	if err != nil {
 		select {
@@ -534,7 +538,7 @@ func (cm *ConnManager) Start() {
 	}
 
 	for i := atomic.LoadUint64(&cm.connReqCount); i < uint64(cm.cfg.TargetOutbound); i++ {
-		go cm.NewConnReq()
+		go cm.newConnReq()
 	}
 }
 
diff --git a/server.go b/server.go
index f1b5fb2f..b1c63d5c 100644
--- a/server.go
+++ b/server.go
@@ -1890,10 +1890,11 @@ func (s *server) handleQuery(state *peerState, querymsg interface{}) {
 		}
 
 		// TODO: if too many, nuke a non-perm peer.
-		go s.connManager.Connect(&connmgr.ConnReq{
-			Addr:      netAddr,
-			Permanent: msg.permanent,
-		})
+		go s.connManager.Connect(context.Background(),
+			&connmgr.ConnReq{
+				Addr:      netAddr,
+				Permanent: msg.permanent,
+			})
 		msg.reply <- nil
 	case removeNodeMsg:
 		found := disconnectPeer(state.persistentPeers, msg.cmp, func(sp *serverPeer) {
@@ -3134,10 +3135,11 @@ func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Param
 			return nil, err
 		}
 
-		go s.connManager.Connect(&connmgr.ConnReq{
-			Addr:      tcpAddr,
-			Permanent: true,
-		})
+		go s.connManager.Connect(context.Background(),
+			&connmgr.ConnReq{
+				Addr:      tcpAddr,
+				Permanent: true,
+			})
 	}
 
 	if !cfg.DisableRPC {


Test patch:
diff --git a/connmgr/connmanager_test.go b/connmgr/connmanager_test.go
index 5a4ca3f3..0107784c 100644
--- a/connmgr/connmanager_test.go
+++ b/connmgr/connmanager_test.go
@@ -6,6 +6,7 @@
 package connmgr
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -62,7 +63,7 @@ func (c mockConn) SetWriteDeadline(t time.Time) error { return nil }
 
 // mockDialer mocks the net.Dial interface by returning a mock connection to
 // the given address.
-func mockDialer(network, addr string) (net.Conn, error) {
+func mockDialer(ctx context.Context, network, addr string) (net.Conn, error) {
 	r, w := io.Pipe()
 	c := &mockConn{rAddr: &mockAddr{network, addr}}
 	c.Reader = r
@@ -72,7 +73,7 @@ func mockDialer(network, addr string) (net.Conn, error) {
 
 // mockDialer mocks the net.Dial interface by returning a mock connection to
 // the given address.
-func mockDialerAddr(addr net.Addr) (net.Conn, error) {
+func mockDialerAddr(ctx context.Context, addr net.Addr) (net.Conn, error) {
 	r, w := io.Pipe()
 	c := &mockConn{rAddr: addr}
 	c.Reader = r
@@ -146,17 +147,22 @@ func TestStartStop(t *testing.T) {
 		},
 		Permanent: true,
 	}
-	cmgr.Connect(cr)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	cmgr.Connect(ctx, cr)
 	if cr.ID() != 0 {
 		t.Fatalf("start/stop: got id: %v, want: 0", cr.ID())
 	}
 	cmgr.Disconnect(gotConnReq.ID())
 	cmgr.Remove(gotConnReq.ID())
-	select {
-	case <-disconnected:
-		t.Fatalf("start/stop: unexpected disconnection")
-	case <-time.Tick(10 * time.Millisecond):
-		break
+out:
+	for {
+		select {
+		case <-disconnected:
+			t.Fatalf("start/stop: unexpected disconnection")
+		case <-ctx.Done():
+			break out
+		}
 	}
 }
 
@@ -184,7 +190,9 @@ func TestConnectMode(t *testing.T) {
 		Permanent: true,
 	}
 	cmgr.Start()
-	cmgr.Connect(cr)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
+	defer cancel()
+	cmgr.Connect(ctx, cr)
 	gotConnReq := <-connected
 	wantID := cr.ID()
 	gotID := gotConnReq.ID()
@@ -199,7 +207,7 @@ func TestConnectMode(t *testing.T) {
 	select {
 	case c := <-connected:
 		t.Fatalf("connect mode: got unexpected connection - %v", c.Addr)
-	case <-time.After(time.Millisecond):
+	case <-ctx.Done():
 		break
 	}
 	cmgr.Stop()
@@ -317,7 +325,7 @@ func TestRetryPermanent(t *testing.T) {
 		},
 		Permanent: true,
 	}
-	go cmgr.Connect(cr)
+	go cmgr.Connect(context.Background(), cr)
 	cmgr.Start()
 	gotConnReq := <-connected
 	wantID := cr.ID()
@@ -380,10 +388,10 @@ func TestMaxRetryDuration(t *testing.T) {
 	time.AfterFunc(5*time.Millisecond, func() {
 		close(networkUp)
 	})
-	timedDialer := func(network, addr string) (net.Conn, error) {
+	timedDialer := func(ctx context.Context, network, addr string) (net.Conn, error) {
 		select {
 		case <-networkUp:
-			return mockDialer(network, addr)
+			return mockDialer(ctx, network, addr)
 		default:
 			return nil, errors.New("network down")
 		}
@@ -409,15 +417,21 @@ func TestMaxRetryDuration(t *testing.T) {
 		},
 		Permanent: true,
 	}
-	go cmgr.Connect(cr)
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	go cmgr.Connect(ctx, cr)
 	cmgr.Start()
 	// retry in 1ms
 	// retry in 2ms - max retry duration reached
 	// retry in 2ms - timedDialer returns mockDial
-	select {
-	case <-connected:
-	case <-time.Tick(100 * time.Millisecond):
-		t.Fatalf("max retry duration: connection timeout")
+out:
+	for {
+		select {
+		case <-connected:
+			break out
+		case <-ctx.Done():
+			t.Fatalf("max retry duration: connection timeout")
+		}
 	}
 }
 
@@ -425,7 +439,7 @@ func TestMaxRetryDuration(t *testing.T) {
 // failure gracefully.
 func TestNetworkFailure(t *testing.T) {
 	var dials uint32
-	errDialer := func(network, addr string) (net.Conn, error) {
+	errDialer := func(ctx context.Context, network, addr string) (net.Conn, error) {
 		atomic.AddUint32(&dials, 1)
 		return nil, errors.New("network down")
 	}
@@ -464,7 +478,7 @@ func TestNetworkFailure(t *testing.T) {
 // the failure.
 func TestStopFailed(t *testing.T) {
 	done := make(chan struct{}, 1)
-	waitDialer := func(network, addr string) (net.Conn, error) {
+	waitDialer := func(ctx context.Context, network, addr string) (net.Conn, error) {
 		done <- struct{}{}
 		time.Sleep(time.Millisecond)
 		return nil, errors.New("network down")
@@ -490,7 +504,7 @@ func TestStopFailed(t *testing.T) {
 		},
 		Permanent: true,
 	}
-	go cmgr.Connect(cr)
+	go cmgr.Connect(context.Background(), cr)
 	cmgr.Wait()
 }
 
@@ -500,7 +514,7 @@ func TestRemovePendingConnection(t *testing.T) {
 	// Create a ConnMgr instance with an instance of a dialer that'll never
 	// succeed.
 	wait := make(chan struct{})
-	indefiniteDialer := func(addr net.Addr) (net.Conn, error) {
+	indefiniteDialer := func(ctx context.Context, addr net.Addr) (net.Conn, error) {
 		<-wait
 		return nil, fmt.Errorf("error")
 	}
@@ -520,7 +534,7 @@ func TestRemovePendingConnection(t *testing.T) {
 		},
 		Permanent: true,
 	}
-	go cmgr.Connect(cr)
+	go cmgr.Connect(context.Background(), cr)
 
 	time.Sleep(10 * time.Millisecond)
 
@@ -555,10 +569,10 @@ func TestCancelIgnoreDelayedConnection(t *testing.T) {
 	// connect chan is signaled. The dial attempt immediately after that
 	// will succeed in returning a connection.
 	connect := make(chan struct{})
-	failingDialer := func(addr net.Addr) (net.Conn, error) {
+	failingDialer := func(ctx context.Context, addr net.Addr) (net.Conn, error) {
 		select {
 		case <-connect:
-			return mockDialerAddr(addr)
+			return mockDialerAddr(ctx, addr)
 		default:
 		}
 
@@ -586,7 +600,9 @@ func TestCancelIgnoreDelayedConnection(t *testing.T) {
 			Port: 18555,
 		},
 	}
-	cmgr.Connect(cr)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*retryTimeout)
+	defer cancel()
+	cmgr.Connect(ctx, cr)
 
 	// Allow for the first retry timeout to elapse.
 	time.Sleep(2 * retryTimeout)
@@ -619,7 +635,7 @@ func TestCancelIgnoreDelayedConnection(t *testing.T) {
 	select {
 	case <-connected:
 		t.Fatalf("on-connect should not be called for canceled req")
-	case <-time.After(5 * retryTimeout):
+	case <-ctx.Done():
 	}
 
 }