From 96702e278bbd3788b6a2a2884e9b2a91e71f1bb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?rich=CE=9Brd?= <info@richardramos.me>
Date: Fri, 22 Nov 2024 09:36:02 -0400
Subject: [PATCH] fix: setup shards (#1256)

---
 cmd/waku/node.go          | 21 ++++++++++++++++-----
 waku/v2/node/wakunode2.go | 18 +++++++++++-------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/cmd/waku/node.go b/cmd/waku/node.go
index 138416bc..61fdf4fb 100644
--- a/cmd/waku/node.go
+++ b/cmd/waku/node.go
@@ -309,11 +309,6 @@ func Execute(options NodeOptions) error {
 		nodeOpts = append(nodeOpts, discv5Opts)
 	}
 
-	wakuNode, err := node.New(nodeOpts...)
-	if err != nil {
-		return fmt.Errorf("could not instantiate waku: %w", err)
-	}
-
 	//Process pubSub and contentTopics specified and arrive at all corresponding pubSubTopics
 	pubSubTopicMap, err := processTopics(options)
 	if err != nil {
@@ -325,6 +320,22 @@ func Execute(options NodeOptions) error {
 		pubSubTopicMapKeys = append(pubSubTopicMapKeys, k)
 	}
 
+	rs, err := wprotocol.TopicsToRelayShards(pubSubTopicMapKeys...)
+	if err == nil {
+		if len(rs) == 1 {
+			nodeOpts = append(nodeOpts, node.WithShards(rs[0].ShardIDs))
+		} else {
+			logger.Warn("could not set ENR shard info", zap.String("error", "invalid number of clusters found"), zap.Int("numClusters", len(rs)))
+		}
+	} else {
+		logger.Warn("could not obtain list of shards", zap.Error(err))
+	}
+
+	wakuNode, err := node.New(nodeOpts...)
+	if err != nil {
+		return fmt.Errorf("could not instantiate waku: %w", err)
+	}
+
 	if err = wakuNode.Start(ctx); err != nil {
 		return nonRecoverError(err)
 	}
diff --git a/waku/v2/node/wakunode2.go b/waku/v2/node/wakunode2.go
index 2de62cc7..1ae5b244 100644
--- a/waku/v2/node/wakunode2.go
+++ b/waku/v2/node/wakunode2.go
@@ -464,13 +464,17 @@ func (w *WakuNode) Start(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		//TODO: setting this up temporarily to improve connectivity success for lightNode in status.
-		//This will have to be removed or changed with community sharding will be implemented.
-		if w.opts.shards != nil {
-			err = w.SetRelayShards(*w.opts.shards)
-			if err != nil {
-				return err
-			}
+	}
+
+	//TODO: setting this up temporarily to improve connectivity success for lightNode
+	//      in status. Also, when executing go-waku service-node as a lightclient
+	//      (using --pubsub-topic and --relay=false)
+	//      This will have to be removed or changed with community sharding will be
+	//      implemented.
+	if w.opts.shards != nil {
+		err = w.SetRelayShards(*w.opts.shards)
+		if err != nil {
+			return err
 		}
 	}