From b5c502398e1a694d4798fc598aa61da60a2161be Mon Sep 17 00:00:00 2001
From: Ran Mishael <ran@lavanet.xyz>
Date: Sun, 24 Mar 2024 17:38:29 +0100
Subject: [PATCH] adding wait groups for reward server claims

---
 protocol/rpcprovider/rewardserver/reward_server.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/protocol/rpcprovider/rewardserver/reward_server.go b/protocol/rpcprovider/rewardserver/reward_server.go
index 295947492f..2f56f93ac2 100644
--- a/protocol/rpcprovider/rewardserver/reward_server.go
+++ b/protocol/rpcprovider/rewardserver/reward_server.go
@@ -265,6 +265,9 @@ func (rws *RewardServer) sendRewardsClaim(ctx context.Context, epoch uint64) err
 		},
 	}
 
+	paymentWaitGroup := sync.WaitGroup{}
+	paymentWaitGroup.Add(newRewardsLength + failedRewardsLength)
+
 	// add expected pay and ask for rewards
 	for _, paymentConfig := range paymentConfiguration {
 		for _, rewardsToClaim := range paymentConfig.relaySessionChunks {
@@ -277,6 +280,7 @@ func (rws *RewardServer) sendRewardsClaim(ctx context.Context, epoch uint64) err
 				continue
 			}
 			go func(rewards []*pairingtypes.RelaySession, payment *PaymentConfiguration) { // send rewards asynchronously
+				defer paymentWaitGroup.Done()
 				specs := map[string]struct{}{}
 				if payment.shouldAddExpectedPayment {
 					for _, relay := range rewards {
@@ -319,6 +323,8 @@ func (rws *RewardServer) sendRewardsClaim(ctx context.Context, epoch uint64) err
 			}(rewardsToClaim, paymentConfig)
 		}
 	}
+	utils.LavaFormatDebug("Waiting for all Payment groups to finish", utils.LogAttr("wait_group_size", newRewardsLength+failedRewardsLength))
+	paymentWaitGroup.Wait()
 	return nil
 }