fix: rebalance issue

- risk of int overflow - rebalance by offset and partition
housepower · May 6, 2024 · 0c42c38 · 0c42c38
1 parent d379713
commit 0c42c38
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/task/task.go b/task/task.go
@@ -209,7 +209,7 @@ func (service *Service) Put(msg *model.InputMessage, traceId string, flushFn fun
 				util.Logger.Fatal("shard number calculation failed", zap.String("task", taskCfg.Name), zap.Error(err))
 			}
 		} else {
-			msgRow.Shard = int(msgRow.Msg.Offset>>service.offShift) % service.sharder.shards
+			msgRow.Shard = int(msgRow.Msg.Offset * (int64(msgRow.Msg.Partition + 1)) >> service.offShift % int64(service.sharder.shards))
 		}
 		service.sharder.PutElement(&msgRow)
 	}