feat: Improve shard chunking (#2163)

Author: murarustefaan

scheduler/scheduler.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"sync"
 	"time"
 
@@ -305,18 +306,16 @@ func shardTableClients(tableClients []tableClient, shard *shard) []tableClient {
 	if shard == nil || len(tableClients) == 0 {
 		return tableClients
 	}
+
 	num := int(shard.num)
 	total := int(shard.total)
-	chunkSize := len(tableClients) / total
-	if chunkSize == 0 {
-		chunkSize = 1
-	}
+
+	chunkSize := int(math.Ceil(float64(len(tableClients)) / float64(total)))
+
 	chunks := lo.Chunk(tableClients, chunkSize)
 	if num > len(chunks) {
 		return nil
 	}
-	if len(chunks) > total && num == total {
-		return append(chunks[num-1], chunks[num]...)
-	}
+
 	return chunks[num-1]
 }

scheduler/scheduler_test.go

@@ -559,6 +559,7 @@ func Test_shardTableClients(t *testing.T) {
 			expected: []tableClient{
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
 			},
 		},
 		{
@@ -571,9 +572,52 @@ func Test_shardTableClients(t *testing.T) {
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
 			},
 			shard: &shard{num: 2, total: 2},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+		},
+		{
+			name: "uneven split 1 of 3",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+			shard: &shard{num: 1, total: 3},
+			expected: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+			},
+		},
+		{
+			name: "uneven split 2 of 3",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+			shard: &shard{num: 2, total: 3},
 			expected: []tableClient{
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+			},
+		},
+		{
+			name: "uneven split 3 of 3",
+			tableClients: []tableClient{
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_1"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_2"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_3"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_4"}},
+				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
+			},
+			shard: &shard{num: 3, total: 3},
+			expected: []tableClient{
 				{client: &testExecutionClient{}, table: &schema.Table{Name: "table_5"}},
 			},
 		},