feat: tools-2646 add --prefer-racks flag (#86)

* feat: add --prefer-racks flag

* test: add --prefer-racks option to integration tests

include/backup_config.h

@@ -72,6 +72,9 @@ extern "C" {
 // The max number of parallel scan calls made at any one time
 #define MAX_PARALLEL 100
 
+// The largest allowable rack id
+#define MAX_RACKID 1000000
+
 /*
  * The global backup configuration and stats shared by all backup threads and the counter thread.
  */
@@ -181,6 +184,9 @@ typedef struct backup_config {
 	// custom b64-encoded filter expression to use in the scan calls
 	char *filter_exp;
 
+	// racks that will be preferred during record backup
+	char *prefer_racks;
+
 	// secret agent client configs
 	sa_cfg secret_cfg;
 } backup_config_t;

include/utils.h

@@ -264,6 +264,7 @@ typedef enum {
 	COMMAND_OPT_DIRECTORY_LIST,
 	COMMAND_OPT_PARENT_DIRECTORY,
 	COMMAND_OPT_VALIDATE,
+	COMMAND_OPT_PREFER_RACKS,
 	COMMAND_SA_ADDRESS,
 	COMMAND_SA_PORT,
 	COMMAND_SA_TIMEOUT,

src/backup_config.c

@@ -145,6 +145,7 @@ backup_config_set(int argc, char* argv[], backup_config_t* conf)
 		{ "sleep-between-retries", required_argument, NULL, COMMAND_OPT_RETRY_DELAY },
 		// support the `--retry-delay` option until a major version bump.
 		{ "retry-delay", required_argument, NULL, COMMAND_OPT_RETRY_DELAY },
+		{ "prefer-racks", required_argument, NULL, COMMAND_OPT_PREFER_RACKS },
 
 		{ "s3-region", required_argument, NULL, COMMAND_OPT_S3_REGION },
 		{ "s3-profile", required_argument, NULL, COMMAND_OPT_S3_PROFILE },
@@ -732,6 +733,10 @@ backup_config_set(int argc, char* argv[], backup_config_t* conf)
 			conf->retry_delay = (uint32_t) tmp;
 			break;
 
+		case COMMAND_OPT_PREFER_RACKS:
+			conf->prefer_racks = strdup(optarg);
+			break;
+
 		case COMMAND_OPT_S3_REGION:
 			conf->s3_region = strdup(optarg);
 			break;
@@ -1003,6 +1008,8 @@ backup_config_init(backup_config_t* conf)
 	conf->max_retries = 5;
 	conf->retry_delay = 0;
 
+	conf->prefer_racks = NULL;
+
 	sa_cfg_init(&conf->secret_cfg);
 }
 
@@ -1093,6 +1100,10 @@ backup_config_destroy(backup_config_t* conf)
 		cf_free(conf->tls_name);
 	}
 
+	if (conf->prefer_racks != NULL) {
+		cf_free(conf->prefer_racks);
+	}
+
 	tls_config_destroy(&conf->tls);
 
 	sa_config_destroy(&conf->secret_cfg);
@@ -1169,6 +1180,7 @@ backup_config_clone(backup_config_t* conf)
 	clone->partition_list = safe_strdup(conf->partition_list);
 	clone->after_digest = safe_strdup(conf->after_digest);
 	clone->filter_exp = safe_strdup(conf->filter_exp);
+	clone->prefer_racks = safe_strdup(conf->prefer_racks);
 
 	sa_config_clone(&clone->secret_cfg, &conf->secret_cfg);
 
@@ -1476,6 +1488,8 @@ usage(const char *name)
 	fprintf(stdout, "                      The default is 5.\n");
 	fprintf(stdout, "      --sleep-between-retries <ms>\n");
 	fprintf(stdout, "                      The amount of time to sleep between retries. Default is 0.\n");
+	fprintf(stdout, "      --prefer-racks <rack id 1>[,<rack id 2>[,...]]\n");
+	fprintf(stdout, "                      A list of Aerospike Server rack IDs to prefer when reading records for a backup.\n");
 	fprintf(stdout, "      --s3-region <region>\n");
 	fprintf(stdout, "                      The S3 region that the bucket(s) exist in.\n");
 	fprintf(stdout, "      --s3-profile <profile_name>\n");

src/backup_status.c

@@ -217,7 +217,7 @@ backup_status_init(backup_status_t* status, backup_config_t* conf)
 
 		if (!cf_b64_validate_and_decode(conf->filter_exp, b64_len,
 					expr->packed, &expr->packed_sz)) {
-			err("Invalide base64 encoded string: %s", conf->filter_exp);
+			err("Invalid base64 encoded string: %s", conf->filter_exp);
 			cf_free(expr);
 			goto cleanup2;
 		}
@@ -348,6 +348,32 @@ backup_status_init(backup_status_t* status, backup_config_t* conf)
 		}
 	}
 
+	if (conf->prefer_racks) {
+		as_vector rackids;
+		as_vector_init(&rackids, sizeof(char*), 1);
+
+		char* racks_clone = strdup(conf->prefer_racks);
+		split_string(racks_clone, ',', true, &rackids);
+
+		for (uint32_t i = 0; i < rackids.size; i++) {
+			char* id_str = (char*) as_vector_get_ptr(&rackids, i);
+			int64_t id = 0;
+			if (!better_atoi(id_str, &id) || id < 0 || id > MAX_RACKID) {
+				err("Invalid rack id %s", id_str);
+
+				cf_free(racks_clone);
+				as_vector_destroy(&rackids);
+
+				goto cleanup2;
+			}
+
+			as_config_add_rack_id(&as_conf, (int)id);
+		}
+
+		cf_free(racks_clone);
+		as_vector_destroy(&rackids);
+	}
+
 	status->as = cf_malloc(sizeof(aerospike));
 	if (status->as == NULL) {
 		err("Failed to malloc aerospike struct");

src/conf.c

@@ -1076,6 +1076,9 @@ config_backup(toml_table_t *config_table, backup_config_t *c, const char *instan
 				status = false;
 			}
 
+		} else if (! strcasecmp("prefer-racks", name)) {
+			status = config_str(config_value, (void*)&c->prefer_racks, override);
+
 		} else if (! strcasecmp("s3-region", name)) {
 			status = config_str(config_value, (void*)&c->s3_region, override);
 

test/integration/test_secrets.py

@@ -324,6 +324,7 @@ class BackupConfigT(ComparableCtStructure):
         ("partition_list", ctypes.c_char_p),
         ("after_digest", ctypes.c_char_p),
         ("filter_exp", ctypes.c_char_p),
+		("prefer_racks", ctypes.c_char_p),
 		("secret_cfg", SCCFG)
     ]
 
@@ -360,6 +361,7 @@ class BackupConfigT(ComparableCtStructure):
 	{"name": "partition-list", "value": string_val, "config_section": "asbackup"},
 	{"name": "after-digest", "value": string_val, "config_section": "asbackup"},
 	{"name": "filter-exp", "value": string_val, "config_section": "asbackup"},
+	{"name": "prefer-racks", "value": string_val, "config_section": "asbackup"},
 	{"name": "modified-after", "value": modified_by_val, "config_section": "asbackup"},
 	{"name": "modified-before", "value": modified_by_val, "config_section": "asbackup"},
 	{"name": "records-per-second", "value": int_val, "config_section": "asbackup"},

test/integration/test_value.py

@@ -82,7 +82,8 @@ def test_string_value():
 	backup_and_restore(
 		lambda context: put_values(lib.SET, "key", STRING_VALUES),
 		None,
-		lambda context: check_values(lib.SET, "key", STRING_VALUES)
+		lambda context: check_values(lib.SET, "key", STRING_VALUES),
+		backup_opts= ["--prefer-racks", "0"]
 	)
 
 def test_blob_value():
@@ -92,7 +93,8 @@ def test_blob_value():
 	backup_and_restore(
 		lambda context: put_values(lib.SET, "key", BLOB_VALUES),
 		None,
-		lambda context: check_values(lib.SET, "key", BLOB_VALUES)
+		lambda context: check_values(lib.SET, "key", BLOB_VALUES),
+		backup_opts= ["--prefer-racks", "0,2,3"]
 	)
 
 def test_blob_value_compact():

test/unit/test_backup_conf.c

@@ -140,6 +140,7 @@ assert_bup_config_eq(backup_config_t *c1, backup_config_t *c2)
 	CMP_STR_FIELD(c1->after_digest, c2->after_digest);
 	CMP_STR_FIELD(c1->partition_list, c2->partition_list);
 	CMP_STR_FIELD(c1->after_digest, c2->after_digest);
+	CMP_STR_FIELD(c1->prefer_racks, c2->prefer_racks);
 
 	CMP_STR_FIELD(c1->s3_region, c2->s3_region);
 	CMP_STR_FIELD(c1->s3_profile, c2->s3_profile);
@@ -374,6 +375,26 @@ START_TEST(test_init_bin_list)
 }
 END_TEST
 
+START_TEST(test_init_prefer_racks)
+{
+	tmp_file_init("", "prefer-racks=\"4,2,0\"", "", "");
+	backup_config_t c1;
+	backup_config_t c2;
+	backup_config_init(&c1);
+	backup_config_init(&c2);
+	backup_config_set_heap_defaults(&c2);
+
+	ck_assert_int_ne(config_from_file(&c1, NULL, file_name, 0, true), 0);
+
+	c2.prefer_racks = strdup("4,2,0");
+
+	assert_bup_config_eq(&c1, &c2);
+
+	backup_config_destroy(&c2);
+	backup_config_destroy(&c1);
+}
+END_TEST
+
 START_TEST(test_init_mod_after)
 {
 	tmp_file_init("", "modified-after=\"2000-01-01_00:00:00\"", "", "");
@@ -896,6 +917,7 @@ Suite* backup_conf_suite()
 	tcase_add_test(tc_init, test_init_mod_after);
 	tcase_add_test(tc_init, test_init_mod_before);
 	tcase_add_test(tc_init, test_init_ttl_zero);
+	tcase_add_test(tc_init, test_init_prefer_racks);
 
 	tcase_add_test(tc_init, test_init_tls_enable);
 	tcase_add_test(tc_init, test_init_tls_protocols);