github.com/ethereum/go-ethereum@v1.16.1/p2p/discover/v5_talk.go (about)

     1  // Copyright 2023 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package discover
    18  
    19  import (
    20  	"net"
    21  	"net/netip"
    22  	"sync"
    23  	"time"
    24  
    25  	"github.com/ethereum/go-ethereum/log"
    26  	"github.com/ethereum/go-ethereum/p2p/discover/v5wire"
    27  	"github.com/ethereum/go-ethereum/p2p/enode"
    28  )
    29  
    30  // This is a limit for the number of concurrent talk requests.
    31  const maxActiveTalkRequests = 1024
    32  
    33  // This is the timeout for acquiring a handler execution slot for a talk request.
    34  // The timeout should be short enough to fit within the request timeout.
    35  const talkHandlerLaunchTimeout = 400 * time.Millisecond
    36  
    37  // TalkRequestHandler callback processes a talk request and returns a response.
    38  //
    39  // Note that talk handlers are expected to come up with a response very quickly, within at
    40  // most 200ms or so. If the handler takes longer than that, the remote end may time out
    41  // and wont receive the response.
    42  type TalkRequestHandler func(*enode.Node, *net.UDPAddr, []byte) []byte
    43  
    44  type talkSystem struct {
    45  	transport *UDPv5
    46  
    47  	mutex     sync.Mutex
    48  	handlers  map[string]TalkRequestHandler
    49  	slots     chan struct{}
    50  	lastLog   time.Time
    51  	dropCount int
    52  }
    53  
    54  func newTalkSystem(transport *UDPv5) *talkSystem {
    55  	t := &talkSystem{
    56  		transport: transport,
    57  		handlers:  make(map[string]TalkRequestHandler),
    58  		slots:     make(chan struct{}, maxActiveTalkRequests),
    59  	}
    60  	for i := 0; i < cap(t.slots); i++ {
    61  		t.slots <- struct{}{}
    62  	}
    63  	return t
    64  }
    65  
    66  // register adds a protocol handler.
    67  func (t *talkSystem) register(protocol string, handler TalkRequestHandler) {
    68  	t.mutex.Lock()
    69  	t.handlers[protocol] = handler
    70  	t.mutex.Unlock()
    71  }
    72  
    73  // handleRequest handles a talk request.
    74  func (t *talkSystem) handleRequest(id enode.ID, addr netip.AddrPort, req *v5wire.TalkRequest) {
    75  	n := t.transport.codec.SessionNode(id, addr.String())
    76  	if n == nil {
    77  		// The node must be contained in the session here, since we wouldn't have
    78  		// received the request otherwise.
    79  		panic("missing node in session")
    80  	}
    81  	t.mutex.Lock()
    82  	handler, ok := t.handlers[req.Protocol]
    83  	t.mutex.Unlock()
    84  
    85  	if !ok {
    86  		resp := &v5wire.TalkResponse{ReqID: req.ReqID}
    87  		t.transport.sendResponse(n.ID(), addr, resp)
    88  		return
    89  	}
    90  
    91  	// Wait for a slot to become available, then run the handler.
    92  	timeout := time.NewTimer(talkHandlerLaunchTimeout)
    93  	defer timeout.Stop()
    94  	select {
    95  	case <-t.slots:
    96  		go func() {
    97  			defer func() { t.slots <- struct{}{} }()
    98  			udpAddr := &net.UDPAddr{IP: addr.Addr().AsSlice(), Port: int(addr.Port())}
    99  			respMessage := handler(n, udpAddr, req.Message)
   100  			resp := &v5wire.TalkResponse{ReqID: req.ReqID, Message: respMessage}
   101  			t.transport.sendFromAnotherThread(n.ID(), addr, resp)
   102  		}()
   103  	case <-timeout.C:
   104  		// Couldn't get it in time, drop the request.
   105  		if time.Since(t.lastLog) > 5*time.Second {
   106  			log.Warn("Dropping TALKREQ due to overload", "ndrop", t.dropCount)
   107  			t.lastLog = time.Now()
   108  			t.dropCount++
   109  		}
   110  	case <-t.transport.closeCtx.Done():
   111  		// Transport closed, drop the request.
   112  	}
   113  }
   114  
   115  // wait blocks until all active requests have finished, and prevents new request
   116  // handlers from being launched.
   117  func (t *talkSystem) wait() {
   118  	for i := 0; i < cap(t.slots); i++ {
   119  		<-t.slots
   120  	}
   121  }