From d12cd8e9fb77e888850488abc5d636a33c475a8a Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Thu, 26 May 2016 19:42:36 +0200
Subject: [PATCH] dnsdist: Keep the servers ordered inside pools

Several policies expect the servers to be ordered on their 'order'.
In addition to that, we keep the servers in a `NumberedVector` to
be able to pass them as a Lua table to Lua custom policies, and
that means we need to get the numbers right there too, especially
when we remove a server from a pool.
---
 pdns/dnsdist.cc                          | 28 ++++++++--
 regression-tests.dnsdist/test_Routing.py | 65 ++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 1035b6beb..a496e2b6e 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -624,6 +624,15 @@ void addServerToPool(pools_t& pools, const string& poolName, std::shared_ptr<Dow
     vinfolog("Adding server to default pool");
   }
   pool->servers.push_back(make_pair(++count, server));
+  /* we need to reorder based on the server 'order' */
+  std::stable_sort(pool->servers.begin(), pool->servers.end(), [](const std::pair<unsigned int,std::shared_ptr<DownstreamState> >& a, const std::pair<unsigned int,std::shared_ptr<DownstreamState> >& b) {
+      return a.second->order < b.second->order;
+    });
+  /* and now we need to renumber for Lua (custom policies) */
+  size_t idx = 1;
+  for (auto& server : pool->servers) {
+    server.first = idx++;
+  }
 }
 
 void removeServerFromPool(pools_t& pools, const string& poolName, std::shared_ptr<DownstreamState> server)
@@ -637,10 +646,21 @@ void removeServerFromPool(pools_t& pools, const string& poolName, std::shared_pt
     vinfolog("Removing server from default pool");
   }
 
-  for (NumberedVector<shared_ptr<DownstreamState> >::iterator it = pool->servers.begin(); it != pool->servers.end(); it++) {
-    if (it->second == server) {
-      pool->servers.erase(it);
-      break;
+  size_t idx = 1;
+  bool found = false;
+  for (NumberedVector<shared_ptr<DownstreamState> >::iterator it = pool->servers.begin(); it != pool->servers.end();) {
+    if (found) {
+      /* we need to renumber the servers placed
+         after the removed one, for Lua (custom policies) */
+      it->first = idx++;
+      it++;
+    }
+    else if (it->second == server) {
+      it = pool->servers.erase(it);
+      found = true;
+    } else {
+      idx++;
+      it++;
     }
   }
 }
diff --git a/regression-tests.dnsdist/test_Routing.py b/regression-tests.dnsdist/test_Routing.py
index f51be3470..807cdea4e 100644
--- a/regression-tests.dnsdist/test_Routing.py
+++ b/regression-tests.dnsdist/test_Routing.py
@@ -296,3 +296,68 @@ class TestRoutingRoundRobinLBOneDown(DNSDistTest):
             total += value
 
         self.assertEquals(total, numberOfQueries * 2)
+
+class TestRoutingOrder(DNSDistTest):
+
+    _testServer2Port = 5351
+    _config_params = ['_testServerPort', '_testServer2Port']
+    _config_template = """
+    setServerPolicy(firstAvailable)
+    s1 = newServer{address="127.0.0.1:%s", order=2}
+    s1:setUp()
+    s2 = newServer{address="127.0.0.1:%s", order=1}
+    s2:setUp()
+    """
+
+    @classmethod
+    def startResponders(cls):
+        print("Launching responders..")
+        cls._UDPResponder = threading.Thread(name='UDP Responder', target=cls.UDPResponder, args=[cls._testServerPort])
+        cls._UDPResponder.setDaemon(True)
+        cls._UDPResponder.start()
+        cls._UDPResponder2 = threading.Thread(name='UDP Responder 2', target=cls.UDPResponder, args=[cls._testServer2Port])
+        cls._UDPResponder2.setDaemon(True)
+        cls._UDPResponder2.start()
+
+        cls._TCPResponder = threading.Thread(name='TCP Responder', target=cls.TCPResponder, args=[cls._testServerPort])
+        cls._TCPResponder.setDaemon(True)
+        cls._TCPResponder.start()
+
+        cls._TCPResponder2 = threading.Thread(name='TCP Responder 2', target=cls.TCPResponder, args=[cls._testServer2Port])
+        cls._TCPResponder2.setDaemon(True)
+        cls._TCPResponder2.start()
+
+    def testOrder(self):
+        """
+        Routing: firstAvailable policy based on 'order'
+
+        Send 50 A queries to "order.routing.tests.powerdns.com.",
+        check that dnsdist routes all of it to the second backend
+        because it has the lower order value.
+        """
+        numberOfQueries = 50
+        name = 'order.routing.tests.powerdns.com.'
+        query = dns.message.make_query(name, 'A', 'IN')
+        response = dns.message.make_response(query)
+        rrset = dns.rrset.from_text(name,
+                                    60,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.A,
+                                    '192.0.2.1')
+        response.answer.append(rrset)
+
+        for _ in range(numberOfQueries):
+            (receivedQuery, receivedResponse) = self.sendUDPQuery(query, response)
+            receivedQuery.id = query.id
+            self.assertEquals(query, receivedQuery)
+            self.assertEquals(response, receivedResponse)
+            (receivedQuery, receivedResponse) = self.sendTCPQuery(query, response)
+            receivedQuery.id = query.id
+            self.assertEquals(query, receivedQuery)
+            self.assertEquals(response, receivedResponse)
+
+        total = 0
+        self.assertEquals(self._responsesCounter['UDP Responder'], 0)
+        self.assertEquals(self._responsesCounter['UDP Responder 2'], numberOfQueries)
+        self.assertEquals(self._responsesCounter['TCP Responder'], 0)
+        self.assertEquals(self._responsesCounter['TCP Responder 2'], numberOfQueries)
-- 
2.49.0