feat: allowing queuing of long-running rewards generation

Layr-Labs · Jan 13, 2025 · d75bb67 · d75bb67
1 parent bcb4b12
commit d75bb67
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/pkg/rpcServer/rewardsHandlers.go b/pkg/rpcServer/rewardsHandlers.go
@@ -18,16 +18,26 @@ func (rpc *RpcServer) GetRewardsRoot(ctx context.Context, req *sidecarV1.GetRewa
 
 func (rpc *RpcServer) GenerateRewards(ctx context.Context, req *sidecarV1.GenerateRewardsRequest) (*sidecarV1.GenerateRewardsResponse, error) {
 	cutoffDate := req.GetCutoffDate()
+	waitForComplete := req.GetWaitForComplete()
 
 	var err error
 	queued := false
 	msg := rewardsCalculatorQueue.RewardsCalculationData{
 		CalculationType: rewardsCalculatorQueue.RewardsCalculationType_CalculateRewards,
 		CutoffDate:      cutoffDate,
 	}
-	data, qErr := rpc.rewardsQueue.EnqueueAndWait(ctx, msg)
-	cutoffDate = data.CutoffDate
-	err = qErr
+
+	if waitForComplete {
+		data, qErr := rpc.rewardsQueue.EnqueueAndWait(ctx, msg)
+		cutoffDate = data.CutoffDate
+		err = qErr
+	} else {
+		rpc.rewardsQueue.Enqueue(&rewardsCalculatorQueue.RewardsCalculationMessage{
+			Data:         msg,
+			ResponseChan: make(chan *rewardsCalculatorQueue.RewardsCalculatorResponse),
+		})
+		queued = true
+	}
 
 	if err != nil {
 		if errors.Is(err, &rewards.ErrRewardsCalculationInProgress{}) {