feat: TOOLS-2514 add --s3-connect-timeout (#58)

* feat: TOOLS-2514 add --s3-connect-timeout

include/backup_config.h

@@ -87,12 +87,14 @@ typedef struct backup_config {
 	char* s3_profile;
 	// An alternative endpoint for S3 compatible storage to send all S3 requests to.
 	char* s3_endpoint_override;
-	// A user override of the minimum part size to use for S3 Multipart Uplaod parts.
+	// A user override of the minimum part size to use for S3 Multipart Upload parts.
 	uint64_t s3_min_part_size;
 	// Max simultaneous download requests from S3 allowed at a time.
 	uint32_t s3_max_async_downloads;
 	// Max simultaneous upload requests from S3 allowed at a time.
 	uint32_t s3_max_async_uploads;
+	// aws-sdk-cpp client connectTimeoutMS.
+	uint32_t s3_connect_timeout;
 	// Logging level of the AWS S3 C+ SDK.
 	s3_log_level_t s3_log_level;
 

include/file_proxy.h

@@ -60,6 +60,7 @@ typedef struct backup_state backup_state_t;
 
 #define S3_DEFAULT_MAX_ASYNC_UPLOADS 16
 #define S3_DEFAULT_MAX_ASYNC_DOWNLOADS 32
+#define S3_DEFAULT_CONNECT_TIMEOUT_MS 1000
 
 #define S3_DEFAULT_LOG_LEVEL ((s3_log_level_t) Fatal)
 
@@ -181,6 +182,12 @@ void s3_set_endpoint(const char* endpoint);
 void s3_set_max_async_downloads(uint32_t max_async_downloads);
 void s3_set_max_async_uploads(uint32_t max_async_uploads);
 
+/*
+ * Sets the S3 client connectTimeoutMS client config field.
+ */
+void
+s3_set_connect_timeout_ms(uint32_t connect_timeout_ms);
+
 /*
  * Sets the logging level of the AWS c++ sdk.
  */

include/restore_config.h

@@ -120,6 +120,8 @@ typedef struct restore_config {
 	char* s3_endpoint_override;
 	// Max simultaneous download requests from S3 allowed at a time.
 	uint32_t s3_max_async_downloads;
+	// aws-sdk-cpp client connectTimeoutMS.
+	uint32_t s3_connect_timeout;
 	// Logging level of the AWS S3 C+ SDK.
 	s3_log_level_t s3_log_level;
 

include/s3_api.h

@@ -74,6 +74,8 @@ class S3API {
 
 	S3API& SetMaxAsyncUploads(uint32_t max_async_uploads);
 
+	S3API& SetConnectTimeoutMS(uint32_t connect_timeout_ms);
+
 	GroupDownloadManager* GetGroupDownloadManager();
 
 	const std::string& GetRegion() const;
@@ -133,9 +135,10 @@ class S3API {
 
 	Aws::S3::S3Client* client;
 
-	// must be initialized with SetMaxAsync... methods
+	// must be initialized with Set<var_name>... methods
 	uint32_t max_async_uploads;
 	uint32_t max_async_downloads;
+	uint32_t connect_timeout_ms;
 
 	// the current number of concurrent async uploads
 	std::atomic<uint32_t> async_uploads;

include/utils.h

@@ -250,6 +250,7 @@ typedef enum {
 	COMMAND_OPT_S3_MAX_ASYNC_DOWNLOADS,
 	COMMAND_OPT_S3_MAX_ASYNC_UPLOADS,
 	COMMAND_OPT_S3_LOG_LEVEL,
+	COMMAND_OPT_S3_CONNECT_TIMEOUT,
 	COMMAND_OPT_DISABLE_BATCH_WRITES,
 	COMMAND_OPT_MAX_ASYNC_BATCHES,
 	COMMAND_OPT_BATCH_SIZE,

src/backup_config.c

@@ -154,6 +154,7 @@ backup_config_init(int argc, char* argv[], backup_config_t* conf)
 		{ "s3-max-async-downloads", required_argument, NULL, COMMAND_OPT_S3_MAX_ASYNC_DOWNLOADS },
 		{ "s3-max-async-uploads", required_argument, NULL, COMMAND_OPT_S3_MAX_ASYNC_UPLOADS },
 		{ "s3-log-level", required_argument, NULL, COMMAND_OPT_S3_LOG_LEVEL },
+		{ "s3-connect-timeout", required_argument, NULL, COMMAND_OPT_S3_CONNECT_TIMEOUT },
 		{ NULL, 0, NULL, 0 }
 	};
 
@@ -626,6 +627,14 @@ backup_config_init(int argc, char* argv[], backup_config_t* conf)
 			conf->s3_max_async_uploads = (uint32_t) tmp;
 			break;
 
+		case COMMAND_OPT_S3_CONNECT_TIMEOUT:
+			if (!better_atoi(optarg, &tmp) || tmp < 0 || tmp > UINT_MAX) {
+				err("Invalid S3 connect timeout value %s", optarg);
+				return BACKUP_CONFIG_INIT_FAILURE;
+			}
+			conf->s3_connect_timeout = (uint32_t) tmp;
+			break;
+
 		case COMMAND_OPT_S3_LOG_LEVEL:
 			if (!s3_parse_log_level(optarg, &s3_log_level)) {
 				err("Invalid S3 log level \"%s\"", optarg);
@@ -758,6 +767,7 @@ backup_config_init(int argc, char* argv[], backup_config_t* conf)
 
 	s3_set_max_async_downloads(conf->s3_max_async_downloads);
 	s3_set_max_async_uploads(conf->s3_max_async_uploads);
+	s3_set_connect_timeout_ms(conf->s3_connect_timeout);
 	s3_set_log_level(conf->s3_log_level);
 
 	if (conf->estimate) {
@@ -794,6 +804,7 @@ backup_config_default(backup_config_t* conf)
 	conf->s3_min_part_size = 0;
 	conf->s3_max_async_downloads = S3_DEFAULT_MAX_ASYNC_DOWNLOADS;
 	conf->s3_max_async_uploads = S3_DEFAULT_MAX_ASYNC_UPLOADS;
+	conf->s3_connect_timeout = S3_DEFAULT_CONNECT_TIMEOUT_MS;
 	conf->s3_log_level = S3_DEFAULT_LOG_LEVEL;
 
 	memset(conf->ns, 0, sizeof(as_namespace));
@@ -1342,7 +1353,11 @@ usage(const char *name)
 	fprintf(stdout, "                       - Info\n");
 	fprintf(stdout, "                       - Debug\n");
 	fprintf(stdout, "                       - Trace\n");
-	fprintf(stdout, "                      The default is Fatal.\n\n");
+	fprintf(stdout, "                      The default is Fatal.\n");
+	fprintf(stdout, "      --s3-connect-timeout <ms>\n");
+	fprintf(stdout, "                      The AWS S3 client's connection timeout in milliseconds.\n");
+	fprintf(stdout, "                      This is equivalent to cli-connect-timeout in the AWS CLI,\n");
+	fprintf(stdout, "                      or connectTimeoutMS in the aws-sdk-cpp client configuration.\n\n");
 
 	fprintf(stdout, "\n\n");
 	fprintf(stdout, "Default configuration files are read from the following files in the given order:\n");

src/conf.c

@@ -847,6 +847,14 @@ config_backup(toml_table_t *conftab, backup_config_t *c, const char *instance,
 				status = false;
 			}
 
+		} else if (! strcasecmp("s3-connect-timeout", name)) {
+			status = config_int64(curtab, name, (void*)&i_val);
+			if (i_val >= 0 && i_val <= UINT_MAX) {
+				c->s3_connect_timeout = (uint32_t) i_val;
+			} else {
+				status = false;
+			}
+
 		} else if (! strcasecmp("s3-log-level", name)) {
 			s = NULL;
 			status = config_str(curtab, name, (void*)&s);
@@ -1098,6 +1106,14 @@ config_restore(toml_table_t *conftab, restore_config_t *c, const char *instance,
 		} else if (! strcasecmp("s3-endpoint-override", name)) {
 			status = config_str(curtab, name, (void*)&c->s3_endpoint_override);
 
+		} else if (! strcasecmp("s3-connect-timeout", name)) {
+			status = config_int64(curtab, name, (void*)&i_val);
+			if (i_val >= 0 && i_val <= UINT_MAX) {
+				c->s3_connect_timeout = (uint32_t) i_val;
+			} else {
+				status = false;
+			}
+
 		} else if (! strcasecmp("s3-max-async-downloads", name)) {
 			status = config_int64(curtab, name, (void*)&i_val);
 			if (i_val > 0 && i_val <= UINT_MAX) {

src/file_proxy_s3.cc

@@ -170,6 +170,12 @@ s3_set_max_async_downloads(uint32_t max_async_downloads)
 	g_api.SetMaxAsyncDownloads(max_async_downloads);
 }
 
+void
+s3_set_connect_timeout_ms(uint32_t connect_timeout_ms)
+{
+	g_api.SetConnectTimeoutMS(connect_timeout_ms);
+}
+
 void
 s3_set_max_async_uploads(uint32_t max_async_uploads)
 {

src/restore_config.c

@@ -150,6 +150,7 @@ restore_config_init(int argc, char* argv[], restore_config_t* conf)
 		{ "s3-endpoint-override", required_argument, NULL, COMMAND_OPT_S3_ENDPOINT_OVERRIDE },
 		{ "s3-max-async-downloads", required_argument, NULL, COMMAND_OPT_S3_MAX_ASYNC_DOWNLOADS },
 		{ "s3-log-level", required_argument, NULL, COMMAND_OPT_S3_LOG_LEVEL },
+		{ "s3-connect-timeout", required_argument, NULL, COMMAND_OPT_S3_CONNECT_TIMEOUT },
 		{ NULL, 0, NULL, 0 }
 	};
 
@@ -573,6 +574,14 @@ restore_config_init(int argc, char* argv[], restore_config_t* conf)
 			conf->s3_max_async_downloads = (uint32_t) tmp;
 			break;
 
+		case COMMAND_OPT_S3_CONNECT_TIMEOUT:
+			if (!better_atoi(optarg, &tmp) || tmp < 0 || tmp > UINT_MAX) {
+				err("Invalid S3 connect timeout value %s", optarg);
+				return RESTORE_CONFIG_INIT_FAILURE;
+			}
+			conf->s3_connect_timeout = (uint32_t) tmp;
+			break;
+
 		case COMMAND_OPT_S3_LOG_LEVEL:
 			if (!s3_parse_log_level(optarg, &s3_log_level)) {
 				err("Invalid S3 log level \"%s\"", optarg);
@@ -651,6 +660,7 @@ restore_config_init(int argc, char* argv[], restore_config_t* conf)
 	}
 
 	s3_set_max_async_downloads(conf->s3_max_async_downloads);
+	s3_set_connect_timeout_ms(conf->s3_connect_timeout);
 	s3_set_log_level(conf->s3_log_level);
 
 	if (conf->nice_list != NULL) {
@@ -708,6 +718,7 @@ restore_config_default(restore_config_t *conf)
 	conf->s3_profile = NULL;
 	conf->s3_endpoint_override = NULL;
 	conf->s3_max_async_downloads = S3_DEFAULT_MAX_ASYNC_DOWNLOADS;
+	conf->s3_connect_timeout = S3_DEFAULT_CONNECT_TIMEOUT_MS;
 	conf->s3_log_level = S3_DEFAULT_LOG_LEVEL;
 
 	conf->parallel = DEFAULT_THREADS;
@@ -1125,7 +1136,11 @@ usage(const char *name)
 	fprintf(stdout, "                       - Info\n");
 	fprintf(stdout, "                       - Debug\n");
 	fprintf(stdout, "                       - Trace\n");
-	fprintf(stdout, "                      The default is Fatal.\n\n");
+	fprintf(stdout, "                      The default is Fatal.\n");
+	fprintf(stdout, "      --s3-connect-timeout <ms>\n");
+	fprintf(stdout, "                      The AWS S3 client's connection timeout in milliseconds.\n");
+	fprintf(stdout, "                      This is equivalent to cli-connect-timeout in the AWS CLI,\n");
+	fprintf(stdout, "                      or connectTimeoutMS in the aws-sdk-cpp client configuration.\n\n");
 
 	fprintf(stdout, "\n\n");
 	fprintf(stdout, "Default configuration files are read from the following files in the given order:\n");

src/s3_api.cc

@@ -145,6 +145,13 @@ S3API::SetMaxAsyncUploads(uint32_t max_async_uploads)
 	return *this;
 }
 
+S3API&
+S3API::SetConnectTimeoutMS(uint32_t connect_timeout_ms)
+{
+	this->connect_timeout_ms = connect_timeout_ms;
+	return *this;
+}
+
 GroupDownloadManager*
 S3API::GetGroupDownloadManager()
 {
@@ -282,6 +289,8 @@ S3API::_init_api(S3API& s3_api)
 
 	conf.maxConnections = std::max(s3_api.max_async_downloads,
 			s3_api.max_async_uploads);
+
+	conf.connectTimeoutMs = s3_api.connect_timeout_ms;
 
 	s3_api.client = new Aws::S3::S3Client(conf,
 			Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Always,

test/integration/test_s3.py

@@ -165,7 +165,7 @@ def test_s3_backup():
 	do_s3_backup(0)
 
 def test_s3_backup_multiple_files():
-	do_s3_backup(0, n_records=10000, backup_opts=['--file-limit', '1'])
+	do_s3_backup(0, n_records=10000, backup_opts=['--file-limit', '1', '--s3-connect-timeout', '2000'], restore_opts=['--s3-connect-timeout', '2000'])
 
 def test_s3_backup_interrupt():
 	do_s3_backup(1, n_records=10000, backup_opts=['--records-per-second', '200'])

test/unit/test_backup_conf.c

@@ -138,6 +138,16 @@ 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->s3_region, c2->s3_region);
+	CMP_STR_FIELD(c1->s3_bucket, c2->s3_bucket);
+	CMP_STR_FIELD(c1->s3_profile, c2->s3_profile);
+	CMP_STR_FIELD(c1->s3_endpoint_override, c2->s3_endpoint_override);
+	CMP_INT_FIELD(c1->s3_min_part_size, c2->s3_min_part_size);
+	CMP_INT_FIELD(c1->s3_max_async_downloads, c2->s3_max_async_downloads);
+	CMP_INT_FIELD(c1->s3_max_async_uploads, c2->s3_max_async_uploads);
+	CMP_INT_FIELD(c1->s3_connect_timeout, c2->s3_connect_timeout);
+	CMP_INT_FIELD(c1->s3_log_level, c2->s3_log_level);
 }
 
 
@@ -366,6 +376,24 @@ START_TEST(test_init_mod_before)
 }
 END_TEST
 
+START_TEST(test_init_s3_log_level)
+{
+	tmp_file_init("", "s3-log-level=\"Debug\"\n", "");
+	backup_config_t c1;
+	backup_config_t c2;
+	backup_config_default(&c1);
+	backup_config_default(&c2);
+
+	ck_assert_int_ne(config_from_file(&c1, NULL, file_name, 0, true), 0);
+
+	c2.s3_log_level = Debug;
+
+	assert_bup_config_eq(&c1, &c2);
+
+	backup_config_destroy(&c2);
+	backup_config_destroy(&c1);
+}
+END_TEST
 
 START_TEST(test_init_compress_mode)
 {
@@ -717,6 +745,15 @@ DEFINE_STR_TEST(test_init_after_digest, "after-digest", after_digest,
 DEFINE_INT_TEST(test_init_max_recs, "max-records", max_records);
 DEFINE_INT_TEST(test_init_records_per_second, "records-per-second", records_per_second);
 
+DEFINE_STR_TEST(test_init_s3_region, "s3-region", s3_region, "us-west-1");
+DEFINE_STR_TEST(test_init_s3_bucket, "s3-bucket", s3_bucket, "my-hosted-content");
+DEFINE_STR_TEST(test_init_s3_profile, "s3-profile", s3_profile, "default");
+DEFINE_STR_TEST(test_init_s3_endpoint_override, "s3-endpoint-override", s3_endpoint_override,
+		"https://<accountid>.r2.test.com");
+DEFINE_INT_TEST_MULT(test_init_s3_min_part_size, "s3-min-part-size", s3_min_part_size, 1024 * 1024);
+DEFINE_INT_TEST(test_init_s3_max_async_downloads, "s3-max-async-downloads", s3_max_async_downloads);
+DEFINE_INT_TEST(test_init_s3_max_async_uploads, "s3-max-async-uploads", s3_max_async_uploads);
+DEFINE_INT_TEST(test_init_s3_connect_timeout, "s3-connect-timeout", s3_connect_timeout);
 
 Suite* backup_conf_suite()
 {
@@ -781,6 +818,17 @@ Suite* backup_conf_suite()
 	tcase_add_test(tc_init, test_init_after_digest);
 	tcase_add_test(tc_init, test_init_max_recs);
 	tcase_add_test(tc_init, test_init_records_per_second);
+	tcase_add_test(tc_init, test_init_compress_mode);
+
+	tcase_add_test(tc_init, test_init_s3_region);
+	tcase_add_test(tc_init, test_init_s3_bucket);
+	tcase_add_test(tc_init, test_init_s3_profile);
+	tcase_add_test(tc_init, test_init_s3_endpoint_override);
+	tcase_add_test(tc_init, test_init_s3_min_part_size);
+	tcase_add_test(tc_init, test_init_s3_max_async_downloads);
+	tcase_add_test(tc_init, test_init_s3_max_async_uploads);
+	tcase_add_test(tc_init, test_init_s3_connect_timeout);
+	tcase_add_test(tc_init, test_init_s3_log_level);
 	suite_add_tcase(s, tc_init);
 
 	return s;