github.com/amazechain/amc@v0.1.3/internal/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  	"github.com/amazechain/amc/internal/p2p/discover/v5wire"
    21  	"github.com/amazechain/amc/internal/p2p/enode"
    22  	"github.com/amazechain/amc/log"
    23  	"net"
    24  	"sync"
    25  	"time"
    26  )
    27  
    28  // This is a limit for the number of concurrent talk requests.
    29  const maxActiveTalkRequests = 1024
    30  
    31  // This is the timeout for acquiring a handler execution slot for a talk request.
    32  // The timeout should be short enough to fit within the request timeout.
    33  const talkHandlerLaunchTimeout = 400 * time.Millisecond
    34  
    35  // TalkRequestHandler callback processes a talk request and returns a response.
    36  //
    37  // Note that talk handlers are expected to come up with a response very quickly, within at
    38  // most 200ms or so. If the handler takes longer than that, the remote end may time out
    39  // and wont receive the response.
    40  type TalkRequestHandler func(enode.ID, *net.UDPAddr, []byte) []byte
    41  
    42  type talkSystem struct {
    43  	transport *UDPv5
    44  
    45  	mutex     sync.Mutex
    46  	handlers  map[string]TalkRequestHandler
    47  	slots     chan struct{}
    48  	lastLog   time.Time
    49  	dropCount int
    50  }
    51  
    52  func newTalkSystem(transport *UDPv5) *talkSystem {
    53  	t := &talkSystem{
    54  		transport: transport,
    55  		handlers:  make(map[string]TalkRequestHandler),
    56  		slots:     make(chan struct{}, maxActiveTalkRequests),
    57  	}
    58  	for i := 0; i < cap(t.slots); i++ {
    59  		t.slots <- struct{}{}
    60  	}
    61  	return t
    62  }
    63  
    64  // register adds a protocol handler.
    65  func (t *talkSystem) register(protocol string, handler TalkRequestHandler) {
    66  	t.mutex.Lock()
    67  	t.handlers[protocol] = handler
    68  	t.mutex.Unlock()
    69  }
    70  
    71  // handleRequest handles a talk request.
    72  func (t *talkSystem) handleRequest(id enode.ID, addr *net.UDPAddr, req *v5wire.TalkRequest) {
    73  	t.mutex.Lock()
    74  	handler, ok := t.handlers[req.Protocol]
    75  	t.mutex.Unlock()
    76  
    77  	if !ok {
    78  		resp := &v5wire.TalkResponse{ReqID: req.ReqID}
    79  		t.transport.sendResponse(id, addr, resp)
    80  		return
    81  	}
    82  
    83  	// Wait for a slot to become available, then run the handler.
    84  	timeout := time.NewTimer(talkHandlerLaunchTimeout)
    85  	defer timeout.Stop()
    86  	select {
    87  	case <-t.slots:
    88  		go func() {
    89  			defer func() { t.slots <- struct{}{} }()
    90  			respMessage := handler(id, addr, req.Message)
    91  			resp := &v5wire.TalkResponse{ReqID: req.ReqID, Message: respMessage}
    92  			t.transport.sendFromAnotherThread(id, addr, resp)
    93  		}()
    94  	case <-timeout.C:
    95  		// Couldn't get it in time, drop the request.
    96  		if time.Since(t.lastLog) > 5*time.Second {
    97  			log.Warn("Dropping TALKREQ due to overload", "ndrop", t.dropCount)
    98  			t.lastLog = time.Now()
    99  			t.dropCount++
   100  		}
   101  	case <-t.transport.closeCtx.Done():
   102  		// Transport closed, drop the request.
   103  	}
   104  }
   105  
   106  // wait blocks until all active requests have finished, and prevents new request
   107  // handlers from being launched.
   108  func (t *talkSystem) wait() {
   109  	for i := 0; i < cap(t.slots); i++ {
   110  		<-t.slots
   111  	}
   112  }