Automatically disconnect and remove the peer from consul registry if peer suddenly dies

Author: anhthii

cmd/mpcium/main.go

@@ -162,7 +162,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	logger.Info("Node is running", "ID", nodeID, "name", nodeName)
 
 	peerNodeIDs := GetPeerIDs(peers)
-	peerRegistry := mpc.NewRegistry(nodeID, peerNodeIDs, consulClient.KV())
+	peerRegistry := mpc.NewRegistry(nodeID, peerNodeIDs, consulClient.KV(), directMessaging)
 
 	mpcNode := mpc.NewNode(
 		nodeID,
@@ -232,6 +232,11 @@ func runNode(ctx context.Context, c *cli.Command) error {
 		if err := ecdhSession.Close(); err != nil {
 			logger.Error("Failed to close ECDH session", err)
 		}
+
+		err := natsConn.Drain()
+		if err != nil {
+			logger.Error("Failed to drain NATS connection", err)
+		}
 	}()
 
 	var wg sync.WaitGroup

pkg/eventconsumer/keygen_consumer.go

@@ -90,10 +90,10 @@ func (sc *keygenConsumer) Run(ctx context.Context) error {
 		sc.handleKeygenEvent,
 	)
 	if err != nil {
-		return fmt.Errorf("failed to subscribe to signing events: %w", err)
+		return fmt.Errorf("failed to subscribe to keygen events: %w", err)
 	}
 	sc.jsSub = sub
-	logger.Info("SigningConsumer: Subscribed to signing events")
+	logger.Info("SigningConsumer: Subscribed to keygen events")
 
 	// Block until context cancellation.
 	<-ctx.Done()

pkg/messaging/point2point.go

@@ -13,9 +13,17 @@ import (
 type DirectMessaging interface {
 	Listen(topic string, handler func(data []byte)) (Subscription, error)
 	SendToOther(topic string, data []byte) error
+	SendToOtherWithRetry(topic string, data []byte, config RetryConfig) error
 	SendToSelf(topic string, data []byte) error
 }
 
+type RetryConfig struct {
+	RetryAttempt       uint
+	ExponentialBackoff bool
+	Delay              time.Duration
+	OnRetry            func(n uint, err error)
+}
+
 type natsDirectMessaging struct {
 	natsConn *nats.Conn
 	handlers map[string][]func([]byte)
@@ -65,6 +73,36 @@ func (d *natsDirectMessaging) SendToOther(topic string, message []byte) error {
 	)
 }
 
+func (d *natsDirectMessaging) SendToOtherWithRetry(topic string, message []byte, config RetryConfig) error {
+	opts := []retry.Option{
+		retry.MaxJitter(80 * time.Millisecond),
+	}
+
+	if config.RetryAttempt > 0 {
+		opts = append(opts, retry.Attempts(config.RetryAttempt))
+	}
+	if config.ExponentialBackoff {
+		opts = append(opts, retry.DelayType(retry.BackOffDelay))
+	}
+	if config.Delay > 0 {
+		opts = append(opts, retry.Delay(config.Delay))
+	}
+	if config.OnRetry != nil {
+		opts = append(opts, retry.OnRetry(config.OnRetry))
+	}
+
+	return retry.Do(
+		func() error {
+			_, err := d.natsConn.Request(topic, message, 3*time.Second)
+			if err != nil {
+				return err
+			}
+			return nil
+		},
+		opts...,
+	)
+}
+
 func (d *natsDirectMessaging) Listen(topic string, handler func(data []byte)) (Subscription, error) {
 	sub, err := d.natsConn.Subscribe(topic, func(m *nats.Msg) {
 		handler(m.Data)
@@ -76,6 +114,14 @@ func (d *natsDirectMessaging) Listen(topic string, handler func(data []byte)) (S
 		return nil, err
 	}
 
+	if err := d.natsConn.Flush(); err != nil {
+		err := sub.Unsubscribe()
+		if err != nil {
+			logger.Error("Failed to unsubscribe", err)
+		}
+		return nil, fmt.Errorf("flush after subscribe failed: %w", err)
+	}
+
 	d.mu.Lock()
 	d.handlers[topic] = append(d.handlers[topic], handler)
 	d.mu.Unlock()

pkg/mpc/registry.go

@@ -2,12 +2,14 @@ package mpc
 
 import (
 	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/fystack/mpcium/pkg/infra"
 	"github.com/fystack/mpcium/pkg/logger"
+	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/hashicorp/consul/api"
 	"github.com/samber/lo"
 )
@@ -35,20 +37,23 @@ type registry struct {
 	mu          sync.RWMutex
 	ready       bool // ready is true when all peers are ready
 
-	consulKV infra.ConsulKV
+	consulKV    infra.ConsulKV
+	healthCheck messaging.DirectMessaging
 }
 
 func NewRegistry(
 	nodeID string,
 	peerNodeIDs []string,
 	consulKV infra.ConsulKV,
+	directMessaging messaging.DirectMessaging,
 ) *registry {
 	return &registry{
 		consulKV:    consulKV,
 		nodeID:      nodeID,
 		peerNodeIDs: getPeerIDsExceptSelf(nodeID, peerNodeIDs),
 		readyMap:    make(map[string]bool),
 		readyCount:  1, // self
+		healthCheck: directMessaging,
 	}
 }
 
@@ -104,12 +109,22 @@ func (r *registry) Ready() error {
 		return fmt.Errorf("Put ready key failed: %w", err)
 	}
 
+	_, err = r.healthCheck.Listen(r.composeHealthCheckTopic(r.nodeID), func(data []byte) {
+		logger.Debug("Health check", "peerID", string(data))
+	})
+	if err != nil {
+		return fmt.Errorf("Listen health check failed: %w", err)
+	}
 	return nil
 }
 
+func (r *registry) composeHealthCheckTopic(nodeID string) string {
+	return fmt.Sprintf("healthcheck:%s", nodeID)
+}
+
 func (r *registry) WatchPeersReady(callback func()) {
 	ticker := time.NewTicker(ReadinessCheckPeriod)
-	go r.logReadyStatus()
+	go r.checkPeersHeath()
 	// first tick is executed immediately
 	for ; true; <-ticker.C {
 		pairs, _, err := r.consulKV.List("ready/", nil)
@@ -146,12 +161,31 @@ func (r *registry) WatchPeersReady(callback func()) {
 
 }
 
-func (r *registry) logReadyStatus() {
+func (r *registry) checkPeersHeath() {
 	for {
 		time.Sleep(5 * time.Second)
 		if !r.ArePeersReady() {
 			logger.Info("Peers are not ready yet", "ready", r.GetReadyPeersCount(), "expected", len(r.peerNodeIDs)+1)
 		}
+
+		pairs, _, err := r.consulKV.List("ready/", nil)
+		if err != nil {
+			logger.Error("List ready keys failed", err)
+			continue
+		}
+		readyPeerIDs := r.getReadyPeersFromKVStore(pairs)
+		for _, peerID := range readyPeerIDs {
+			err := r.healthCheck.SendToOtherWithRetry(r.composeHealthCheckTopic(peerID), []byte(peerID), messaging.RetryConfig{
+				RetryAttempt: 2,
+			})
+			if err != nil && strings.Contains(err.Error(), "no responders") {
+				logger.Info("No response from peer", "peerID", peerID)
+				_, err := r.consulKV.Delete(r.readyKey(peerID), nil)
+				if err != nil {
+					logger.Error("Delete ready key failed", err)
+				}
+			}
+		}
 	}
 }