FMWK-564-prefer-racks

- added prefer racks
- fixed error with nodes and partitions backup

cmd/asbackup/readme.md

@@ -109,6 +109,7 @@ Backup Flags:
                                     This argument is mutually exclusive to partition-list/after-digest arguments.
                                     Default: backup all nodes in the cluster
       --no-ttl-only                 Only include records that have no ttl set (persistent records).
+      --prefer-racks string         <rack id 1>[,<rack id 2>[,...]]\nA list of Aerospike Server rack IDs to prefer when reading records for a backup.
 
 Compression Flags:
   -z, --compress string         Enables compressing of backup files using the specified compression algorithm.
@@ -185,9 +186,6 @@ Azure Flags:
 
 --estimate-samples  The number of samples to take when running a backup estimate.
 
---prefer-racks      <rack id 1>[,<rack id 2>[,...]]
-                    A list of Aerospike Server rack IDs to prefer when reading records for a backup.
-
 --no-config-file    Do not read any config file. Default: disabled
 
 --instance          Section with these instance is read. e.g in case instance `a` is specified

cmd/internal/app/asbackup.go

@@ -61,7 +61,7 @@ func NewASBackup(
 		return nil, nil
 	}
 
-	aerospikeClient, err := newAerospikeClient(clientConfig)
+	aerospikeClient, err := newAerospikeClient(clientConfig, backupParams.PreferRacks)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create aerospike client: %w", err)
 	}

cmd/internal/app/asrestore.go

@@ -56,7 +56,7 @@ func NewASRestore(
 		return nil, fmt.Errorf("failed to create backup reader: %w", err)
 	}
 
-	aerospikeClient, err := newAerospikeClient(clientConfig)
+	aerospikeClient, err := newAerospikeClient(clientConfig, "")
 	if err != nil {
 		return nil, fmt.Errorf("failed to create aerospike client: %w", err)
 	}

cmd/internal/app/clients.go

@@ -17,6 +17,7 @@ package app
 import (
 	"context"
 	"fmt"
+	"strconv"
 
 	"cloud.google.com/go/storage"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -29,7 +30,9 @@ import (
 	"google.golang.org/api/option"
 )
 
-func newAerospikeClient(cfg *client.AerospikeConfig) (*aerospike.Client, error) {
+const maxRack = 1000000
+
+func newAerospikeClient(cfg *client.AerospikeConfig, racks string) (*aerospike.Client, error) {
 	if len(cfg.Seeds) < 1 {
 		return nil, fmt.Errorf("at least one seed must be provided")
 	}
@@ -39,6 +42,16 @@ func newAerospikeClient(cfg *client.AerospikeConfig) (*aerospike.Client, error)
 		return nil, fmt.Errorf("failed to create Aerospike client policy: %w", err)
 	}
 
+	if racks != "" {
+		racksIDs, err := parseRacks(racks)
+		if err != nil {
+			return nil, err
+		}
+
+		p.RackIds = racksIDs
+		p.RackAware = true
+	}
+
 	asClient, err := aerospike.NewClientWithPolicyAndHost(p, toHosts(cfg.Seeds)...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create Aerospike client: %w", err)
@@ -135,3 +148,27 @@ func toHosts(htpSlice client.HostTLSPortSlice) []*aerospike.Host {
 
 	return hosts
 }
+
+func parseRacks(racks string) ([]int, error) {
+	racksStringSlice := stringSplit(racks)
+	racksIntSlice := make([]int, len(racksStringSlice))
+
+	for i := range racksStringSlice {
+		rackID, err := strconv.Atoi(racksStringSlice[i])
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse racks: %w", err)
+		}
+
+		if rackID < 0 {
+			return nil, fmt.Errorf("rack id %d invalid, should be positive number", rackID)
+		}
+
+		if rackID > maxRack {
+			return nil, fmt.Errorf("rack id %d invalid, should not exceed %d", rackID, maxRack)
+		}
+
+		racksIntSlice = append(racksIntSlice, rackID)
+	}
+
+	return racksIntSlice, nil
+}

cmd/internal/app/clients_test.go

@@ -37,14 +37,14 @@ func TestClients_newAerospikeClient(t *testing.T) {
 		User:     testASLoginPassword,
 		Password: testASLoginPassword,
 	}
-	_, err := newAerospikeClient(cfg)
+	_, err := newAerospikeClient(cfg, "1")
 	require.NoError(t, err)
 
 	cfg = &client.AerospikeConfig{
 		User:     testASLoginPassword,
 		Password: testASLoginPassword,
 	}
-	_, err = newAerospikeClient(cfg)
+	_, err = newAerospikeClient(cfg, "")
 	require.ErrorContains(t, err, "at least one seed must be provided")
 
 	cfg = &client.AerospikeConfig{
@@ -57,7 +57,7 @@ func TestClients_newAerospikeClient(t *testing.T) {
 			Cert: []byte("error"),
 		},
 	}
-	_, err = newAerospikeClient(cfg)
+	_, err = newAerospikeClient(cfg, "")
 	require.ErrorContains(t, err, "failed to create Aerospike client policy")
 
 	hostPort.Host = "255.255.255.255"
@@ -68,7 +68,7 @@ func TestClients_newAerospikeClient(t *testing.T) {
 		User:     testASLoginPassword,
 		Password: testASLoginPassword,
 	}
-	_, err = newAerospikeClient(cfg)
+	_, err = newAerospikeClient(cfg, "")
 	require.ErrorContains(t, err, "failed to create Aerospike client")
 }
 

cmd/internal/app/configs.go

@@ -25,8 +25,6 @@ import (
 	bModels "github.com/aerospike/backup-go/models"
 )
 
-const sliceSeparator = ","
-
 func mapBackupConfig(
 	backupParams *models.Backup,
 	commonParams *models.Common,
@@ -40,8 +38,8 @@ func mapBackupConfig(
 
 	c := backup.NewDefaultBackupConfig()
 	c.Namespace = commonParams.Namespace
-	c.SetList = strings.Split(commonParams.SetList, sliceSeparator)
-	c.BinList = strings.Split(commonParams.BinList, sliceSeparator)
+	c.SetList = stringSplit(commonParams.SetList)
+	c.BinList = stringSplit(commonParams.BinList)
 	c.NoRecords = commonParams.NoRecords
 	c.NoIndexes = commonParams.NoIndexes
 	c.RecordsPerSecond = commonParams.RecordsPerSecond
@@ -54,11 +52,16 @@ func mapBackupConfig(
 	// As we set --nice in MiB we must convert it to bytes
 	// TODO: make Bandwidth int64 to avoid overflow.
 	c.Bandwidth = commonParams.Nice * 1024 * 1024
-	c.ParallelNodes = backupParams.ParallelNodes
 	c.Compact = backupParams.Compact
-	c.NodeList = strings.Split(backupParams.NodeList, sliceSeparator)
 	c.NoTTLOnly = backupParams.NoTTLOnly
 
+	// Overwrite partitions if we use nodes.
+	if backupParams.ParallelNodes || backupParams.NodeList != "" {
+		c.Partitions = backup.PartitionRange{}
+		c.ParallelNodes = backupParams.ParallelNodes
+		c.NodeList = stringSplit(backupParams.NodeList)
+	}
+
 	sp, err := mapScanPolicy(backupParams, commonParams)
 	if err != nil {
 		return nil, err
@@ -103,8 +106,8 @@ func mapRestoreConfig(
 
 	c := backup.NewDefaultRestoreConfig()
 	c.Namespace = mapRestoreNamespace(commonParams.Namespace)
-	c.SetList = strings.Split(commonParams.SetList, sliceSeparator)
-	c.BinList = strings.Split(commonParams.BinList, sliceSeparator)
+	c.SetList = stringSplit(commonParams.SetList)
+	c.BinList = stringSplit(commonParams.BinList)
 	c.NoRecords = commonParams.NoRecords
 	c.NoIndexes = commonParams.NoIndexes
 	c.RecordsPerSecond = commonParams.RecordsPerSecond
@@ -124,7 +127,7 @@ func mapRestoreConfig(
 }
 
 func mapRestoreNamespace(n string) *backup.RestoreNamespaceConfig {
-	nsArr := strings.Split(n, ",")
+	nsArr := stringSplit(n)
 
 	var source, destination string
 
@@ -216,6 +219,10 @@ func mapScanPolicy(b *models.Backup, c *models.Common) (*aerospike.ScanPolicy, e
 	p.SleepBetweenRetries = time.Duration(b.SleepBetweenRetries) * time.Millisecond
 	p.TotalTimeout = time.Duration(c.TotalTimeout) * time.Millisecond
 	p.SocketTimeout = time.Duration(c.SocketTimeout) * time.Millisecond
+	// If we selected racks we must set replica policy to aerospike.PREFER_RACK
+	if b.PreferRacks != "" {
+		p.ReplicaPolicy = aerospike.PREFER_RACK
+	}
 
 	if b.NoBins {
 		p.IncludeBinData = false
@@ -275,3 +282,11 @@ func mapRetryPolicy(r *models.Restore) *bModels.RetryPolicy {
 		MaxRetries:  r.RetryMaxRetries,
 	}
 }
+
+func stringSplit(s string) []string {
+	if s == "" {
+		return nil
+	}
+
+	return strings.Split(s, ",")
+}

cmd/internal/flags/backup.go

@@ -98,6 +98,10 @@ func (f *Backup) NewFlagSet() *pflag.FlagSet {
 	flagSet.BoolVar(&f.NoTTLOnly, "no-ttl-only",
 		false,
 		"Only include records that have no ttl set (persistent records).")
+	flagSet.StringVar(&f.PreferRacks, "prefer-racks",
+		"",
+		"<rack id 1>[,<rack id 2>[,...]]\n"+
+			"A list of Aerospike Server rack IDs to prefer when reading records for a backup.")
 
 	return flagSet
 }

cmd/internal/flags/backup_test.go

@@ -41,6 +41,7 @@ func TestBackup_NewFlagSet(t *testing.T) {
 		"--compact",
 		"--node-list", "node1,node2",
 		"--no-ttl-only",
+		"--prefer-racks", "1,2,3,4",
 	}
 
 	err := flagSet.Parse(args)
@@ -62,6 +63,7 @@ func TestBackup_NewFlagSet(t *testing.T) {
 	assert.Equal(t, true, result.Compact, "The compact flag should be parsed correctly")
 	assert.Equal(t, "node1,node2", result.NodeList, "The node-list flag should be parsed correctly")
 	assert.Equal(t, true, result.NoTTLOnly, "The no-ttl-only flag should be parsed correctly")
+	assert.Equal(t, "1,2,3,4", result.PreferRacks, "The prefer-racks flag should be parsed correctly")
 }
 
 func TestBackup_NewFlagSet_DefaultValues(t *testing.T) {
@@ -89,4 +91,5 @@ func TestBackup_NewFlagSet_DefaultValues(t *testing.T) {
 	assert.Equal(t, false, result.Compact, "The default value for compact should be false")
 	assert.Equal(t, "", result.NodeList, "The default value for node-list should be empty")
 	assert.Equal(t, false, result.NoTTLOnly, "The default value for no-ttl-only should be false")
+	assert.Equal(t, "", result.PreferRacks, "The default value for prefer-racks should be false")
 }

cmd/internal/models/backup.go

@@ -30,6 +30,7 @@ type Backup struct {
 	Compact             bool
 	NodeList            string
 	NoTTLOnly           bool
+	PreferRacks         string
 }
 
 // ShouldClearTarget check if we should clean target directory.