p2p/discover: concurrent TALKREQ handling (#27112)

This changes TALKREQ message processing to run the handler on separate goroutine,
instead of running on the main discv5 dispatcher goroutine. It's better this way because
it allows the handler to perform blocking actions.

I'm also adding a new method TalkRequestToID here. The method allows implementing
a request flow where one node A sends TALKREQ to another node B, and node B later
sends a TALKREQ back. With TalkRequestToID, node B does not need the ENR of A to
send its request.

p2p/discover/table_util_test.go

@@ -52,6 +52,7 @@ func newTestTable(t transport) (*Table, *enode.DB) {
 func nodeAtDistance(base enode.ID, ld int, ip net.IP) *node {
 	var r enr.Record
 	r.Set(enr.IP(ip))
+	r.Set(enr.UDP(30303))
 	return wrapNode(enode.SignNull(&r, idAtDistance(base, ld)))
 }
 

p2p/discover/v5_talk.go

+// Copyright 2023 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package discover
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/p2p/discover/v5wire"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+)
+
+// This is a limit for the number of concurrent talk requests.
+const maxActiveTalkRequests = 1024
+
+// This is the timeout for acquiring a handler execution slot for a talk request.
+// The timeout should be short enough to fit within the request timeout.
+const talkHandlerLaunchTimeout = 400 * time.Millisecond
+
+// TalkRequestHandler callback processes a talk request and returns a response.
+//
+// Note that talk handlers are expected to come up with a response very quickly, within at
+// most 200ms or so. If the handler takes longer than that, the remote end may time out
+// and wont receive the response.
+type TalkRequestHandler func(enode.ID, *net.UDPAddr, []byte) []byte
+
+type talkSystem struct {
+	transport *UDPv5
+
+	mutex     sync.Mutex
+	handlers  map[string]TalkRequestHandler
+	slots     chan struct{}
+	lastLog   time.Time
+	dropCount int
+}
+
+func newTalkSystem(transport *UDPv5) *talkSystem {
+	t := &talkSystem{
+		transport: transport,
+		handlers:  make(map[string]TalkRequestHandler),
+		slots:     make(chan struct{}, maxActiveTalkRequests),
+	}
+	for i := 0; i < cap(t.slots); i++ {
+		t.slots <- struct{}{}
+	}
+	return t
+}
+
+// register adds a protocol handler.
+func (t *talkSystem) register(protocol string, handler TalkRequestHandler) {
+	t.mutex.Lock()
+	t.handlers[protocol] = handler
+	t.mutex.Unlock()
+}
+
+// handleRequest handles a talk request.
+func (t *talkSystem) handleRequest(id enode.ID, addr *net.UDPAddr, req *v5wire.TalkRequest) {
+	t.mutex.Lock()
+	handler, ok := t.handlers[req.Protocol]
+	t.mutex.Unlock()
+
+	if !ok {
+		resp := &v5wire.TalkResponse{ReqID: req.ReqID}
+		t.transport.sendResponse(id, addr, resp)
+		return
+	}
+
+	// Wait for a slot to become available, then run the handler.
+	timeout := time.NewTimer(talkHandlerLaunchTimeout)
+	defer timeout.Stop()
+	select {
+	case <-t.slots:
+		go func() {
+			defer func() { t.slots <- struct{}{} }()
+			respMessage := handler(id, addr, req.Message)
+			resp := &v5wire.TalkResponse{ReqID: req.ReqID, Message: respMessage}
+			t.transport.sendFromAnotherThread(id, addr, resp)
+		}()
+	case <-timeout.C:
+		// Couldn't get it in time, drop the request.
+		if time.Since(t.lastLog) > 5*time.Second {
+			log.Warn("Dropping TALKREQ due to overload", "ndrop", t.dropCount)
+			t.lastLog = time.Now()
+			t.dropCount++
+		}
+	case <-t.transport.closeCtx.Done():
+		// Transport closed, drop the request.
+	}
+}
+
+// wait blocks until all active requests have finished, and prevents new request
+// handlers from being launched.
+func (t *talkSystem) wait() {
+	for i := 0; i < cap(t.slots); i++ {
+		<-t.slots
+	}
+}

p2p/discover/v5_udp_test.go

@@ -186,7 +186,7 @@ func TestUDPv5_findnodeHandling(t *testing.T) {
 
 	// This request gets all the distance-253 nodes.
 	test.packetIn(&v5wire.Findnode{ReqID: []byte{4}, Distances: []uint{253}})
-	test.expectNodes([]byte{4}, 1, nodes253)
+	test.expectNodes([]byte{4}, 2, nodes253)
 
 	// This request gets all the distance-249 nodes and some more at 248 because
 	// the bucket at 249 is not full.
@@ -515,6 +515,27 @@ func TestUDPv5_talkRequest(t *testing.T) {
 	if err := <-done; err != nil {
 		t.Fatal(err)
 	}
+
+	// Also check requesting without ENR.
+	go func() {
+		_, err := test.udp.TalkRequestToID(remote.ID(), test.remoteaddr, "test", []byte("test request 2"))
+		done <- err
+	}()
+	test.waitPacketOut(func(p *v5wire.TalkRequest, addr *net.UDPAddr, _ v5wire.Nonce) {
+		if p.Protocol != "test" {
+			t.Errorf("wrong protocol ID in talk request: %q", p.Protocol)
+		}
+		if string(p.Message) != "test request 2" {
+			t.Errorf("wrong message talk request: %q", p.Message)
+		}
+		test.packetInFrom(test.remotekey, test.remoteaddr, &v5wire.TalkResponse{
+			ReqID:   p.ReqID,
+			Message: []byte("test response 2"),
+		})
+	})
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
 }
 
 // This test checks that lookupDistances works.

p2p/discover/v5wire/msg.go

@@ -216,7 +216,7 @@ func (p *TalkRequest) RequestID() []byte      { return p.ReqID }
 func (p *TalkRequest) SetRequestID(id []byte) { p.ReqID = id }
 
 func (p *TalkRequest) AppendLogInfo(ctx []interface{}) []interface{} {
-	return append(ctx, "proto", p.Protocol, "reqid", hexutil.Bytes(p.ReqID), "len", len(p.Message))
+	return append(ctx, "proto", p.Protocol, "req", hexutil.Bytes(p.ReqID), "len", len(p.Message))
 }
 
 func (*TalkResponse) Name() string             { return "TALKRESP/v5" }
@@ -225,5 +225,5 @@ func (p *TalkResponse) RequestID() []byte      { return p.ReqID }
 func (p *TalkResponse) SetRequestID(id []byte) { p.ReqID = id }
 
 func (p *TalkResponse) AppendLogInfo(ctx []interface{}) []interface{} {
-	return append(ctx, "req", p.ReqID, "len", len(p.Message))
+	return append(ctx, "req", hexutil.Bytes(p.ReqID), "len", len(p.Message))
 }