Add basic (non weight adjusted) shuffling for SrvRecords selection

Also add "Debug::Acquire::SrvRecs" debug option and the option
"Acquire::EnableSrvRecods" to allow disabling this lookup.

4 files changed, 97 insertions, 5 deletions

diff --git a/apt-pkg/contrib/srvrec.cc b/apt-pkg/contrib/srvrec.cc
index 85241c1d7..b4a3d97d2 100644
--- a/apt-pkg/contrib/srvrec.cc
+++ b/apt-pkg/contrib/srvrec.cc
@@ -13,13 +13,18 @@
 #include <netinet/in.h>
 #include <arpa/nameser.h>
 #include <resolv.h>
+#include <chrono>
 
 #include <algorithm>
 
-#include <apt-pkg/strutl.h>
+#include <apt-pkg/configuration.h>
 #include <apt-pkg/error.h>
+#include <apt-pkg/strutl.h>
+
+
 #include "srvrec.h"
 
+
 bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result)
 {
    std::string target;
@@ -112,6 +117,29 @@ bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
    // sort them by priority
    std::stable_sort(Result.begin(), Result.end());
 
+   for(std::vector<SrvRec>::iterator I = Result.begin();
+      I != Result.end(); ++I)
+   {
+      if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+      {
+         std::cerr << "SrvRecs: got " << I->target
+                   << " prio: " << I->priority
+                   << " weight: " << I->weight
+                   << std::endl;
+      }
+   }
+
+   return true;
+}
+
+SrvRec PopFromSrvRecs(std::vector<SrvRec> &Recs)
+{
+   // FIXME: instead of the simplistic shuffle below use the algorithm
+   //        described in rfc2782 (with weights)
+   //        and figure out how the weights need to be adjusted if
+   //        a host refuses connections
+
+#if 0  // all code below is only needed for the weight adjusted selection 
    // assign random number ranges
    int prev_weight = 0;
    int prev_priority = 0;
@@ -124,6 +152,12 @@ bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
       I->random_number_range_end = prev_weight + I->weight;
       prev_weight = I->random_number_range_end;
       prev_priority = I->priority;
+
+      if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+         std::cerr << "SrvRecs: got " << I->target
+                   << " prio: " << I->priority
+                   << " weight: " << I->weight
+                   << std::endl;
    }
 
    // go over the code in reverse order and note the max random range
@@ -136,8 +170,27 @@ bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result)
          max = I->random_number_range_end;
       I->random_number_range_max = max;
    }
+#endif
 
-   // FIXME: now shuffle 
+   // shuffle in a very simplistic way for now (equal weights)
+   std::vector<SrvRec>::iterator I, J;
+   I = J = Recs.begin();
+   for(;I != Recs.end(); ++I)
+   {
+      if(I->priority != J->priority)
+         break;
+   }
 
-   return true;
+   // FIXME: meeeeh, where to init this properly
+   unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
+   std::shuffle(J, I, std::default_random_engine(seed));
+
+   // meh, no pop_front() in std::vector?
+   SrvRec selected = *Recs.begin();
+   Recs.erase(Recs.begin());
+
+   if (_config->FindB("Debug::Acquire::SrvRecs", false) == true)
+      std::cerr << "PopFromSrvRecs: selecting " << selected.target << std::endl;
+
+   return selected;
 }
diff --git a/apt-pkg/contrib/srvrec.h b/apt-pkg/contrib/srvrec.h
index 9323623a5..e07edc683 100644
--- a/apt-pkg/contrib/srvrec.h
+++ b/apt-pkg/contrib/srvrec.h
@@ -39,4 +39,9 @@ bool GetSrvRecords(std::string name, std::vector<SrvRec> &Result);
  */
 bool GetSrvRecords(std::string host, int port, std::vector<SrvRec> &Result);
 
+/** \brief Pop a single SRV record from the vector of SrvRec taking
+ *         priority and weight into account
+ */
+SrvRec PopFromSrvRecs(std::vector<SrvRec> &Recs);
+
 #endif
diff --git a/methods/connect.cc b/methods/connect.cc
index 9e0f01ee3..5612af6ec 100644
--- a/methods/connect.cc
+++ b/methods/connect.cc
@@ -272,7 +272,8 @@ bool Connect(std::string Host,int Port,const char *Service,
    if(LastHost != Host || LastPort != Port)
    {
       SrvRecords.clear();
-      bool res = GetSrvRecords(Host, DefPort, SrvRecords);
+      if (_config->FindB("Acquire::EnableSrvRecods", true) == true)
+         GetSrvRecords(Host, DefPort, SrvRecords);
    }
    // we have no SrvRecords for this host, connect right away
    if(SrvRecords.size() == 0)
@@ -282,7 +283,7 @@ bool Connect(std::string Host,int Port,const char *Service,
    // try to connect in the priority order of the srv records
    while(SrvRecords.size() > 0)
    {
-      Host = SrvRecords[0].target;
+      Host = PopFromSrvRecs(SrvRecords).target;
       if(ConnectToHostname(Host, Port, Service, DefPort, Fd, TimeOut, Owner))
          return true;
 
diff --git a/test/libapt/srvrecs_test.cc b/test/libapt/srvrecs_test.cc
new file mode 100644
index 000000000..7e43cc757
--- /dev/null
+++ b/test/libapt/srvrecs_test.cc
@@ -0,0 +1,33 @@
+#include <config.h>
+
+#include <apt-pkg/srvrec.h>
+
+#include <string>
+#include <iostream>
+
+#include <gtest/gtest.h>
+
+TEST(SrvRecTest, PopFromSrvRecs)
+{
+   // the PopFromSrvRecs() is using a random number so we must
+   // run it a bunch of times to ensure we are not fooled by randomness
+   std::set<std::string> selected;
+   for(int i=0;i<100;i++)
+   {
+      std::vector<SrvRec> Meep;
+      SrvRec foo = {target:"foo", priority: 20, weight: 0, port: 80};
+      Meep.push_back(foo);
+      
+      SrvRec bar = {target:"bar", priority: 20, weight: 0, port: 80};
+      Meep.push_back(bar);
+
+      EXPECT_EQ(Meep.size(), 2);
+      SrvRec result = PopFromSrvRecs(Meep);
+      selected.insert(result.target);
+      // ensure that pop removed one element
+      EXPECT_EQ(Meep.size(), 1);
+   }
+
+   // ensure that after enough runs we end up with both selected
+   EXPECT_EQ(selected.size(), 2);
+}