mixpool: Refactor activeInEpoch to scale linearly.

mixing/message.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2023-2024 The Decred developers
+// Copyright (c) 2023-2026 The Decred developers
 // Use of this source code is governed by an ISC
 // license that can be found in the LICENSE file.
 
@@ -14,21 +14,28 @@ import (
 // Message is a mixing message.  In addition to implementing wire encoding,
 // these messages are signed by an ephemeral mixing participant identity,
 // declare the previous messages that have been observed by a peer in a mixing
-// session, and include expiry information to increase resilience to replay
-// and denial-of-service attacks.
+// session, and include expiry information to increase resilience to replay and
+// denial-of-service attacks.
 //
 // All mixing messages satisify this interface, however, the pair request
 // message returns nil for some fields that do not apply, as it is the first
 // message in the protocol.
 type Message interface {
 	wire.Message
 
+	// Pub returns the message sender's public key identity.
 	Pub() []byte
+	// Sig returns the message signature.
 	Sig() []byte
 	WriteHash(hash.Hash)
 	Hash() chainhash.Hash
 	WriteSignedData(hash.Hash)
-	PrevMsgs() []chainhash.Hash // PR, FP returns nil
-	Sid() []byte                // PR returns nil
-	GetRun() uint32             // PR returns 0
+	// PrevMsgs returns messages from the previous stage of the mixing session
+	// seen by the peer. For example, PrevMsgs of a KE message will return PR
+	// messages. PR and FP messages return nil.
+	PrevMsgs() []chainhash.Hash
+	// Sid returns the session ID. PR messages return nil.
+	Sid() []byte
+	// GetRun returns the run number. PR messages return 0.
+	GetRun() uint32
 }

mixing/mixpool/mixpool.go

@@ -809,35 +809,43 @@ func (p *Pool) activeInEpoch(epoch uint64) map[[33]byte]activePeer {
 	p.mtx.RLock()
 	defer p.mtx.RUnlock()
 
+	type idPair struct {
+		peerID    idPubKey
+		sessionID [32]byte
+	}
+
+	// Gather all of the KE messages of the epoch, and create a peer/session
+	// idPair for every CT message of the epoch.
 	var epochKEs []*wire.MsgMixKeyExchange
+	idPairs := make(map[idPair]struct{})
 	for _, e := range p.pool {
-		ke, ok := e.msg.(*wire.MsgMixKeyExchange)
-		if !ok {
-			continue
-		}
-		if ke.Epoch != epoch {
-			continue
-		}
-		epochKEs = append(epochKEs, ke)
-	}
-	kes := make([]*wire.MsgMixKeyExchange, 0, len(epochKEs))
-NextKE:
-	for _, ke := range epochKEs {
-		for _, msgHash := range p.messagesByIdentity[ke.Identity] {
-			e := p.pool[msgHash]
-			if e.msgtype == msgtypeCT && e.sid == ke.SessionID {
-				kes = append(kes, ke)
-				continue NextKE
+		switch e.msgtype {
+		case msgtypeKE:
+			ke, ok := e.msg.(*wire.MsgMixKeyExchange)
+			if !ok || ke.Epoch != epoch {
+				continue
 			}
+			epochKEs = append(epochKEs, ke)
+
+		case msgtypeCT:
+			idPairs[idPair{
+				peerID:    *(*idPubKey)(e.msg.Pub()),
+				sessionID: e.sid,
+			}] = struct{}{}
 		}
 	}
 
-	// TODO: sorting the key exchanges by identity and subslicing would be
-	// more memory efficient.
-	activeKEs := make(map[[33]byte][]*wire.MsgMixKeyExchange)
-	for _, ke := range kes {
+	// Select only KE messages of peers that formed a session, as indicated by
+	// the existence of an idPair.
+	activeKEs := make(map[idPubKey][]*wire.MsgMixKeyExchange)
+	for _, ke := range epochKEs {
+		ids := idPair{peerID: ke.Identity, sessionID: ke.SessionID}
+		if _, ok := idPairs[ids]; !ok {
+			continue
+		}
 		activeKEs[ke.Identity] = append(activeKEs[ke.Identity], ke)
 	}
+
 	active := make(map[idPubKey]activePeer)
 	for _, pr := range p.prs {
 		if kes, ok := activeKEs[pr.Identity]; ok {