From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 26 Sep 2016 16:08:43 +0000 (+0100)
Subject: [Fix] Do not replan retransmits if merely one server is defined
X-Git-Tag: 1.4.0~370
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=502559ceab99ae36ae0200f98aa2ac14c99b0255;p=rspamd.git

[Fix] Do not replan retransmits if merely one server is defined
---

diff --git a/contrib/librdns/rdns.h b/contrib/librdns/rdns.h
index 7e5843183..e5fc82fe4 100644
--- a/contrib/librdns/rdns.h
+++ b/contrib/librdns/rdns.h
@@ -164,6 +164,7 @@ struct rdns_upstream_context {
 			size_t len, void *ups_data);
 	struct rdns_upstream_elt* (*select_retransmit)(const char *name,
 			size_t len, void *ups_data);
+	unsigned int (*count)(void *ups_data);
 	void (*ok)(struct rdns_upstream_elt *elt, void *ups_data);
 	void (*fail)(struct rdns_upstream_elt *elt, void *ups_data);
 };
diff --git a/contrib/librdns/resolver.c b/contrib/librdns/resolver.c
index 329caa620..8c5df913f 100644
--- a/contrib/librdns/resolver.c
+++ b/contrib/librdns/resolver.c
@@ -295,6 +295,7 @@ rdns_process_timer (void *arg)
 	bool renew = false;
 	struct rdns_resolver *resolver;
 	struct rdns_server *serv = NULL;
+	unsigned cnt;
 
 	req->retransmits --;
 	resolver = req->resolver;
@@ -318,44 +319,57 @@ rdns_process_timer (void *arg)
 	}
 
 	if (!req->io->active || req->retransmits == 1) {
-		/* Do not reschedule IO requests on inactive sockets */
-		rdns_debug ("reschedule request with id: %d", (int)req->id);
-		rdns_request_unschedule (req);
-		REF_RELEASE (req->io);
 
 		if (resolver->ups) {
-			struct rdns_upstream_elt *elt;
+			cnt = resolver->ups->count (resolver->ups->data);
+		}
+		else {
+			cnt = 0;
+			UPSTREAM_FOREACH (resolver->servers, serv) {
+				cnt ++;
+			}
+		}
 
-			elt = resolver->ups->select_retransmit (req->requested_names[0].name,
-					req->requested_names[0].len, resolver->ups->data);
+		if (!req->io->active || cnt > 1) {
+			/* Do not reschedule IO requests on inactive sockets */
+			rdns_debug ("reschedule request with id: %d", (int)req->id);
+			rdns_request_unschedule (req);
+			REF_RELEASE (req->io);
 
-			if (elt) {
-				serv = elt->server;
-				serv->ups_elt = elt;
+			if (resolver->ups) {
+				struct rdns_upstream_elt *elt;
+
+				elt = resolver->ups->select_retransmit (req->requested_names[0].name,
+						req->requested_names[0].len, resolver->ups->data);
+
+				if (elt) {
+					serv = elt->server;
+					serv->ups_elt = elt;
+				}
+				else {
+					UPSTREAM_SELECT_ROUND_ROBIN (resolver->servers, serv);
+				}
 			}
 			else {
 				UPSTREAM_SELECT_ROUND_ROBIN (resolver->servers, serv);
 			}
-		}
-		else {
-			UPSTREAM_SELECT_ROUND_ROBIN (resolver->servers, serv);
-		}
 
-		if (serv == NULL) {
-			rdns_warn ("cannot find suitable server for request");
-			rep = rdns_make_reply (req, RDNS_RC_SERVFAIL);
-			req->state = RDNS_REQUEST_REPLIED;
-			req->func (rep, req->arg);
-			REF_RELEASE (req);
+			if (serv == NULL) {
+				rdns_warn ("cannot find suitable server for request");
+				rep = rdns_make_reply (req, RDNS_RC_SERVFAIL);
+				req->state = RDNS_REQUEST_REPLIED;
+				req->func (rep, req->arg);
+				REF_RELEASE (req);
 
-			return;
-		}
+				return;
+			}
 
-		/* Select random IO channel */
-		req->io = serv->io_channels[ottery_rand_uint32 () % serv->io_cnt];
-		req->io->uses ++;
-		REF_RETAIN (req->io);
-		renew = true;
+			/* Select random IO channel */
+			req->io = serv->io_channels[ottery_rand_uint32 () % serv->io_cnt];
+			req->io->uses ++;
+			REF_RETAIN (req->io);
+			renew = true;
+		}
 	}
 
 	/*
diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index 13d1b8309..34b685260 100644
--- a/src/libserver/dns.c
+++ b/src/libserver/dns.c
@@ -30,12 +30,14 @@ static void rspamd_dns_upstream_ok (struct rdns_upstream_elt *elt,
 		void *ups_data);
 static void rspamd_dns_upstream_fail (struct rdns_upstream_elt *elt,
 		void *ups_data);
+static unsigned int rspamd_dns_upstream_count (void *ups_data);
 
 static struct rdns_upstream_context rspamd_ups_ctx = {
 		.select = rspamd_dns_select_upstream,
 		.select_retransmit = rspamd_dns_select_upstream_retransmit,
 		.ok = rspamd_dns_upstream_ok,
 		.fail = rspamd_dns_upstream_fail,
+		.count = rspamd_dns_upstream_count,
 		.data = NULL
 };
 
@@ -340,3 +342,11 @@ rspamd_dns_upstream_fail (struct rdns_upstream_elt *elt,
 
 	rspamd_upstream_fail (up);
 }
+
+static unsigned int
+rspamd_dns_upstream_count (void *ups_data)
+{
+	struct upstream_list *ups = ups_data;
+
+	return rspamd_upstreams_alive (ups);
+}