feat(restore): implement native restore batching

Previously, batches restored with both rclone and native method
used the same batching mechanism. Although batching to smaller
chunks makes sense for rclone restore, as it reduces the risk
of running out of disk space and increases resource util by
having nodes which work at potentially different restore stages
(rclone download / load&stream), this logic doesn't apply to
native restore where sstables are not saved on disk before
they are streamed to the primary replicas. In native restore
case, we want to send as big batches as possible while still
keeping balanced workload distribution.

To achieve that, we just simply reuse the --batch-size=0
behavior making batches containing 5% of expected node
workload, so that batches restored with native method aim
to contain 100% of expected node workload.

To do that, we need to make the decision about restore method
used for given batch restoration during batch dispatch.
This resulted in small code refactor - mainly populating
hostInfo before accessing batching mechanism and adding
restore method as a part of the batch definition.

As working with batches containing 100% of expected node
workload is not granular at all, we might run into problems
with large workload imbalance, as giving one node one
additional batch could push it towards 200% of expected
workload. To fight with that, we create batches containing
up to the remaining expected node bytes, or up to the 5%
as in the --batch-size=0 case, so that when the node already
pushes over the 100% expected workload, it will be receiving
smaller batches that should result in more even distribution.

Fixes #4592

Author: Michal-Leszczynski

docs/source/sctool/partials/sctool_restore.yaml

@@ -15,7 +15,7 @@ options:
     - name: batch-size
       default_value: "2"
       usage: |
-        Number of SSTables per shard to process in one request by one node.
+        Number of SSTables per shard to process in one request by one node when restoring with --method=rclone.
         Increasing the default batch size might significantly increase restore performance, as only one shard can work on restoring a single SSTable bundle.
         Set to 0 for best performance (batches will contain sstables of total size up to 5% of expected total node workload).
     - name: cluster

docs/source/sctool/partials/sctool_restore_update.yaml

@@ -13,7 +13,7 @@ options:
     - name: batch-size
       default_value: "2"
       usage: |
-        Number of SSTables per shard to process in one request by one node.
+        Number of SSTables per shard to process in one request by one node when restoring with --method=rclone.
         Increasing the default batch size might significantly increase restore performance, as only one shard can work on restoring a single SSTable bundle.
         Set to 0 for best performance (batches will contain sstables of total size up to 5% of expected total node workload).
     - name: cluster

pkg/command/restore/res.yaml

@@ -26,7 +26,7 @@ snapshot-tag: |
   Snapshot tags can be obtained from backup listing ('./sctool backup list' command - e.g. sm_20060102150405UTC).
 
 batch-size: |
-  Number of SSTables per shard to process in one request by one node.
+  Number of SSTables per shard to process in one request by one node when restoring with --method=rclone.
   Increasing the default batch size might significantly increase restore performance, as only one shard can work on restoring a single SSTable bundle.
   Set to 0 for best performance (batches will contain sstables of total size up to 5% of expected total node workload).
 

pkg/service/restore/batch.go

@@ -12,6 +12,12 @@ import (
 	"github.com/scylladb/scylla-manager/v3/pkg/sstable"
 )
 
+type hostBatchInfo struct {
+	host                 string
+	shardCnt             uint
+	nativeRestoreSupport bool
+}
+
 // batchDispatcher is a tool for batching SSTables from
 // Workload across different hosts during restore.
 // It follows a few rules:
@@ -53,18 +59,21 @@ type batchDispatcher struct {
 	// For batchSize 0, batches contain N*node_shard_cnt SSTables
 	// of total size up to 5% of node expected workload
 	// (expectedShardWorkload*node_shard_cnt).
+	// For hosts and batches supporting native restore, regardless of batchSize,
+	// batches contain N*node_shard_cnt SSTables of total
+	// size up to 100% of node expected workload.
 	batchSize int
 	// Equals total_backup_size/($\sum_{node} shard_cnt(node)$)
 	expectedShardWorkload int64
-	// Stores host shard count
-	hostShardCnt map[string]uint
+	hostInfo              map[string]hostBatchInfo
+	method                Method
 }
 
-func newBatchDispatcher(workload Workload, batchSize int, hostShardCnt map[string]uint, locationInfo []LocationInfo) *batchDispatcher {
+func newBatchDispatcher(workload Workload, batchSize int, hostInfo map[string]hostBatchInfo, locationInfo []LocationInfo, method Method) *batchDispatcher {
 	sortWorkload(workload)
 	var shards uint
-	for _, sh := range hostShardCnt {
-		shards += sh
+	for _, hi := range hostInfo {
+		shards += hi.shardCnt
 	}
 	if shards == 0 {
 		shards = 1
@@ -76,7 +85,8 @@ func newBatchDispatcher(workload Workload, batchSize int, hostShardCnt map[strin
 		workloadProgress:      newWorkloadProgress(workload, locationInfo),
 		batchSize:             batchSize,
 		expectedShardWorkload: workload.TotalSize / int64(shards),
-		hostShardCnt:          hostShardCnt,
+		hostInfo:              hostInfo,
+		method:                method,
 	}
 }
 
@@ -94,6 +104,10 @@ type workloadProgress struct {
 	// It assumes that the whole DC is backed up to a single
 	// backup location.
 	hostDCAccess map[string][]string
+	// Stores the amount of bytes assigned for given host to restore.
+	// It's used when as an additional safety net when batching sstables
+	// eligible for native restore (expected batch size is 100 % of expected node workload).
+	hostAssignedBytes map[string]int64
 	// SSTables grouped by RemoteSSTableDir that are yet to
 	// be batched. They are removed on batch dispatch, but can
 	// be re-added when batch failed to be restored.
@@ -129,6 +143,7 @@ func newWorkloadProgress(workload Workload, locationInfo []LocationInfo) workloa
 		dcSSTableToBeRestored: dcSSTable,
 		hostFailedDC:          make(map[string][]string),
 		hostDCAccess:          getHostDCAccess(locationInfo),
+		hostAssignedBytes:     make(map[string]int64),
 		remoteDir:             p,
 	}
 }
@@ -164,6 +179,7 @@ type batch struct {
 	*ManifestInfo
 
 	batchType        batchType
+	method           Method // that should be used for this batch restoration
 	RemoteSSTableDir string
 	Size             int64
 	SSTables         []RemoteSSTable
@@ -286,6 +302,7 @@ func (bd *batchDispatcher) DispatchBatch(ctx context.Context, host string) (batc
 
 func (bd *batchDispatcher) dispatchBatch(host string) (batch, bool) {
 	dirIdx := -1
+
 	for i := range bd.workloadProgress.remoteDir {
 		rdw := bd.workload.RemoteDir[i]
 		// Skip empty dir
@@ -312,9 +329,9 @@ func (bd *batchDispatcher) dispatchBatch(host string) (batch, bool) {
 // Returns batch from given RemoteSSTableDir and updates workloadProgress.
 func (bd *batchDispatcher) createBatch(dirIdx int, host string) (batch, bool) {
 	rdp := &bd.workloadProgress.remoteDir[dirIdx]
-	shardCnt := bd.hostShardCnt[host]
-	if shardCnt == 0 {
-		shardCnt = 1
+	hi := bd.hostInfo[host]
+	if hi.shardCnt == 0 {
+		hi.shardCnt = 1
 	}
 
 	// Choose batch type and candidate sstables
@@ -331,15 +348,36 @@ func (bd *batchDispatcher) createBatch(dirIdx int, host string) (batch, bool) {
 		return batch{}, false
 	}
 
+	// At this point, if method is explicitly set to MethodNative,
+	// we already verified that all hosts and the entire workload
+	// supports native restore, so we can safely fall back to rclone if needed.
+	batchMethod := MethodRclone
+	if slices.Contains([]Method{MethodAuto, MethodNative}, bd.method) &&
+		hi.nativeRestoreSupport &&
+		sstables[0].NativeRestoreSupport() == nil {
+		batchMethod = MethodNative
+	}
+
 	var i int
 	var size int64
-	if bd.batchSize == maxBatchSize {
+	if bd.batchSize == maxBatchSize || batchMethod == MethodNative {
 		// Create batch containing multiple of node shard count sstables
-		// and size up to 5% of expected node workload.
-		expectedNodeWorkload := bd.expectedShardWorkload * int64(shardCnt)
+		// and size up to a percentage of expected node workload.
+		// For rclone restore with --batch-size=0, aim for 5% of expected node workload.
+		// For native restore, aim for 100% of expected node workload.
+		expectedNodeWorkload := bd.expectedShardWorkload * int64(hi.shardCnt)
 		sizeLimit := expectedNodeWorkload / 20
+		if batchMethod == MethodNative {
+			// For 5% rclone batches, we expect imbalance in workload distribution
+			// to be less significant even in the more problematic scenarios.
+			// The risk of noticeable imbalance in 100% native batches is much higher,
+			// so we fall back to 5% batches when node already restored its expected workload.
+			alreadyAssigned := bd.workloadProgress.hostAssignedBytes[host]
+			remainingExpected := max(expectedNodeWorkload-alreadyAssigned, 0)
+			sizeLimit = max(remainingExpected, sizeLimit)
+		}
 		for {
-			for range shardCnt {
+			for range hi.shardCnt {
 				if i >= len(sstables) {
 					break
 				}
@@ -355,7 +393,7 @@ func (bd *batchDispatcher) createBatch(dirIdx int, host string) (batch, bool) {
 		}
 	} else {
 		// Create batch containing node_shard_count*batch_size sstables.
-		i = min(bd.batchSize*int(shardCnt), len(sstables))
+		i = min(bd.batchSize*int(hi.shardCnt), len(sstables))
 		for j := range i {
 			size += sstables[j].Size
 		}
@@ -366,19 +404,21 @@ func (bd *batchDispatcher) createBatch(dirIdx int, host string) (batch, bool) {
 	}
 	// Extend batch if it was to leave less than
 	// 1 sstable per shard for the next one.
-	if len(sstables)-i < int(shardCnt) {
+	if len(sstables)-i < int(hi.shardCnt) {
 		for ; i < len(sstables); i++ {
 			size += sstables[i].Size
 		}
 	}
 
 	rdp.RemainingSSTables[batchT] = sstables[i:]
 	rdw := bd.workload.RemoteDir[dirIdx]
+	bd.workloadProgress.hostAssignedBytes[host] += size
 
 	return batch{
 		TableName:        rdw.TableName,
 		ManifestInfo:     rdw.ManifestInfo,
 		batchType:        batchT,
+		method:           batchMethod,
 		RemoteSSTableDir: rdw.RemoteSSTableDir,
 		Size:             size,
 		SSTables:         sstables[:i],
@@ -405,6 +445,7 @@ func (bd *batchDispatcher) ReportFailure(host string, b batch) error {
 
 	// Mark failed DC for host
 	bd.workloadProgress.hostFailedDC[host] = append(bd.workloadProgress.hostFailedDC[host], b.DC)
+	bd.workloadProgress.hostAssignedBytes[host] -= b.Size
 
 	dirIdx := -1
 	for i := range bd.workload.RemoteDir {