fix: fixes the backward compatibility when dc-mapping is not set

This fixes the behavior when dc-mapping is not set.
scylladb · Feb 4, 2025 · f0df68a · f0df68a
1 parent d237454
commit f0df68a
Show file tree

Hide file tree

Showing 2 changed files with 124 additions and 12 deletions.
diff --git a/pkg/service/restore/worker.go b/pkg/service/restore/worker.go
@@ -94,10 +94,14 @@ func (w *worker) getLocationInfo(ctx context.Context, target Target) ([]Location
 			return nil, errors.Errorf("no snapshot with tag %s", target.SnapshotTag)
 		}
 
+		manifests = filterManifests(manifests, sourceDC2TargetDCMap)
+		locationDCs := collectDCsFromManifests(manifests)
+		dcHosts := hostsByDC(nodes, targetDC2SourceDCMap, locationDCs)
+
 		result = append(result, LocationInfo{
-			DC:       collectAllDCs(manifests),
-			DCHosts:  hostsByDC(nodes, targetDC2SourceDCMap),
-			Manifest: filterManifests(manifests, sourceDC2TargetDCMap),
+			DC:       locationDCs,
+			DCHosts:  dcHosts,
+			Manifest: manifests,
 			Location: l,
 		})
 	}
@@ -124,27 +128,42 @@ func (w *worker) getNodesWithAccess(ctx context.Context, nodeStatus scyllaclient
 }
 
 // hostsByDC creates map of which hosts are responsible for which DC, also applies DCMappings if available.
-func hostsByDC(nodes scyllaclient.NodeStatusInfoSlice, targetDC2sourceDCMap map[string]string) map[string][]string {
-	dcHosts := map[string][]string{}
+func hostsByDC(nodes scyllaclient.NodeStatusInfoSlice, targetDC2SourceDCMap map[string]string, locationDCs []string) map[string][]string {
+	dc2HostsMap := map[string][]string{}
+	// when --dc-mapping is not set all nodes with access to the location can handle all DCs from it
+	if len(targetDC2SourceDCMap) == 0 {
+		var hosts []string
+		for _, node := range nodes {
+			hosts = append(hosts, node.Addr)
+		}
+		for _, dc := range locationDCs {
+			dc2HostsMap[dc] = hosts
+		}
+		return dc2HostsMap
+	}
+	// when --dc-mapping is set, nodes can handle DCs only accordingly to mappings
 	for _, n := range nodes {
-		sourceDC, ok := targetDC2sourceDCMap[n.Datacenter]
+		sourceDC, ok := targetDC2SourceDCMap[n.Datacenter]
 		if !ok {
-			sourceDC = n.Datacenter
+			continue
+		}
+		if !slices.Contains(locationDCs, sourceDC) {
+			continue
 		}
-		dcHosts[sourceDC] = append(dcHosts[sourceDC], n.Addr)
+		dc2HostsMap[sourceDC] = append(dc2HostsMap[sourceDC], n.Addr)
 	}
-	return dcHosts
+	return dc2HostsMap
 }
 
-func collectAllDCs(manifests []*ManifestInfo) []string {
+func collectDCsFromManifests(manifests []*ManifestInfo) []string {
 	dcs := strset.New()
 	for _, m := range manifests {
 		dcs.Add(m.DC)
 	}
 	return dcs.List()
 }
 
-// keep only manifests that have dc mapping.
+// keep only manifests that have dc mapping. if --dc-mapping is not set it will return all manifests.
 func filterManifests(manifests []*ManifestInfo, sourceDC2TargetDCMap map[string]string) []*ManifestInfo {
 	if len(sourceDC2TargetDCMap) == 0 {
 		return manifests

diff --git a/pkg/service/restore/worker_test.go b/pkg/service/restore/worker_test.go
@@ -1,7 +1,12 @@
 // Copyright (C) 2025 ScyllaDB
 package restore
 
-import "testing"
+import (
+	"testing"
+
+	gocmp "github.com/google/go-cmp/cmp"
+	"github.com/scylladb/scylla-manager/v3/pkg/scyllaclient"
+)
 
 func TestValidateDCMappings(t *testing.T) {
 	testCases := []struct {
@@ -86,3 +91,91 @@ func TestValidateDCMappings(t *testing.T) {
 		})
 	}
 }
+
+func TestHostsByDC(t *testing.T) {
+	testCases := []struct {
+		name                 string
+		nodes                scyllaclient.NodeStatusInfoSlice
+		targetDC2SourceDCMap map[string]string
+		locationDCs          []string
+
+		expected map[string][]string
+	}{
+		{
+			name: "When --dc-mapping is provided will all DCs",
+			nodes: scyllaclient.NodeStatusInfoSlice{
+				{Addr: "n1", Datacenter: "target_dc1"},
+				{Addr: "n2", Datacenter: "target_dc1"},
+				{Addr: "n3", Datacenter: "target_dc2"},
+				{Addr: "n4", Datacenter: "target_dc2"},
+			},
+			targetDC2SourceDCMap: map[string]string{
+				"target_dc1": "source_dc1",
+				"target_dc2": "source_dc2",
+			},
+			locationDCs: []string{"source_dc1", "source_dc2"},
+
+			expected: map[string][]string{
+				"source_dc1": {"n1", "n2"},
+				"source_dc2": {"n3", "n4"},
+			},
+		},
+		{
+			name: "When --dc-mapping is provided will some DCs",
+			nodes: scyllaclient.NodeStatusInfoSlice{
+				{Addr: "n1", Datacenter: "target_dc1"},
+				{Addr: "n2", Datacenter: "target_dc1"},
+				{Addr: "n3", Datacenter: "target_dc2"},
+				{Addr: "n4", Datacenter: "target_dc2"},
+			},
+			targetDC2SourceDCMap: map[string]string{
+				"target_dc1": "source_dc1",
+			},
+			locationDCs: []string{"source_dc1", "source_dc2"},
+
+			expected: map[string][]string{
+				"source_dc1": {"n1", "n2"},
+			},
+		},
+		{
+			name: "When --dc-mapping is empty and location contains one DC",
+			nodes: scyllaclient.NodeStatusInfoSlice{
+				{Addr: "n1", Datacenter: "target_dc1"},
+				{Addr: "n2", Datacenter: "target_dc1"},
+				{Addr: "n3", Datacenter: "target_dc2"},
+				{Addr: "n4", Datacenter: "target_dc2"},
+			},
+			targetDC2SourceDCMap: map[string]string{},
+			locationDCs:          []string{"source_dc1"},
+
+			expected: map[string][]string{
+				"source_dc1": {"n1", "n2", "n3", "n4"},
+			},
+		},
+		{
+			name: "When --dc-mapping is empty and location contains multiple DCs",
+			nodes: scyllaclient.NodeStatusInfoSlice{
+				{Addr: "n1", Datacenter: "target_dc1"},
+				{Addr: "n2", Datacenter: "target_dc1"},
+				{Addr: "n3", Datacenter: "target_dc2"},
+				{Addr: "n4", Datacenter: "target_dc2"},
+			},
+			targetDC2SourceDCMap: map[string]string{},
+			locationDCs:          []string{"source_dc1", "source_dc2"},
+
+			expected: map[string][]string{
+				"source_dc1": {"n1", "n2", "n3", "n4"},
+				"source_dc2": {"n1", "n2", "n3", "n4"},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actual := hostsByDC(tc.nodes, tc.targetDC2SourceDCMap, tc.locationDCs)
+			if diff := gocmp.Diff(actual, tc.expected); diff != "" {
+				t.Fatalf("Unexpected result: %v", diff)
+			}
+		})
+	}
+}