use signal dispatcher + refactor

go.mod

@@ -161,6 +161,7 @@ require (
 	github.com/moby/sys/user v0.1.0 // indirect
 	github.com/moby/term v0.5.0 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20240801161502-2171ea0babf1 // indirect
 	github.com/nxadm/tail v1.4.8 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0 // indirect

go.sum

@@ -337,6 +337,8 @@ github.com/netbirdio/management-integrations/integrations v0.0.0-20240703085513-
 github.com/netbirdio/management-integrations/integrations v0.0.0-20240703085513-32605f7ffd8e/go.mod h1:nykwWZnxb+sJz2Z//CEq45CMRWSHllH8pODKRB8eY7Y=
 github.com/netbirdio/service v0.0.0-20230215170314-b923b89432b0 h1:hirFRfx3grVA/9eEyjME5/z3nxdJlN9kfQpvWWPk32g=
 github.com/netbirdio/service v0.0.0-20230215170314-b923b89432b0/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
+github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20240801161502-2171ea0babf1 h1:YMPC80v0/bCWXjpCraSSpgh7/5lsNUZ1qoxVzzmRu5k=
+github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20240801161502-2171ea0babf1/go.mod h1:5/sjFmLb8O96B5737VCqhHyGRzNFIaN/Bu7ZodXc3qQ=
 github.com/netbirdio/systray v0.0.0-20231030152038-ef1ed2a27949 h1:xbWM9BU6mwZZLHxEjxIX/V8Hv3HurQt4mReIE4mY4DM=
 github.com/netbirdio/systray v0.0.0-20231030152038-ef1ed2a27949/go.mod h1:AecygODWIsBquJCJFop8MEQcJbWFfw/1yWbVabNgpCM=
 github.com/netbirdio/wireguard-go v0.0.0-20240105182236-6c340dd55aed h1:t0UADZUJDaaZgfKrt8JUPrOLL9Mg/ryjP85RAH53qgs=

signal/peer/peer.go

@@ -18,16 +18,20 @@ type Peer struct {
 
 	StreamID int64
 
-	//a gRpc connection stream to the Peer
+	// a gRpc connection stream to the Peer
 	Stream proto.SignalExchange_ConnectStreamServer
+
+	// registration time
+	RegisteredAt time.Time
 }
 
 // NewPeer creates a new instance of a connected Peer
 func NewPeer(id string, stream proto.SignalExchange_ConnectStreamServer) *Peer {
 	return &Peer{
-		Id:       id,
+		Id:           id,
-		Stream:   stream,
+		Stream:       stream,
-		StreamID: time.Now().UnixNano(),
+		StreamID:     time.Now().UnixNano(),
+		RegisteredAt: time.Now(),
 	}
 }
 

signal/server/signal.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"time"
 
+	"github.com/netbirdio/signal-dispatcher/dispatcher"
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
@@ -40,8 +41,8 @@ const (
 type Server struct {
 	registry *peer.Registry
 	proto.UnimplementedSignalExchangeServer
-
+	dispatcher *dispatcher.Dispatcher
-	metrics *metrics.AppMetrics
+	metrics    *metrics.AppMetrics
 }
 
 // NewServer creates a new Signal server
@@ -51,9 +52,14 @@ func NewServer(meter metric.Meter) (*Server, error) {
 		return nil, fmt.Errorf("creating app metrics: %v", err)
 	}
 
+	dispatcher, err := dispatcher.NewDispatcher()
+	if err != nil {
+		return nil, fmt.Errorf("creating dispatcher: %v", err)
+	}
 	s := &Server{
-		registry: peer.NewRegistry(appMetrics),
+		dispatcher: dispatcher,
-		metrics:  appMetrics,
+		registry:   peer.NewRegistry(appMetrics),
+		metrics:    appMetrics,
 	}
 
 	return s, nil
@@ -61,57 +67,26 @@ func NewServer(meter metric.Meter) (*Server, error) {
 
 // Send forwards a message to the signal peer
 func (s *Server) Send(ctx context.Context, msg *proto.EncryptedMessage) (*proto.EncryptedMessage, error) {
-	if !s.registry.IsPeerRegistered(msg.Key) {
+	log.Debugf("received a new message to send from peer [%s] to peer [%s]", msg.Key, msg.RemoteKey)
-		s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeNotRegistered)))
 
-		return nil, fmt.Errorf("peer %s is not registered", msg.Key)
+	if msg.RemoteKey == "dummy" {
+		// Test message send during netbird status
+		return &proto.EncryptedMessage{}, nil
 	}
 
-	getRegistrationStart := time.Now()
+	return s.dispatcher.SendMessage(ctx, msg)
-
-	if dstPeer, found := s.registry.Get(msg.RemoteKey); found {
-		s.metrics.GetRegistrationDelay.Record(ctx, float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeMessage), attribute.String(labelRegistrationStatus, labelRegistrationFound)))
-		start := time.Now()
-		//forward the message to the target peer
-		if err := dstPeer.Stream.Send(msg); err != nil {
-			log.Errorf("error while forwarding message from peer [%s] to peer [%s] %v", msg.Key, msg.RemoteKey, err)
-			//todo respond to the sender?
-
-			s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeError)))
-		} else {
-			s.metrics.MessageForwardLatency.Record(ctx, float64(time.Since(start).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeMessage)))
-			s.metrics.MessagesForwarded.Add(context.Background(), 1)
-		}
-	} else {
-		s.metrics.GetRegistrationDelay.Record(ctx, float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeMessage), attribute.String(labelRegistrationStatus, labelRegistrationNotFound)))
-		log.Debugf("message from peer [%s] can't be forwarded to peer [%s] because destination peer is not connected", msg.Key, msg.RemoteKey)
-		//todo respond to the sender?
-
-		s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeNotConnected)))
-	}
-	return &proto.EncryptedMessage{}, nil
 }
 
 // ConnectStream connects to the exchange stream
 func (s *Server) ConnectStream(stream proto.SignalExchange_ConnectStreamServer) error {
-	p, err := s.connectPeer(stream)
+	p, err := s.RegisterPeer(stream)
 	if err != nil {
 		return err
 	}
 
-	startRegister := time.Now()
+	defer s.DeregisterPeer(p)
 
-	s.metrics.ActivePeers.Add(stream.Context(), 1)
+	// needed to confirm that the peer has been registered so that the client can proceed
-
-	defer func() {
-		log.Infof("peer disconnected [%s] [streamID %d] ", p.Id, p.StreamID)
-		s.registry.Deregister(p)
-
-		s.metrics.PeerConnectionDuration.Record(stream.Context(), int64(time.Since(startRegister).Seconds()))
-		s.metrics.ActivePeers.Add(context.Background(), -1)
-	}()
-
-	//needed to confirm that the peer has been registered so that the client can proceed
 	header := metadata.Pairs(proto.HeaderRegistered, "1")
 	err = stream.SendHeader(header)
 	if err != nil {
@@ -119,11 +94,10 @@ func (s *Server) ConnectStream(stream proto.SignalExchange_ConnectStreamServer)
 		return err
 	}
 
-	log.Infof("peer connected [%s] [streamID %d] ", p.Id, p.StreamID)
+	log.Debugf("peer connected [%s] [streamID %d] ", p.Id, p.StreamID)
 
 	for {
-
+		// read incoming messages
-		//read incoming messages
 		msg, err := stream.Recv()
 		if err == io.EOF {
 			break
@@ -131,44 +105,25 @@ func (s *Server) ConnectStream(stream proto.SignalExchange_ConnectStreamServer)
 			return err
 		}
 
-		log.Debugf("received a new message from peer [%s] to peer [%s]", p.Id, msg.RemoteKey)
+		log.Debugf("Received a response from peer [%s] to peer [%s]", msg.Key, msg.RemoteKey)
 
-		getRegistrationStart := time.Now()
+		s.dispatcher.SendMessage(stream.Context(), msg)
-
-		// lookup the target peer where the message is going to
-		if dstPeer, found := s.registry.Get(msg.RemoteKey); found {
-			s.metrics.GetRegistrationDelay.Record(stream.Context(), float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream), attribute.String(labelRegistrationStatus, labelRegistrationFound)))
-			start := time.Now()
-			//forward the message to the target peer
-			if err := dstPeer.Stream.Send(msg); err != nil {
-				log.Errorf("error while forwarding message from peer [%s] to peer [%s] %v", p.Id, msg.RemoteKey, err)
-				//todo respond to the sender?
-				s.metrics.MessageForwardFailures.Add(stream.Context(), 1, metric.WithAttributes(attribute.String(labelType, labelTypeError)))
-			} else {
-				// in milliseconds
-				s.metrics.MessageForwardLatency.Record(stream.Context(), float64(time.Since(start).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream)))
-				s.metrics.MessagesForwarded.Add(stream.Context(), 1)
-			}
-		} else {
-			s.metrics.GetRegistrationDelay.Record(stream.Context(), float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream), attribute.String(labelRegistrationStatus, labelRegistrationNotFound)))
-			s.metrics.MessageForwardFailures.Add(stream.Context(), 1, metric.WithAttributes(attribute.String(labelType, labelTypeNotConnected)))
-			log.Debugf("message from peer [%s] can't be forwarded to peer [%s] because destination peer is not connected", p.Id, msg.RemoteKey)
-			//todo respond to the sender?
-		}
 	}
+
 	<-stream.Context().Done()
 	return stream.Context().Err()
 }
 
-// Handles initial Peer connection.
+func (s *Server) RegisterPeer(stream proto.SignalExchange_ConnectStreamServer) (*peer.Peer, error) {
-// Each connection must provide an Id header.
+	log.Debugf("registering new peer")
-// At this moment the connecting Peer will be registered in the peer.Registry
-func (s Server) connectPeer(stream proto.SignalExchange_ConnectStreamServer) (*peer.Peer, error) {
 	if meta, hasMeta := metadata.FromIncomingContext(stream.Context()); hasMeta {
 		if id, found := meta[proto.HeaderId]; found {
 			p := peer.NewPeer(id[0], stream)
 
 			s.registry.Register(p)
+			s.dispatcher.ListenForMessages(stream.Context(), p.Id, s.forwardMessageToPeer)
+
+			s.metrics.ActivePeers.Add(stream.Context(), 1)
 
 			return p, nil
 		} else {
@@ -180,3 +135,38 @@ func (s Server) connectPeer(stream proto.SignalExchange_ConnectStreamServer) (*p
 		return nil, status.Errorf(codes.FailedPrecondition, "missing connection stream meta")
 	}
 }
+
+func (s *Server) DeregisterPeer(p *peer.Peer) {
+	log.Debugf("peer disconnected [%s] [streamID %d] ", p.Id, p.StreamID)
+	s.registry.Deregister(p)
+
+	s.metrics.PeerConnectionDuration.Record(p.Stream.Context(), int64(time.Since(p.RegisteredAt).Seconds()))
+	s.metrics.ActivePeers.Add(context.Background(), -1)
+}
+
+func (s *Server) forwardMessageToPeer(ctx context.Context, msg *proto.EncryptedMessage) {
+	log.Debugf("forwarding a new message from peer [%s] to peer [%s]", msg.Key, msg.RemoteKey)
+
+	getRegistrationStart := time.Now()
+
+	// lookup the target peer where the message is going to
+	if dstPeer, found := s.registry.Get(msg.RemoteKey); found {
+		s.metrics.GetRegistrationDelay.Record(ctx, float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream), attribute.String(labelRegistrationStatus, labelRegistrationFound)))
+		start := time.Now()
+		// forward the message to the target peer
+		if err := dstPeer.Stream.Send(msg); err != nil {
+			log.Warnf("error while forwarding message from peer [%s] to peer [%s] %v", msg.Key, msg.RemoteKey, err)
+			// todo respond to the sender?
+			s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeError)))
+		} else {
+			// in milliseconds
+			s.metrics.MessageForwardLatency.Record(ctx, float64(time.Since(start).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream)))
+			s.metrics.MessagesForwarded.Add(ctx, 1)
+		}
+	} else {
+		s.metrics.GetRegistrationDelay.Record(ctx, float64(time.Since(getRegistrationStart).Nanoseconds())/1e6, metric.WithAttributes(attribute.String(labelType, labelTypeStream), attribute.String(labelRegistrationStatus, labelRegistrationNotFound)))
+		s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeNotConnected)))
+		log.Debugf("message from peer [%s] can't be forwarded to peer [%s] because destination peer is not connected", msg.Key, msg.RemoteKey)
+		// todo respond to the sender?
+	}
+}