From d6f5cc651367dc16177c199cbbc09bab53557b31 Mon Sep 17 00:00:00 2001
From: anhthii <nguyendanganhthi247@gmail.com>
Date: Sun, 2 Nov 2025 13:29:21 +0700
Subject: [PATCH] Gracefully handle context cancelled

---
 cmd/mpcium/main.go                   | 16 ++++++++++++----
 pkg/eventconsumer/keygen_consumer.go |  5 +++++
 pkg/eventconsumer/sign_consumer.go   |  5 +++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/cmd/mpcium/main.go b/cmd/mpcium/main.go
index c8c10de0..95ae4224 100644
--- a/cmd/mpcium/main.go
+++ b/cmd/mpcium/main.go
@@ -277,8 +277,12 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	go func() {
 		defer wg.Done()
 		if err := keygenConsumer.Run(appContext); err != nil {
-			logger.Error("error running keygen consumer", err)
-			errChan <- fmt.Errorf("keygen consumer error: %w", err)
+			if appContext.Err() != context.Canceled {
+				logger.Error("error running keygen consumer", err)
+				errChan <- fmt.Errorf("keygen consumer error: %w", err)
+			} else {
+				logger.Info("Keygen consumer finished successfully")
+			}
 			return
 		}
 		logger.Info("Keygen consumer finished successfully")
@@ -288,8 +292,12 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	go func() {
 		defer wg.Done()
 		if err := signingConsumer.Run(appContext); err != nil {
-			logger.Error("error running signing consumer", err)
-			errChan <- fmt.Errorf("signing consumer error: %w", err)
+			if appContext.Err() != context.Canceled {
+				logger.Error("error running signing consumer", err)
+				errChan <- fmt.Errorf("signing consumer error: %w", err)
+			} else {
+				logger.Info("Signing consumer finished successfully")
+			}
 			return
 		}
 		logger.Info("Signing consumer finished successfully")
diff --git a/pkg/eventconsumer/keygen_consumer.go b/pkg/eventconsumer/keygen_consumer.go
index 98d9f072..8af027d0 100644
--- a/pkg/eventconsumer/keygen_consumer.go
+++ b/pkg/eventconsumer/keygen_consumer.go
@@ -72,6 +72,7 @@ func (sc *keygenConsumer) waitForAllPeersReadyToGenKey(ctx context.Context) erro
 		select {
 		case <-ctx.Done():
 			if ctx.Err() == context.Canceled {
+				logger.Info("KeygenConsumer: Shutdown signal received during peer readiness wait")
 				return nil
 			}
 			return ctx.Err()
@@ -107,6 +108,10 @@ func (sc *keygenConsumer) Run(ctx context.Context) error {
 		sc.handleKeygenEvent,
 	)
 	if err != nil {
+		if ctx.Err() == context.Canceled {
+			logger.Info("KeygenConsumer: Shutdown during subscription setup")
+			return nil
+		}
 		return fmt.Errorf("failed to subscribe to keygen events: %w", err)
 	}
 	sc.jsSub = sub
diff --git a/pkg/eventconsumer/sign_consumer.go b/pkg/eventconsumer/sign_consumer.go
index 24f21fec..e0ea23e2 100644
--- a/pkg/eventconsumer/sign_consumer.go
+++ b/pkg/eventconsumer/sign_consumer.go
@@ -75,6 +75,7 @@ func (sc *signingConsumer) waitForSufficientPeers(ctx context.Context) error {
 		select {
 		case <-ctx.Done():
 			if ctx.Err() == context.Canceled {
+				logger.Info("SigningConsumer: Shutdown signal received during peer readiness wait")
 				return nil
 			}
 			return ctx.Err()
@@ -110,6 +111,10 @@ func (sc *signingConsumer) Run(ctx context.Context) error {
 		sc.handleSigningEvent,
 	)
 	if err != nil {
+		if ctx.Err() == context.Canceled {
+			logger.Info("SigningConsumer: Shutdown during subscription setup")
+			return nil
+		}
 		return fmt.Errorf("failed to subscribe to signing events: %w", err)
 	}
 	sc.jsSub = sub