Merge pull request #43 from fystack/feat/improve-keygen-concurrency

Code refactoring, improve concurrency

README.md

@@ -148,7 +148,7 @@ func main () {
 		NatsConn: natsConn,
 		KeyPath:  "./event_initiator.key",
 	})
-	err = mpcClient.OnWalletCreationResult(func(event mpc.KeygenSuccessEvent) {
+	err = mpcClient.OnWalletCreationResult(func(event event.KeygenSuccessEvent) {
 		logger.Info("Received wallet creation result", "event", event)
 	})
 	if err != nil {

config.yaml.template

@@ -7,3 +7,4 @@ mpc_threshold: 2
 environment: development
 badger_password: "your_badger_password"
 event_initiator_pubkey: "event_initiator_pubkey"
+max_concurrent_keygen: 2

examples/generate/main.go

@@ -1,15 +1,20 @@
 package main
 
 import (
+	"encoding/json"
+	"flag"
 	"fmt"
 	"os"
 	"os/signal"
+	"sync"
+	"sync/atomic"
 	"syscall"
+	"time"
 
 	"github.com/fystack/mpcium/pkg/client"
 	"github.com/fystack/mpcium/pkg/config"
+	"github.com/fystack/mpcium/pkg/event"
 	"github.com/fystack/mpcium/pkg/logger"
-	"github.com/fystack/mpcium/pkg/mpc"
 	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
 	"github.com/spf13/viper"
@@ -17,6 +22,9 @@ import (
 
 func main() {
 	const environment = "development"
+	numWallets := flag.Int("n", 1, "Number of wallets to generate")
+	flag.Parse()
+
 	config.InitViperConfig()
 	logger.Init(environment, false)
 
@@ -25,25 +33,95 @@ func main() {
 	if err != nil {
 		logger.Fatal("Failed to connect to NATS", err)
 	}
-	defer natsConn.Drain() // drain inflight msgs
+	defer natsConn.Drain()
 	defer natsConn.Close()
 
 	mpcClient := client.NewMPCClient(client.Options{
 		NatsConn: natsConn,
 		KeyPath:  "./event_initiator.key",
 	})
-	err = mpcClient.OnWalletCreationResult(func(event mpc.KeygenSuccessEvent) {
-		logger.Info("Received wallet creation result", "event", event)
+
+	var walletStartTimes sync.Map
+	var walletIDs []string
+	var walletIDsMu sync.Mutex
+	var wg sync.WaitGroup
+	var completedCount int32
+
+	startAll := time.Now()
+
+	// STEP 1: Pre-generate wallet IDs and store start times
+	for i := 0; i < *numWallets; i++ {
+		walletID := uuid.New().String()
+		walletStartTimes.Store(walletID, time.Now())
+
+		walletIDsMu.Lock()
+		walletIDs = append(walletIDs, walletID)
+		walletIDsMu.Unlock()
+	}
+
+	// STEP 2: Register the result handler AFTER all walletIDs are stored
+	err = mpcClient.OnWalletCreationResult(func(event event.KeygenSuccessEvent) {
+		now := time.Now()
+		startTimeAny, ok := walletStartTimes.Load(event.WalletID)
+		if ok {
+			startTime := startTimeAny.(time.Time)
+			duration := now.Sub(startTime).Seconds()
+			accumulated := now.Sub(startAll).Seconds()
+			countSoFar := atomic.AddInt32(&completedCount, 1)
+
+			logger.Info("Wallet created",
+				"walletID", event.WalletID,
+				"duration_seconds", fmt.Sprintf("%.3f", duration),
+				"accumulated_time_seconds", fmt.Sprintf("%.3f", accumulated),
+				"count_so_far", countSoFar,
+			)
+
+			walletStartTimes.Delete(event.WalletID)
+		} else {
+			logger.Warn("Received wallet result but no start time found", "walletID", event.WalletID)
+		}
+		wg.Done()
 	})
 	if err != nil {
 		logger.Fatal("Failed to subscribe to wallet-creation results", err)
 	}
 
-	walletID := uuid.New().String()
-	if err := mpcClient.CreateWallet(walletID); err != nil {
-		logger.Fatal("CreateWallet failed", err)
+	// STEP 3: Create wallets
+	for _, walletID := range walletIDs {
+		wg.Add(1)
+		if err := mpcClient.CreateWallet(walletID); err != nil {
+			logger.Error("CreateWallet failed", err)
+			walletStartTimes.Delete(walletID)
+			wg.Done()
+			continue
+		}
+		logger.Info("CreateWallet sent, awaiting result...", "walletID", walletID)
 	}
-	logger.Info("CreateWallet sent, awaiting result...", "walletID", walletID)
+
+	// Wait until all wallet creations complete
+	go func() {
+		wg.Wait()
+		totalDuration := time.Since(startAll).Seconds()
+		logger.Info("All wallets generated", "count", completedCount, "total_duration_seconds", fmt.Sprintf("%.3f", totalDuration))
+
+		// Save wallet IDs to wallets.json
+		walletIDsMu.Lock()
+		data, err := json.MarshalIndent(walletIDs, "", "  ")
+		walletIDsMu.Unlock()
+		if err != nil {
+			logger.Error("Failed to marshal wallet IDs", err)
+		} else {
+			err = os.WriteFile("wallets.json", data, 0644)
+			if err != nil {
+				logger.Error("Failed to write wallets.json", err)
+			} else {
+				logger.Info("wallets.json written", "count", len(walletIDs))
+			}
+		}
+		os.Exit(0)
+	}()
+
+	// Block on SIGINT/SIGTERM (Ctrl+C etc.)
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
 	<-stop

examples/sign/main.go

@@ -40,7 +40,7 @@ func main() {
 
 	txMsg := &types.SignTxMessage{
 		KeyType:             types.KeyTypeEd25519,
-		WalletID:            "77dd7e23-9d5c-4ff1-8759-f119d1b19b36",
+		WalletID:            "c47cd6f4-8ef4-4d77-9d2b-37f9d062e615",
 		NetworkInternalCode: "solana-devnet",
 		TxID:                txID,
 		Tx:                  dummyTx,

pkg/client/client.go

@@ -15,7 +15,6 @@ import (
 	"github.com/fystack/mpcium/pkg/eventconsumer"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
-	"github.com/fystack/mpcium/pkg/mpc"
 	"github.com/fystack/mpcium/pkg/types"
 	"github.com/nats-io/nats.go"
 )
@@ -26,7 +25,7 @@ const (
 
 type MPCClient interface {
 	CreateWallet(walletID string) error
-	OnWalletCreationResult(callback func(event mpc.KeygenSuccessEvent)) error
+	OnWalletCreationResult(callback func(event event.KeygenSuccessEvent)) error
 
 	SignTransaction(msg *types.SignTxMessage) error
 	OnSignResult(callback func(event event.SigningResultEvent)) error
@@ -186,9 +185,9 @@ func (c *mpcClient) CreateWallet(walletID string) error {
 }
 
 // The callback will be invoked whenever a wallet creation result is received.
-func (c *mpcClient) OnWalletCreationResult(callback func(event mpc.KeygenSuccessEvent)) error {
+func (c *mpcClient) OnWalletCreationResult(callback func(event event.KeygenSuccessEvent)) error {
 	err := c.genKeySuccessQueue.Dequeue(GenerateWalletSuccessTopic, func(msg []byte) error {
-		var event mpc.KeygenSuccessEvent
+		var event event.KeygenSuccessEvent
 		err := json.Unmarshal(msg, &event)
 		if err != nil {
 			return err

pkg/eventconsumer/event_consumer.go

@@ -23,6 +23,9 @@ import (
 const (
 	MPCGenerateEvent = "mpc:generate"
 	MPCSignEvent     = "mpc:sign"
+
+	DefaultConcurrentKeygen     = 2
+	DefaultKeyGenStartupDelayMs = 500
 )
 
 type EventConsumer interface {
@@ -42,6 +45,9 @@ type eventConsumer struct {
 	signingSub       messaging.Subscription
 	identityStore    identity.Store
 
+	msgBuffer           chan *nats.Msg
+	maxConcurrentKeygen int
+
 	// Track active sessions with timestamps for cleanup
 	activeSessions  map[string]time.Time // Maps "walletID-txID" to creation time
 	sessionsLock    sync.RWMutex
@@ -57,19 +63,27 @@ func NewEventConsumer(
 	signingResultQueue messaging.MessageQueue,
 	identityStore identity.Store,
 ) EventConsumer {
-	ec := &eventConsumer{
-		node:               node,
-		pubsub:             pubsub,
-		genKeySucecssQueue: genKeySucecssQueue,
-		signingResultQueue: signingResultQueue,
-		activeSessions:     make(map[string]time.Time),
-		cleanupInterval:    5 * time.Minute,  // Run cleanup every 5 minutes
-		sessionTimeout:     30 * time.Minute, // Consider sessions older than 30 minutes stale
-		cleanupStopChan:    make(chan struct{}),
-		mpcThreshold:       viper.GetInt("mpc_threshold"),
-		identityStore:      identityStore,
+	maxConcurrentKeygen := viper.GetInt("max_concurrent_keygen")
+	if maxConcurrentKeygen == 0 {
+		maxConcurrentKeygen = DefaultConcurrentKeygen
 	}
 
+	ec := &eventConsumer{
+		node:                node,
+		pubsub:              pubsub,
+		genKeySucecssQueue:  genKeySucecssQueue,
+		signingResultQueue:  signingResultQueue,
+		activeSessions:      make(map[string]time.Time),
+		cleanupInterval:     5 * time.Minute,  // Run cleanup every 5 minutes
+		sessionTimeout:      30 * time.Minute, // Consider sessions older than 30 minutes stale
+		cleanupStopChan:     make(chan struct{}),
+		mpcThreshold:        viper.GetInt("mpc_threshold"),
+		maxConcurrentKeygen: maxConcurrentKeygen,
+		identityStore:       identityStore,
+		msgBuffer:           make(chan *nats.Msg, 100),
+	}
+
+	go ec.startKeyGenEventWorker()
 	// Start background cleanup goroutine
 	go ec.sessionCleanupRoutine()
 
@@ -90,100 +104,120 @@ func (ec *eventConsumer) Run() {
 	logger.Info("MPC Event consumer started...!")
 }
 
+func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
+	raw := natMsg.Data
+	var msg types.GenerateKeyMessage
+	err := json.Unmarshal(raw, &msg)
+	if err != nil {
+		logger.Error("Failed to unmarshal signing message", err)
+		return
+	}
+	err = ec.identityStore.VerifyInitiatorMessage(&msg)
+	if err != nil {
+		logger.Error("Failed to verify initiator message", err)
+		return
+	}
+
+	walletID := msg.WalletID
+	ecdsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeECDSA, walletID, ec.mpcThreshold, ec.genKeySucecssQueue)
+	if err != nil {
+		logger.Error("Failed to create key generation session", err, "walletID", walletID)
+		return
+	}
+	eddsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeEDDSA, walletID, ec.mpcThreshold, ec.genKeySucecssQueue)
+	if err != nil {
+		logger.Error("Failed to create key generation session", err, "walletID", walletID)
+		return
+	}
+
+	ecdsaSession.Init()
+	eddsaSession.Init()
+
+	ctx := context.Background()
+	ctxEcdsa, doneEcdsa := context.WithCancel(ctx)
+	ctxEddsa, doneEddsa := context.WithCancel(ctx)
+
+	successEvent := &event.KeygenSuccessEvent{
+		WalletID: walletID,
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		for {
+			select {
+			case <-ctxEcdsa.Done():
+				successEvent.ECDSAPubKey = ecdsaSession.GetPubKeyResult()
+				wg.Done()
+				return
+			case err := <-ecdsaSession.ErrChan():
+				logger.Error("Keygen session error", err)
+			}
+		}
+	}()
+
+	go func() {
+		for {
+			select {
+			case <-ctxEddsa.Done():
+				successEvent.EDDSAPubKey = eddsaSession.GetPubKeyResult()
+				wg.Done()
+				return
+			case err := <-eddsaSession.ErrChan():
+				logger.Error("Keygen session error", err)
+			}
+		}
+	}()
+
+	ecdsaSession.ListenToIncomingMessageAsync()
+	eddsaSession.ListenToIncomingMessageAsync()
+
+	// Temporary delay to allow peer nodes to subscribe and prepare before starting key generation.
+	// This should be replaced with a proper distributed coordination mechanism later (e.g., Consul lock).
+	time.Sleep(DefaultKeyGenStartupDelayMs * time.Millisecond)
+
+	go ecdsaSession.GenerateKey(doneEcdsa)
+	go eddsaSession.GenerateKey(doneEddsa)
+
+	wg.Wait()
+	logger.Info("Closing session successfully!", "event", successEvent)
+
+	successEventBytes, err := json.Marshal(successEvent)
+	if err != nil {
+		logger.Error("Failed to marshal keygen success event", err)
+		return
+	}
+
+	err = ec.genKeySucecssQueue.Enqueue(fmt.Sprintf(mpc.TypeGenerateWalletSuccess, walletID), successEventBytes, &messaging.EnqueueOptions{
+		IdempotententKey: fmt.Sprintf(mpc.TypeGenerateWalletSuccess, walletID),
+	})
+	if err != nil {
+		logger.Error("Failed to publish key generation success message", err)
+		return
+	}
+
+	logger.Info("[COMPLETED KEY GEN] Key generation completed successfully", "walletID", walletID)
+}
+
+func (ec *eventConsumer) startKeyGenEventWorker() {
+	// semaphore to limit concurrency
+	semaphore := make(chan struct{}, ec.maxConcurrentKeygen)
+
+	for {
+		select {
+		case natMsg := <-ec.msgBuffer:
+			semaphore <- struct{}{} // acquire a slot
+			go func(msg *nats.Msg) {
+				defer func() { <-semaphore }() // release the slot when done
+				ec.handleKeyGenEvent(msg)
+			}(natMsg)
+		}
+	}
+}
+
 func (ec *eventConsumer) consumeKeyGenerationEvent() error {
 	sub, err := ec.pubsub.Subscribe(MPCGenerateEvent, func(natMsg *nats.Msg) {
-		raw := natMsg.Data
-		var msg types.GenerateKeyMessage
-		err := json.Unmarshal(raw, &msg)
-		if err != nil {
-			logger.Error("Failed to unmarshal signing message", err)
-			return
-		}
-		logger.Info("Received key generation event", "msg", msg)
-
-		err = ec.identityStore.VerifyInitiatorMessage(&msg)
-		if err != nil {
-			logger.Error("Failed to verify initiator message", err)
-			return
-		}
-
-		walletID := msg.WalletID
-		session, err := ec.node.CreateKeyGenSession(walletID, ec.mpcThreshold, ec.genKeySucecssQueue)
-		if err != nil {
-			logger.Error("Failed to create key generation session", err, "walletID", walletID)
-			return
-		}
-		eddsaSession, err := ec.node.CreateEDDSAKeyGenSession(walletID, ec.mpcThreshold, ec.genKeySucecssQueue)
-		if err != nil {
-			logger.Error("Failed to create key generation session", err, "walletID", walletID)
-			return
-		}
-
-		session.Init()
-		eddsaSession.Init()
-
-		ctx, done := context.WithCancel(context.Background())
-		ctxEddsa, doneEddsa := context.WithCancel(context.Background())
-
-		successEvent := &mpc.KeygenSuccessEvent{
-			WalletID: walletID,
-		}
-
-		var wg sync.WaitGroup
-		wg.Add(2)
-		go func() {
-			for {
-				select {
-				case <-ctx.Done():
-					successEvent.ECDSAPubKey = session.GetPubKeyResult()
-					wg.Done()
-					return
-				case err := <-session.ErrCh:
-					logger.Error("Keygen session error", err)
-				}
-			}
-		}()
-
-		go func() {
-			for {
-				select {
-				case <-ctxEddsa.Done():
-					successEvent.EDDSAPubKey = eddsaSession.GetPubKeyResult()
-					wg.Done()
-					return
-				case err := <-eddsaSession.ErrCh:
-					logger.Error("Keygen session error", err)
-				}
-			}
-		}()
-
-		session.ListenToIncomingMessageAsync()
-		eddsaSession.ListenToIncomingMessageAsync()
-		// TODO: replace sleep with distributed lock
-		time.Sleep(1 * time.Second)
-
-		go session.GenerateKey(done)
-		go eddsaSession.GenerateKey(doneEddsa)
-
-		wg.Wait()
-		logger.Info("Closing session successfully!", "event", successEvent)
-
-		successEventBytes, err := json.Marshal(successEvent)
-		if err != nil {
-			logger.Error("Failed to marshal keygen success event", err)
-			return
-		}
-
-		err = ec.genKeySucecssQueue.Enqueue(fmt.Sprintf(mpc.TypeGenerateWalletSuccess, walletID), successEventBytes, &messaging.EnqueueOptions{
-			IdempotententKey: fmt.Sprintf(mpc.TypeGenerateWalletSuccess, walletID),
-		})
-		if err != nil {
-			logger.Error("Failed to publish key generation success message", err)
-			return
-		}
-
-		logger.Info("[COMPLETED KEY GEN] Key generation completed successfully", "walletID", walletID)
-
+		ec.msgBuffer <- natMsg
 	})
 
 	ec.keyGenerationSub = sub
@@ -227,10 +261,11 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
 			return
 		}
 
-		var session mpc.ISigningSession
+		var session mpc.SigningSession
 		switch msg.KeyType {
 		case types.KeyTypeSecp256k1:
 			session, err = ec.node.CreateSigningSession(
+				mpc.SessionTypeECDSA,
 				msg.WalletID,
 				msg.TxID,
 				msg.NetworkInternalCode,
@@ -238,7 +273,8 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
 				ec.signingResultQueue,
 			)
 		case types.KeyTypeEd25519:
-			session, err = ec.node.CreateEDDSASigningSession(
+			session, err = ec.node.CreateSigningSession(
+				mpc.SessionTypeECDSA,
 				msg.WalletID,
 				msg.TxID,
 				msg.NetworkInternalCode,

pkg/kvstore/badger.go

@@ -3,9 +3,9 @@ package kvstore
 import (
 	"errors"
 
-	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/dgraph-io/badger/v4"
 	"github.com/dgraph-io/badger/v4/options"
+	"github.com/fystack/mpcium/pkg/logger"
 )
 
 var (

pkg/mpc/ecdsa_keygen_session.go

@@ -15,22 +15,20 @@ import (
 	"github.com/fystack/mpcium/pkg/messaging"
 )
 
-const (
-	TypeGenerateWalletSuccess = "mpc.mpc_keygen_success.%s"
-)
-
-type KeygenSession struct {
+type KeyGenSession interface {
 	Session
+
+	Init()
+	GenerateKey(done func())
+	GetPubKeyResult() []byte
+}
+
+type ecdsaKeygenSession struct {
+	session
 	endCh chan *keygen.LocalPartySaveData
 }
 
-type KeygenSuccessEvent struct {
-	WalletID    string `json:"wallet_id"`
-	ECDSAPubKey []byte `json:"ecdsa_pub_key"`
-	EDDSAPubKey []byte `json:"eddsa_pub_key"`
-}
-
-func NewKeygenSession(
+func newECDSAKeygenSession(
 	walletID string,
 	pubSub messaging.PubSub,
 	direct messaging.DirectMessaging,
@@ -43,9 +41,9 @@ func NewKeygenSession(
 	keyinfoStore keyinfo.Store,
 	resultQueue messaging.MessageQueue,
 	identityStore identity.Store,
-) *KeygenSession {
-	return &KeygenSession{
-		Session: Session{
+) *ecdsaKeygenSession {
+	return &ecdsaKeygenSession{
+		session: session{
 			walletID:           walletID,
 			pubSub:             pubSub,
 			direct:             direct,
@@ -71,14 +69,14 @@ func NewKeygenSession(
 			},
 			getRoundFunc:  GetEcdsaMsgRound,
 			resultQueue:   resultQueue,
-			sessionType:   SessionTypeEcdsa,
+			sessionType:   SessionTypeECDSA,
 			identityStore: identityStore,
 		},
 		endCh: make(chan *keygen.LocalPartySaveData),
 	}
 }
 
-func (s *KeygenSession) Init() {
+func (s *ecdsaKeygenSession) Init() {
 	logger.Infof("Initializing session with partyID: %s, peerIDs %s", s.selfPartyID, s.partyIDs)
 	ctx := tss.NewPeerContext(s.partyIDs)
 	params := tss.NewParameters(tss.S256(), ctx, s.selfPartyID, len(s.partyIDs), s.threshold)
@@ -86,8 +84,8 @@ func (s *KeygenSession) Init() {
 	logger.Infof("[INITIALIZED] Initialized session successfully partyID: %s, peerIDs %s, walletID %s, threshold = %d", s.selfPartyID, s.partyIDs, s.walletID, s.threshold)
 }
 
-func (s *KeygenSession) GenerateKey(done func()) {
-	logger.Info("Starting to generate key", "walletID", s.walletID)
+func (s *ecdsaKeygenSession) GenerateKey(done func()) {
+	logger.Info("Starting to generate key ECDSA", "walletID", s.walletID)
 	go func() {
 		if err := s.party.Start(); err != nil {
 			s.ErrCh <- err
@@ -139,12 +137,11 @@ func (s *KeygenSession) GenerateKey(done func()) {
 				return
 			}
 			s.pubkeyBytes = pubKeyBytes
-			done()
 			err = s.Close()
 			if err != nil {
 				logger.Error("Failed to close session", err)
 			}
-			// done()
+			done()
 			return
 		}
 	}

pkg/mpc/ecdsa_resharing_session.go

@@ -0,0 +1 @@
+package mpc

pkg/mpc/ecdsa_signing_session.go

@@ -20,9 +20,16 @@ import (
 	"github.com/samber/lo"
 )
 
-// Ecdsa signing session
-type SigningSession struct {
+type SigningSession interface {
 	Session
+
+	Init(tx *big.Int) error
+	Sign(onSuccess func(data []byte))
+}
+
+// Ecdsa signing session
+type ecdsaSigningSession struct {
+	session
 	endCh               chan *common.SignatureData
 	data                *keygen.LocalPartySaveData
 	tx                  *big.Int
@@ -30,19 +37,7 @@ type SigningSession struct {
 	networkInternalCode string
 }
 
-type ISession interface {
-	ErrChan() <-chan error
-	ListenToIncomingMessageAsync()
-}
-
-type ISigningSession interface {
-	ISession
-
-	Init(tx *big.Int) error
-	Sign(onSuccess func(data []byte))
-}
-
-func NewSigningSession(
+func newECDSASigningSession(
 	walletID string,
 	txID string,
 	networkInternalCode string,
@@ -57,9 +52,9 @@ func NewSigningSession(
 	keyinfoStore keyinfo.Store,
 	resultQueue messaging.MessageQueue,
 	identityStore identity.Store,
-) *SigningSession {
-	return &SigningSession{
-		Session: Session{
+) *ecdsaSigningSession {
+	return &ecdsaSigningSession{
+		session: session{
 			walletID:           walletID,
 			pubSub:             pubSub,
 			direct:             direct,
@@ -93,7 +88,7 @@ func NewSigningSession(
 	}
 }
 
-func (s *SigningSession) Init(tx *big.Int) error {
+func (s *ecdsaSigningSession) Init(tx *big.Int) error {
 	logger.Infof("Initializing signing session with partyID: %s, peerIDs %s", s.selfPartyID, s.partyIDs)
 	ctx := tss.NewPeerContext(s.partyIDs)
 	params := tss.NewParameters(tss.S256(), ctx, s.selfPartyID, len(s.partyIDs), s.threshold)
@@ -138,7 +133,7 @@ func (s *SigningSession) Init(tx *big.Int) error {
 	return nil
 }
 
-func (s *SigningSession) Sign(onSuccess func(data []byte)) {
+func (s *ecdsaSigningSession) Sign(onSuccess func(data []byte)) {
 	logger.Info("Starting signing", "walletID", s.walletID)
 	go func() {
 		if err := s.party.Start(); err != nil {

pkg/mpc/eddsa_keygen_session.go

@@ -14,17 +14,12 @@ import (
 	"github.com/fystack/mpcium/pkg/messaging"
 )
 
-type EDDSAKeygenSession struct {
-	Session
+type eddsaKeygenSession struct {
+	session
 	endCh chan *keygen.LocalPartySaveData
 }
 
-type EDDSAKeygenSuccessEvent struct {
-	WalletID string `json:"wallet_id"`
-	PubKey   []byte `json:"pub_key"`
-}
-
-func NewEDDSAKeygenSession(
+func newEDDSAKeygenSession(
 	walletID string,
 	pubSub messaging.PubSub,
 	direct messaging.DirectMessaging,
@@ -36,8 +31,8 @@ func NewEDDSAKeygenSession(
 	keyinfoStore keyinfo.Store,
 	resultQueue messaging.MessageQueue,
 	identityStore identity.Store,
-) *EDDSAKeygenSession {
-	return &EDDSAKeygenSession{Session: Session{
+) *eddsaKeygenSession {
+	return &eddsaKeygenSession{session: session{
 		walletID:           walletID,
 		pubSub:             pubSub,
 		direct:             direct,
@@ -62,14 +57,14 @@ func NewEDDSAKeygenSession(
 		},
 		getRoundFunc:  GetEddsaMsgRound,
 		resultQueue:   resultQueue,
-		sessionType:   SessionTypeEddsa,
+		sessionType:   SessionTypeEDDSA,
 		identityStore: identityStore,
 	},
 		endCh: make(chan *keygen.LocalPartySaveData),
 	}
 }
 
-func (s *EDDSAKeygenSession) Init() {
+func (s *eddsaKeygenSession) Init() {
 	logger.Infof("Initializing session with partyID: %s, peerIDs %s", s.selfPartyID, s.partyIDs)
 	ctx := tss.NewPeerContext(s.partyIDs)
 	params := tss.NewParameters(tss.Edwards(), ctx, s.selfPartyID, len(s.partyIDs), s.threshold)
@@ -77,8 +72,8 @@ func (s *EDDSAKeygenSession) Init() {
 	logger.Infof("[INITIALIZED] Initialized session successfully partyID: %s, peerIDs %s, walletID %s, threshold = %d", s.selfPartyID, s.partyIDs, s.walletID, s.threshold)
 }
 
-func (s *EDDSAKeygenSession) GenerateKey(done func()) {
-	logger.Info("Starting to generate key", "walletID", s.walletID)
+func (s *eddsaKeygenSession) GenerateKey(done func()) {
+	logger.Info("Starting to generate key EDDSA", "walletID", s.walletID)
 	go func() {
 		if err := s.party.Start(); err != nil {
 			s.ErrCh <- err

pkg/mpc/eddsa_resharing_session.go

@@ -0,0 +1 @@
+package mpc

pkg/mpc/eddsa_signing_session.go

@@ -9,6 +9,7 @@ import (
 	"github.com/bnb-chain/tss-lib/v2/eddsa/keygen"
 	"github.com/bnb-chain/tss-lib/v2/eddsa/signing"
 	"github.com/bnb-chain/tss-lib/v2/tss"
+	"github.com/decred/dcrd/dcrec/edwards/v2"
 	"github.com/fystack/mpcium/pkg/common/errors"
 	"github.com/fystack/mpcium/pkg/event"
 	"github.com/fystack/mpcium/pkg/identity"
@@ -16,12 +17,11 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
-	"github.com/decred/dcrd/dcrec/edwards/v2"
 	"github.com/samber/lo"
 )
 
-type EDDSASigningSession struct {
-	Session
+type eddsaSigningSession struct {
+	session
 	endCh               chan *common.SignatureData
 	data                *keygen.LocalPartySaveData
 	tx                  *big.Int
@@ -43,9 +43,9 @@ func NewEDDSASigningSession(
 	keyinfoStore keyinfo.Store,
 	resultQueue messaging.MessageQueue,
 	identityStore identity.Store,
-) *EDDSASigningSession {
-	return &EDDSASigningSession{
-		Session: Session{
+) *eddsaSigningSession {
+	return &eddsaSigningSession{
+		session: session{
 			walletID:           walletID,
 			pubSub:             pubSub,
 			direct:             direct,
@@ -79,7 +79,7 @@ func NewEDDSASigningSession(
 	}
 }
 
-func (s *EDDSASigningSession) Init(tx *big.Int) error {
+func (s *eddsaSigningSession) Init(tx *big.Int) error {
 	logger.Infof("Initializing signing session with partyID: %s, peerIDs %s", s.selfPartyID, s.partyIDs)
 	ctx := tss.NewPeerContext(s.partyIDs)
 	params := tss.NewParameters(tss.Edwards(), ctx, s.selfPartyID, len(s.partyIDs), s.threshold)
@@ -124,7 +124,7 @@ func (s *EDDSASigningSession) Init(tx *big.Int) error {
 	return nil
 }
 
-func (s *EDDSASigningSession) Sign(onSuccess func(data []byte)) {
+func (s *eddsaSigningSession) Sign(onSuccess func(data []byte)) {
 	logger.Info("Starting signing", "walletID", s.walletID)
 	go func() {
 		if err := s.party.Start(); err != nil {

pkg/mpc/node.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/bnb-chain/tss-lib/v2/ecdsa/keygen"
 	"github.com/bnb-chain/tss-lib/v2/tss"
+	"github.com/fystack/mpcium/pkg/common/errors"
 	"github.com/fystack/mpcium/pkg/identity"
 	"github.com/fystack/mpcium/pkg/keyinfo"
 	"github.com/fystack/mpcium/pkg/kvstore"
@@ -65,11 +66,13 @@ func NewNode(
 	peerRegistry PeerRegistry,
 	identityStore identity.Store,
 ) *Node {
+	start := time.Now()
 	preParams, err := keygen.GeneratePreParams(5 * time.Minute)
 	if err != nil {
 		logger.Fatal("Generate pre params failed", err)
 	}
-	logger.Info("Starting new node, preparams is generated successfully!")
+	elapsed := time.Since(start)
+	logger.Info("Starting new node, preparams is generated successfully!", "elapsed", elapsed.Milliseconds())
 
 	go peerRegistry.WatchPeersReady()
 
@@ -90,18 +93,30 @@ func (p *Node) ID() string {
 	return p.nodeID
 }
 
-func composeReadyTopic(nodeID string) string {
-	return fmt.Sprintf("%s-%s", nodeID, "ready")
-}
-
-func (p *Node) CreateKeyGenSession(walletID string, threshold int, successQueue messaging.MessageQueue) (*KeygenSession, error) {
-	if p.peerRegistry.GetReadyPeersCount() < int64(threshold+1) {
+func (p *Node) CreateKeyGenSession(
+	sessionType SessionType,
+	walletID string,
+	threshold int,
+	successQueue messaging.MessageQueue,
+) (KeyGenSession, error) {
+	if !p.peerRegistry.ArePeersReady() {
 		return nil, fmt.Errorf("Not enough peers to create gen session! Expected %d, got %d", threshold+1, p.peerRegistry.GetReadyPeersCount())
 	}
 
+	switch sessionType {
+	case SessionTypeECDSA:
+		return p.createECDSAKeyGenSession(walletID, threshold, successQueue)
+	case SessionTypeEDDSA:
+		return p.createEDDSAKeyGenSession(walletID, threshold, successQueue)
+	default:
+		return nil, fmt.Errorf("Unknown session type: %s", sessionType)
+	}
+}
+
+func (p *Node) createECDSAKeyGenSession(walletID string, threshold int, successQueue messaging.MessageQueue) (KeyGenSession, error) {
 	readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
 	selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs)
-	session := NewKeygenSession(
+	session := newECDSAKeygenSession(
 		walletID,
 		p.pubSub,
 		p.direct,
@@ -118,14 +133,10 @@ func (p *Node) CreateKeyGenSession(walletID string, threshold int, successQueue
 	return session, nil
 }
 
-func (p *Node) CreateEDDSAKeyGenSession(walletID string, threshold int, successQueue messaging.MessageQueue) (*EDDSAKeygenSession, error) {
-	if p.peerRegistry.GetReadyPeersCount() < int64(threshold+1) {
-		return nil, fmt.Errorf("Not enough peers to create gen session! Expected %d, got %d", threshold+1, p.peerRegistry.GetReadyPeersCount())
-	}
-
+func (p *Node) createEDDSAKeyGenSession(walletID string, threshold int, successQueue messaging.MessageQueue) (KeyGenSession, error) {
 	readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
 	selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs)
-	session := NewEDDSAKeygenSession(
+	session := newEDDSAKeygenSession(
 		walletID,
 		p.pubSub,
 		p.direct,
@@ -142,58 +153,52 @@ func (p *Node) CreateEDDSAKeyGenSession(walletID string, threshold int, successQ
 }
 
 func (p *Node) CreateSigningSession(
+	sessionType SessionType,
 	walletID string,
 	txID string,
 	networkInternalCode string,
 	threshold int,
 	resultQueue messaging.MessageQueue,
-) (*SigningSession, error) {
+) (SigningSession, error) {
 	readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
 	selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs)
-	session := NewSigningSession(
-		walletID,
-		txID,
-		networkInternalCode,
-		p.pubSub,
-		p.direct,
-		readyPeerIDs,
-		selfPartyID,
-		allPartyIDs,
-		threshold,
-		p.ecdsaPreParams,
-		p.kvstore,
-		p.keyinfoStore,
-		resultQueue,
-		p.identityStore,
-	)
-	return session, nil
-}
+	switch sessionType {
+	case SessionTypeECDSA:
+		return newECDSASigningSession(
+			walletID,
+			txID,
+			networkInternalCode,
+			p.pubSub,
+			p.direct,
+			readyPeerIDs,
+			selfPartyID,
+			allPartyIDs,
+			threshold,
+			p.ecdsaPreParams,
+			p.kvstore,
+			p.keyinfoStore,
+			resultQueue,
+			p.identityStore,
+		), nil
+	case SessionTypeEDDSA:
+		return NewEDDSASigningSession(
+			walletID,
+			txID,
+			networkInternalCode,
+			p.pubSub,
+			p.direct,
+			readyPeerIDs,
+			selfPartyID,
+			allPartyIDs,
+			threshold,
+			p.kvstore,
+			p.keyinfoStore,
+			resultQueue,
+			p.identityStore,
+		), nil
+	}
 
-func (p *Node) CreateEDDSASigningSession(
-	walletID string,
-	txID string,
-	networkInternalCode string,
-	threshold int,
-	resultQueue messaging.MessageQueue,
-) (*EDDSASigningSession, error) {
-	readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
-	selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs)
-	session := NewEDDSASigningSession(
-		walletID,
-		txID,
-		networkInternalCode,
-		p.pubSub,
-		p.direct,
-		readyPeerIDs,
-		selfPartyID,
-		allPartyIDs,
-		threshold,
-		p.kvstore,
-		p.keyinfoStore,
-		resultQueue,
-		p.identityStore,
-	)
-	return session, nil
+	return nil, errors.New("Unknown session type")
 }
 
 func (p *Node) generatePartyIDs(purpose string, readyPeerIDs []string) (self *tss.PartyID, all []*tss.PartyID) {

pkg/mpc/registry.go

@@ -24,6 +24,7 @@ type PeerRegistry interface {
 	Resign() error
 	GetReadyPeersCount() int64
 	GetReadyPeersIncludeSelf() []string // get ready peers include self
+	GetTotalPeersCount() int64
 }
 
 type registry struct {
@@ -195,6 +196,11 @@ func (r *registry) ArePeersReady() bool {
 	return r.ready
 }
 
+func (r *registry) GetTotalPeersCount() int64 {
+	var self int64 = 1
+	return int64(len(r.peerNodeIDs)) + self
+}
+
 func (r *registry) Resign() error {
 	k := r.readyKey(r.nodeID)
 

pkg/mpc/session.go

@@ -17,6 +17,14 @@ import (
 	"github.com/nats-io/nats.go"
 )
 
+type SessionType string
+
+const (
+	TypeGenerateWalletSuccess             = "mpc.mpc_keygen_success.%s"
+	SessionTypeECDSA          SessionType = "session_ecdsa"
+	SessionTypeEDDSA          SessionType = "session_eddsa"
+)
+
 var (
 	ErrNotEnoughParticipants = errors.New("Not enough participants to sign")
 )
@@ -28,14 +36,12 @@ type TopicComposer struct {
 
 type KeyComposerFn func(id string) string
 
-type SessionType string
+type Session interface {
+	ListenToIncomingMessageAsync()
+	ErrChan() <-chan error
+}
 
-const (
-	SessionTypeEcdsa SessionType = "session_ecdsa"
-	SessionTypeEddsa SessionType = "session_eddsa"
-)
-
-type Session struct {
+type session struct {
 	walletID           string
 	pubSub             messaging.PubSub
 	direct             messaging.DirectMessaging
@@ -66,19 +72,19 @@ type Session struct {
 	sessionType SessionType
 }
 
-func (s *Session) PartyID() *tss.PartyID {
+func (s *session) PartyID() *tss.PartyID {
 	return s.selfPartyID
 }
 
-func (s *Session) PartyIDs() []*tss.PartyID {
+func (s *session) PartyIDs() []*tss.PartyID {
 	return s.partyIDs
 }
 
-func (s *Session) PartyCount() int {
+func (s *session) PartyCount() int {
 	return len(s.partyIDs)
 }
 
-func (s *Session) handleTssMessage(keyshare tss.Message) {
+func (s *session) handleTssMessage(keyshare tss.Message) {
 	data, routing, err := keyshare.WireBytes()
 	if err != nil {
 		s.ErrCh <- err
@@ -118,7 +124,7 @@ func (s *Session) handleTssMessage(keyshare tss.Message) {
 	}
 }
 
-func (s *Session) receiveTssMessage(rawMsg []byte) {
+func (s *session) receiveTssMessage(rawMsg []byte) {
 	msg, err := types.UnmarshalTssMessage(rawMsg)
 	if err != nil {
 		s.ErrCh <- fmt.Errorf("Failed to unmarshal message: %w", err)
@@ -157,7 +163,7 @@ func (s *Session) receiveTssMessage(rawMsg []byte) {
 	}
 }
 
-func (s *Session) SendReplySignSuccess(natMsg *nats.Msg) {
+func (s *session) SendReplySignSuccess(natMsg *nats.Msg) {
 	msg := natMsg.Data
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -170,7 +176,7 @@ func (s *Session) SendReplySignSuccess(natMsg *nats.Msg) {
 	logger.Info("Sent reply sign sucess message", "reply", natMsg.Reply)
 }
 
-func (s *Session) ListenToIncomingMessageAsync() {
+func (s *session) ListenToIncomingMessageAsync() {
 	go func() {
 		sub, err := s.pubSub.Subscribe(s.topicComposer.ComposeBroadcastTopic(), func(natMsg *nats.Msg) {
 			msg := natMsg.Data
@@ -197,7 +203,7 @@ func (s *Session) ListenToIncomingMessageAsync() {
 
 }
 
-func (s *Session) Close() error {
+func (s *session) Close() error {
 	err := s.broadcastSub.Unsubscribe()
 	if err != nil {
 		return err
@@ -209,10 +215,10 @@ func (s *Session) Close() error {
 	return nil
 }
 
-func (s *Session) GetPubKeyResult() []byte {
+func (s *session) GetPubKeyResult() []byte {
 	return s.pubkeyBytes
 }
 
-func (s *Session) ErrChan() <-chan error {
+func (s *session) ErrChan() <-chan error {
 	return s.ErrCh
 }