Skip logging error on duplicate key gen

Author: anhthii

pkg/eventconsumer/event_consumer.go

@@ -137,13 +137,11 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 	walletID := msg.WalletID
 	ecdsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeECDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue)
 	if err != nil {
-		logger.Error("Failed to create ECDSA key generation session", err, "walletID", walletID)
 		ec.handleKeygenSessionError(walletID, err, "Failed to create ECDSA key generation session", natMsg)
 		return
 	}
 	eddsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeEDDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue)
 	if err != nil {
-		logger.Error("Failed to create EdDSA key generation session", err, "walletID", walletID)
 		ec.handleKeygenSessionError(walletID, err, "Failed to create EdDSA key generation session", natMsg)
 		return
 	}
@@ -225,7 +223,11 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 	}
 
 	key := fmt.Sprintf(mpc.TypeGenerateWalletResultFmt, walletID)
-	if err := ec.genKeyResultQueue.Enqueue(key, payload, &messaging.EnqueueOptions{IdempotententKey: key}); err != nil {
+	if err := ec.genKeyResultQueue.Enqueue(
+		key,
+		payload,
+		&messaging.EnqueueOptions{IdempotententKey: composeKeygenIdempotentKey(walletID, natMsg)},
+	); err != nil {
 		logger.Error("Failed to publish key generation success message", err)
 		ec.handleKeygenSessionError(walletID, err, "Failed to publish key generation success message", natMsg)
 		return
@@ -238,14 +240,6 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 func (ec *eventConsumer) handleKeygenSessionError(walletID string, err error, contextMsg string, natMsg *nats.Msg) {
 	fullErrMsg := fmt.Sprintf("%s: %v", contextMsg, err)
 	errorCode := event.GetErrorCodeFromError(err)
-
-	logger.Warn("Keygen session error",
-		"walletID", walletID,
-		"error", err.Error(),
-		"errorCode", errorCode,
-		"context", contextMsg,
-	)
-
 	keygenResult := event.KeygenResultEvent{
 		ResultType:  event.ResultTypeError,
 		ErrorCode:   string(errorCode),
@@ -263,7 +257,7 @@ func (ec *eventConsumer) handleKeygenSessionError(walletID string, err error, co
 
 	key := fmt.Sprintf(mpc.TypeGenerateWalletResultFmt, walletID)
 	err = ec.genKeyResultQueue.Enqueue(key, keygenResultBytes, &messaging.EnqueueOptions{
-		IdempotententKey: key,
+		IdempotententKey: composeKeygenIdempotentKey(walletID, natMsg),
 	})
 	if err != nil {
 		logger.Error("Failed to enqueue keygen result event", err,
@@ -795,3 +789,14 @@ func sessionTypeFromKeyType(keyType types.KeyType) (mpc.SessionType, error) {
 		return "", fmt.Errorf("unsupported key type: %v", keyType)
 	}
 }
+
+func composeKeygenIdempotentKey(walletID string, natMsg *nats.Msg) string {
+	var uniqueKey string
+	sid := natMsg.Header.Get("SessionID")
+	if sid != "" {
+		uniqueKey = fmt.Sprintf("%s:%s", walletID, sid)
+	} else {
+		uniqueKey = walletID
+	}
+	return fmt.Sprintf(mpc.TypeGenerateWalletResultFmt, uniqueKey)
+}

pkg/eventconsumer/keygen_consumer.go

@@ -9,6 +9,7 @@ import (
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/fystack/mpcium/pkg/mpc"
+	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -142,7 +143,10 @@ func (sc *keygenConsumer) handleKeygenEvent(msg jetstream.Msg) {
 	}()
 
 	// Publish the signing event with the reply inbox.
-	if err := sc.pubsub.PublishWithReply(MPCGenerateEvent, replyInbox, msg.Data()); err != nil {
+	headers := map[string]string{
+		"SessionID": uuid.New().String(),
+	}
+	if err := sc.pubsub.PublishWithReply(MPCGenerateEvent, replyInbox, msg.Data(), headers); err != nil {
 		logger.Error("KeygenConsumer: Failed to publish signing event with reply", err)
 		_ = msg.Nak()
 		return

pkg/eventconsumer/sign_consumer.go

@@ -9,6 +9,7 @@ import (
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/fystack/mpcium/pkg/mpc"
+	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 	"github.com/spf13/viper"
@@ -158,7 +159,10 @@ func (sc *signingConsumer) handleSigningEvent(msg jetstream.Msg) {
 	}()
 
 	// Publish the signing event with the reply inbox.
-	if err := sc.pubsub.PublishWithReply(MPCSignEvent, replyInbox, msg.Data()); err != nil {
+	headers := map[string]string{
+		"SessionID": uuid.New().String(),
+	}
+	if err := sc.pubsub.PublishWithReply(MPCSignEvent, replyInbox, msg.Data(), headers); err != nil {
 		logger.Error("SigningConsumer: Failed to publish signing event with reply", err)
 		_ = msg.Nak()
 		return

pkg/messaging/pubsub.go

@@ -11,7 +11,7 @@ type Subscription interface {
 
 type PubSub interface {
 	Publish(topic string, message []byte) error
-	PublishWithReply(topic, reply string, data []byte) error
+	PublishWithReply(topic, reply string, data []byte, headers map[string]string) error
 	Subscribe(topic string, handler func(msg *nats.Msg)) (Subscription, error)
 }
 
@@ -36,12 +36,18 @@ func (n *natsPubSub) Publish(topic string, message []byte) error {
 	return n.natsConn.Publish(topic, message)
 }
 
-func (n *natsPubSub) PublishWithReply(topic, reply string, data []byte) error {
-	return n.natsConn.PublishMsg(&nats.Msg{
+func (n *natsPubSub) PublishWithReply(topic, reply string, data []byte, headers map[string]string) error {
+	msg := &nats.Msg{
 		Subject: topic,
 		Reply:   reply,
 		Data:    data,
-	})
+		Header:  nats.Header{},
+	}
+	for k, v := range headers {
+		msg.Header.Set(k, v)
+	}
+	err := n.natsConn.PublishMsg(msg)
+	return err
 }
 
 func (n *natsPubSub) Subscribe(topic string, handler func(msg *nats.Msg)) (Subscription, error) {