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