Repository: decred/dcrd
Pull request #2562 (2021-01-20 08:26:26)

connmgr: Add func to iterate conn reqs.
**This is rebased on #2561**.

This adds a new func named `ForEachConnReq` to the connection manager which accepts a closure to call with each connection req known to the connection manager, including pending ones.

A test is also added to ensure proper functionality.


Patch:
diff --git a/connmgr/connmanager.go b/connmgr/connmanager.go
index 7efbcb14..4d801a18 100644
--- a/connmgr/connmanager.go
+++ b/connmgr/connmanager.go
@@ -1,5 +1,5 @@
 // Copyright (c) 2016 The btcsuite developers
-// Copyright (c) 2017-2020 The Decred developers
+// Copyright (c) 2017-2021 The Decred developers
 // Use of this source code is governed by an ISC
 // license that can be found in the LICENSE file.
 
@@ -200,6 +200,13 @@ type handleCancelPending struct {
 	done chan error
 }
 
+// handleForEachConnReq is used to iterate all known connection requests to
+// include pending ones.
+type handleForEachConnReq struct {
+	f    func(c *ConnReq) error
+	done chan error
+}
+
 // ConnManager provides a manager to handle network connections.
 type ConnManager struct {
 	// The following variables must only be used atomically.
@@ -419,6 +426,26 @@ out:
 				} else {
 					msg.done <- fmt.Errorf("no pending connection to %v", msg.addr)
 				}
+
+			case handleForEachConnReq:
+				var err error
+				for _, connReq := range pending {
+					err = msg.f(connReq)
+					if err != nil {
+						break
+					}
+				}
+				if err != nil {
+					msg.done <- err
+					continue
+				}
+				for _, connReq := range conns {
+					err = msg.f(connReq)
+					if err != nil {
+						break
+					}
+				}
+				msg.done <- err
 			}
 
 		case <-ctx.Done():
@@ -598,6 +625,30 @@ func (cm *ConnManager) CancelPending(addr net.Addr) error {
 	}
 }
 
+// ForEachConnReq calls the provided function with each connection request known
+// to the connection manager, including pending requests.  Returning an error
+// from the provided function will stop the iteration early and return said
+// error from this function.
+//
+// This function is safe for concurrent access.
+//
+// NOTE: This must not call any other connection manager methods during
+// iteration or it will result in a deadlock.
+func (cm *ConnManager) ForEachConnReq(f func(c *ConnReq) error) error {
+	done := make(chan error, 1)
+	select {
+	case cm.requests <- handleForEachConnReq{f, done}:
+	case <-cm.quit:
+	}
+
+	select {
+	case err := <-done:
+		return err
+	case <-cm.quit:
+		return fmt.Errorf("connection manager stopped")
+	}
+}
+
 // listenHandler accepts incoming connections on a given listener.  It must be
 // run as a goroutine.
 func (cm *ConnManager) listenHandler(ctx context.Context, listener net.Listener) {
@@ -619,7 +670,7 @@ func (cm *ConnManager) listenHandler(ctx context.Context, listener net.Listener)
 }
 
 // Run starts the connection manager along with its configured listeners and
-// begin connecting to the network.  It blocks until the provided context is
+// begins connecting to the network.  It blocks until the provided context is
 // cancelled.
 func (cm *ConnManager) Run(ctx context.Context) {
 	log.Trace("Starting connection manager")


Test patch:
diff --git a/connmgr/connmanager_test.go b/connmgr/connmanager_test.go
index 79b1d65c..48b26ede 100644
--- a/connmgr/connmanager_test.go
+++ b/connmgr/connmanager_test.go
@@ -1,5 +1,5 @@
 // Copyright (c) 2016 The btcsuite developers
-// Copyright (c) 2019-2020 The Decred developers
+// Copyright (c) 2019-2021 The Decred developers
 // Use of this source code is governed by an ISC
 // license that can be found in the LICENSE file.
 
@@ -571,7 +571,7 @@ func TestShutdownFailedConns(t *testing.T) {
 }
 
 // TestRemovePendingConnection tests that it's possible to cancel a pending
-// connection, removing its internal state from the ConnMgr.
+// connection, removing its internal state from the connection manager.
 func TestRemovePendingConnection(t *testing.T) {
 	// Create a ConnMgr instance with an instance of a dialer that'll never
 	// succeed.
@@ -803,6 +803,113 @@ func TestConnectContext(t *testing.T) {
 	wg.Wait()
 }
 
+// TestForEachConnReq tests the connection request iteration logic work as
+// expected including for normal, permanent, and pending connections.
+func TestForEachConnReq(t *testing.T) {
+	// Create a connection manager instance with a dialer that recognizes a
+	// special address to delay on in order to keep it pending.
+	targetOutbound := uint32(5)
+	connected := make(chan *ConnReq)
+	pending := make(chan struct{})
+	delayDialer := func(ctx context.Context, addr net.Addr) (net.Conn, error) {
+		if addr.String() == "127.0.0.1:18557" {
+			close(pending)
+			time.Sleep(time.Second)
+			return nil, errors.New("error")
+		}
+		return mockDialerAddr(ctx, addr)
+	}
+	cmgr, err := New(&Config{
+		TargetOutbound: targetOutbound,
+		DialAddr:       delayDialer,
+		GetNewAddress: func() (net.Addr, error) {
+			return &net.TCPAddr{
+				IP:   net.ParseIP("127.0.0.1"),
+				Port: 18555,
+			}, nil
+		},
+		OnConnection: func(c *ConnReq, conn net.Conn) {
+			connected <- c
+		},
+	})
+	if err != nil {
+		t.Fatalf("New error: %v", err)
+	}
+	_, shutdown, wg := runConnMgrAsync(context.Background(), cmgr)
+
+	// Wait for the expected number of target outbound conns to be established.
+	allConnected := make(chan struct{})
+	go func() {
+		for i := uint32(0); i < targetOutbound; i++ {
+			<-connected
+		}
+		close(allConnected)
+	}()
+	select {
+	case <-allConnected:
+	case <-time.After(time.Millisecond * 5 * time.Duration(targetOutbound)):
+		t.Fatal("timeout waiting for connections")
+	}
+
+	// Create a permanent connection.
+	cr := &ConnReq{
+		Permanent: true,
+		Addr: &net.TCPAddr{
+			IP:   net.ParseIP("127.0.0.1"),
+			Port: 18556,
+		},
+	}
+	go cmgr.Connect(context.Background(), cr)
+	select {
+	case <-connected:
+	case <-time.After(time.Millisecond * 5):
+		t.Fatal("timeout waiting for permanent connection")
+	}
+
+	// Create a connection that triggers the mock dialer to keep it pending.
+	cr = &ConnReq{
+		Addr: &net.TCPAddr{
+			IP:   net.ParseIP("127.0.0.1"),
+			Port: 18557,
+		},
+	}
+	go cmgr.Connect(context.Background(), cr)
+	select {
+	case <-pending:
+	case <-time.After(time.Millisecond * 5):
+		t.Fatal("timeout waiting for pending connection")
+	}
+
+	// Ensure the expected number of each type of connection exists.
+	var numConnected, numPermanent, numPending uint32
+	cmgr.ForEachConnReq(func(cr *ConnReq) error {
+		numConnected++
+		if cr.State() == ConnPending {
+			numPending++
+		}
+		if cr.Permanent {
+			numPermanent++
+		}
+		return nil
+	})
+	if numConnected != targetOutbound+2 {
+		t.Fatalf("unexpected number of iterated conn reqs -- got %d, want %d",
+			numConnected, targetOutbound+2)
+	}
+	if numPermanent != 1 {
+		t.Fatalf("unexpected number of permanent conn reqs -- got %d, want %d",
+			numPermanent, 1)
+	}
+	if numPending != 1 {
+		t.Fatalf("unexpected number of pending conn reqs -- got %d, want %d",
+			numPending, 1)
+	}
+
+	// Ensure clean shutdown of connection manager.
+	shutdown()
+	wg.Wait()
+}
+
 // mockListener implements the net.Listener interface and is used to test
 // code that deals with net.Listeners without having to actually make any real
 // connections.