Added connection timeout to check_nscp_nrpe

clients/nrpe/check_nrpe.cpp

@@ -207,7 +207,9 @@ boost::program_options::options_description add_client_options(client::destinati
 	desc.add_options()
 		("log", po::value<std::string>()->notifier(boost::bind(&client::destination_container::set_string_data, &source, "log", _1)),
 			"Set log level")
-		;
+        ("con-timeout", po::value<std::string>()->notifier(boost::bind(&client::destination_container::set_string_data, &source, "con_timeout", _1)),
+            "Set connection timeout in seconds")
+        ;
 	return desc;
 }
 

include/client/command_line_parser.hpp

@@ -47,9 +47,10 @@ namespace client {
 		net::url address;
 		int timeout;
 		int retry;
+		int connection_timeout;
 		data_map data;
 
-		destination_container() : timeout(10), retry(2) {}
+		destination_container() : timeout(10), retry(2), connection_timeout(-1) {}
 
 		void apply(nscapi::settings_objects::object_instance obj) {
 			BOOST_FOREACH(const nscapi::settings_objects::options_map::value_type &k, obj->get_options()) {
@@ -90,6 +91,9 @@ namespace client {
 		bool has_protocol() const {
 			return !address.protocol.empty();
 		}
+		void set_connection_timeout(int conn_timeout) {
+			connection_timeout = conn_timeout;
+		}
 
 		static bool to_bool(std::string value, bool def = false) {
 			if (value.empty())
@@ -135,6 +139,8 @@ namespace client {
 				timeout = to_int(value, timeout);
 			else if (key == "retry")
 				retry = to_int(value, retry);
+			else if (key == "con_timeout")
+				set_connection_timeout(std::stoi(value));
 			else
 				data[key] = value;
 		}

include/socket/client.hpp

@@ -25,6 +25,7 @@
 #include <iostream>
 
 using boost::asio::ip::tcp;
+using boost::asio::deadline_timer;
 
 namespace socket_helpers {
 	namespace client {
@@ -33,6 +34,7 @@ namespace socket_helpers {
 		private:
 			boost::asio::io_service &io_service_;
 			boost::asio::deadline_timer timer_;
+			boost::asio::deadline_timer con_timer_;
 			boost::posix_time::time_duration timeout_;
 			boost::shared_ptr<typename protocol_type::client_handler> handler_;
 			protocol_type protocol_;
@@ -44,6 +46,7 @@ namespace socket_helpers {
 			connection(boost::asio::io_service &io_service, boost::posix_time::time_duration timeout, boost::shared_ptr<typename protocol_type::client_handler> handler)
 				: io_service_(io_service)
 				, timer_(io_service)
+				, con_timer_(io_service)
 				, timeout_(timeout)
 				, handler_(handler)
 				, protocol_(handler) {}
@@ -78,11 +81,25 @@ namespace socket_helpers {
 					timer_result_.reset(ec);
 				}
 			}
+			// Connection timeout functions
+			void connected( const boost::system::error_code ec,
+							boost::system::error_code *err) {
+				*err = ec;
+			}
+			void check_deadline() {
+				if (con_timer_.expires_at() <= deadline_timer::traits_type::now()) {
+					boost::system::error_code ignored_ec;
+					get_socket().close(ignored_ec);
+					con_timer_.expires_at(boost::posix_time::pos_infin);
+				}
 
+				// Put the actor back to sleep.
+				con_timer_.async_wait(bind(&connection::check_deadline, this));
+			}
 			//////////////////////////////////////////////////////////////////////////
 			// External API functions
 			//
-			virtual boost::system::error_code connect(std::string host, std::string port) {
+			virtual boost::system::error_code connect(std::string host, std::string port, int connection_timeout=-1) {
 				trace("connect(" + host + ", " + port + ")");
 				tcp::resolver resolver(io_service_);
 				tcp::resolver::query query(host, port, boost::asio::ip::resolver_query_base::numeric_service);
@@ -91,14 +108,39 @@ namespace socket_helpers {
 				tcp::resolver::iterator end;
 
 				boost::system::error_code error = boost::asio::error::host_not_found;
-				while (error && endpoint_iterator != end) {
-					get_socket().close();
-					get_socket().lowest_layer().connect(*endpoint_iterator++, error);
-				}
-				if (error) {
-					trace("Failed to connect to: " + host + ":" + port);
-					return error;
+
+				if (connection_timeout == -1) {
+					while (error && endpoint_iterator != end) {
+						get_socket().close();
+						get_socket().lowest_layer().connect(*endpoint_iterator++, error);
+					}
+
+					if (error) {
+						trace("Failed to connect to: " + host + ":" + port);
+						return error;
+					}
+				} else {
+					// Asynchronous connection
+					con_timer_.expires_at(boost::posix_time::pos_infin);
+					check_deadline();
+					con_timer_.expires_from_now(boost::posix_time::seconds(connection_timeout));
+					error = boost::asio::error::would_block;
+
+					boost::asio::async_connect( get_socket(),
+												endpoint_iterator,
+												boost::bind(&connection::connected,
+															this->shared_from_this(),
+															_1,
+															&error));
+					do io_service_.run_one();
+					while (error == boost::asio::error::would_block);
+
+					if (error || !get_socket().is_open()) {
+						trace("Failed to connect to: " + host + ":" + port);
+						return error;
+					}
 				}
+
 				protocol_.on_connect();
 				return error;
 			}
@@ -270,8 +312,8 @@ namespace socket_helpers {
 				}
 			}
 
-			virtual boost::system::error_code connect(std::string host, std::string port) {
-				boost::system::error_code error = connection_type::connect(host, port);
+			virtual boost::system::error_code connect(std::string host, std::string port, int connection_timeout=-1) {
+				boost::system::error_code error = connection_type::connect(host, port, connection_timeout);
 				if (error) {
 					this->log_error(__FILE__, __LINE__, "Failed to connect to server: " + utf8::utf8_from_native(error.message()));
 				}
@@ -336,7 +378,7 @@ namespace socket_helpers {
 
 			void connect() {
 				connection_.reset(create_connection());
-				boost::system::error_code error = connection_->connect(info_.get_address(), info_.get_port());
+				boost::system::error_code error = connection_->connect(info_.get_address(), info_.get_port(), info_.get_connection_timeout());
 				if (error) {
 					connection_.reset();
 					throw socket_helpers::socket_exception("Failed to connect to: " + info_.get_endpoint_string() + " :" + utf8::utf8_from_native(error.message()));

include/socket/socket_helpers.hpp

@@ -218,10 +218,11 @@ namespace socket_helpers {
 		unsigned int timeout;
 		int retry;
 		bool reuse;
+		unsigned int con_timeout;
 		ssl_opts ssl;
 		allowed_hosts_manager allowed_hosts;
 
-		connection_info() : back_log(backlog_default), port_("0"), thread_pool_size(0), timeout(30), retry(2), reuse(true) {}
+		connection_info() : back_log(backlog_default), port_("0"), thread_pool_size(0), timeout(30), retry(2), reuse(true), con_timeout(-1) {}
 
 		connection_info(const connection_info &other)
 			: address(other.address)
@@ -231,6 +232,7 @@ namespace socket_helpers {
 			, timeout(other.timeout)
 			, retry(other.retry)
 			, reuse(other.reuse)
+			, con_timeout(other.con_timeout)
 			, ssl(other.ssl)
 			, allowed_hosts(other.allowed_hosts) {}
 		connection_info& operator=(const connection_info &other) {
@@ -241,6 +243,7 @@ namespace socket_helpers {
 			timeout = other.timeout;
 			retry = other.retry;
 			reuse = other.reuse;
+			con_timeout = other.con_timeout;
 			ssl = other.ssl;
 			allowed_hosts = other.allowed_hosts;
 			return *this;
@@ -256,6 +259,9 @@ namespace socket_helpers {
 		std::string get_endpoint_string() const {
 			return address + ":" + get_port();
 		}
+		unsigned int get_connection_timeout() const {
+			return con_timeout;
+		}
 		long get_ctx_opts();
 
 		std::string to_string() const {

modules/NRPEClient/nrpe_client.hpp

@@ -70,6 +70,7 @@ namespace nrpe_client {
 			retry = target.retry;
 			buffer_length = target.get_int_data("payload length", 1024);
 			encoding = target.get_string_data("encoding");
+			con_timeout = source.connection_timeout;
 
 			if (target.has_data("no ssl"))
 				ssl.enabled = !target.get_bool_data("no ssl");