github.com/turingchain2020/turingchain@v1.1.21/system/p2p/dht/protocol/manager.go (about)

     1  package protocol
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/turingchain2020/turingchain/queue"
     8  	"github.com/libp2p/go-libp2p-core/host"
     9  	"github.com/libp2p/go-libp2p-core/network"
    10  	"github.com/libp2p/go-libp2p-core/protocol"
    11  )
    12  
    13  // RegisterStreamHandler registers stream handler
    14  func RegisterStreamHandler(h host.Host, p protocol.ID, handler network.StreamHandler) {
    15  	if handler == nil {
    16  		panic(fmt.Sprintf("addEventHandler, handler is nil, protocol=%s", p))
    17  	}
    18  	h.SetStreamHandler(p, HandlerWithClose(handler))
    19  }
    20  
    21  //Initializer is a initial function which any protocol should have.
    22  type Initializer func(env *P2PEnv)
    23  
    24  var (
    25  	protocolInitializerArray []Initializer
    26  )
    27  
    28  //RegisterProtocolInitializer registers the initial function.
    29  func RegisterProtocolInitializer(initializer Initializer) {
    30  	protocolInitializerArray = append(protocolInitializerArray, initializer)
    31  }
    32  
    33  //InitAllProtocol initials all protocols.
    34  func InitAllProtocol(env *P2PEnv) {
    35  	for _, initializer := range protocolInitializerArray {
    36  		initializer(env)
    37  	}
    38  }
    39  
    40  // EventHandler handles turingchain event
    41  type EventHandler func(*queue.Message)
    42  
    43  var (
    44  	eventHandlers = make(map[int64]EventHandler)
    45  	mu            sync.RWMutex
    46  )
    47  
    48  // RegisterEventHandler registers a handler with an event ID.
    49  func RegisterEventHandler(eventID int64, handler EventHandler) {
    50  	if handler == nil {
    51  		panic(fmt.Sprintf("addEventHandler, handler is nil, id=%d", eventID))
    52  	}
    53  	if _, dup := eventHandlers[eventID]; dup {
    54  		panic(fmt.Sprintf("addEventHandler, duplicate handler, id=%d, len=%d", eventID, len(eventHandlers)))
    55  	}
    56  	eventHandlers[eventID] = EventHandlerWithRecover(handler)
    57  }
    58  
    59  // GetEventHandler gets event handler by event ID.
    60  func GetEventHandler(eventID int64) EventHandler {
    61  	mu.RLock()
    62  	defer mu.RUnlock()
    63  	return eventHandlers[eventID]
    64  }
    65  
    66  // ClearEventHandler clear event handler map, plugin存在多个p2p实例测试，会导致重复注册，需要清除
    67  func ClearEventHandler() {
    68  	mu.Lock()
    69  	defer mu.Unlock()
    70  	for k := range eventHandlers {
    71  		delete(eventHandlers, k)
    72  	}
    73  }