[PATCH] dnscache: merge similar outgoing udp packets

This patch takes a slightly different approach to merging
than the previous qmerge patch: rather than merging between
the "query" and "dns_transmit" layers, it actually
incorporates the merging into the dns_transmit layer.

This makes it a much more invasive and messy patch, but
means that it should correctly handle all situations
(in particular, ones where sending to two different IPs will
gives different results, generally because one server is
down or lame).

The general strategy is this:

  - when a dns_transmit is about to send() a UDP query, it
    checks for an equivalent inprogress query.

    If it finds one, it marks the inprogress dns_transmit as
    its "master", and itself as a "slave" of the master.

    If it doesn't, it sends the packet and registers itself
    as inprogress.

  - when a dns_transmit recv()s a UDP packet, it creates a
    copy of the packet for each of its slaves

Two outgoing packets are considered equivalent if:

  1. they are going to the same IP
  2. they have the same qtype
  3. they have the same qname

Because this change affects the dns library itself, this
change can potentially affect not just dnscache, but all of
the diagnostic tools. To address this, merging must be
enabled explicitly by the caller; dnscache will enable
merging if the MERGEQUERIES environment variable is set.

I tried to keep the patch as small and simple as possible so
that its correctness could be verified by readers.  There
are a few places where performance might benefit from making
it more complex:

  - dns_transmit now knows the value of MAXUDP, since it is
    an upper bound on the number of slaves and inprogress
    queries. As a result:

      - each non-merging program which uses the dns library
        wastes MAXUDP*sizeof(pointer) bytes of static memory
        for the inprogress list

      - each dns_transmit uses an extra MAXUDP*sizeof(pointer)
        bytes for its slaves list. For dnscache, this
        translates to 160K total on a 32-bit platform with
        default MAXUDP.

    Both could be avoided by using dynamic allocation.

  - We have to do an O(MAXUDP) linear search to find similar
    inprogress queries (non-merge-enabled users of the
    library avoid paying this price, though).

    This could be avoided by using a data structure with a
    fast key lookup for inprogress.

This patch is only lightly tested. Use on production servers at your own
risk (and please report to the list if you have success using it).

---
 clients.h      |    7 ++++
 dns.h          |    7 ++++
 dns_transmit.c |  105 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 dnscache.c     |    4 +-
 log.c          |    6 +++
 log.h          |    1 +
 6 files changed, 126 insertions(+), 4 deletions(-)
 create mode 100644 clients.h

diff --git a/clients.h b/clients.h
new file mode 100644
index 0000000..983a4ad
--- /dev/null
+++ b/clients.h
@@ -0,0 +1,7 @@
+#ifndef CLIENTS_H
+#define CLIENTS_H
+
+#define MAXUDP 200
+#define MAXTCP 20
+
+#endif /* CLIENTS_H */
diff --git a/dns.h b/dns.h
index 2f899ef..157d39b 100644
--- a/dns.h
+++ b/dns.h
@@ -4,6 +4,7 @@
 #include "stralloc.h"
 #include "iopause.h"
 #include "taia.h"
+#include "clients.h"
 
 #define DNS_C_IN "\0\1"
 #define DNS_C_ANY "\0\377"
@@ -37,8 +38,14 @@ struct dns_transmit {
   const char *servers;
   char localip[4];
   char qtype[2];
+  struct dns_transmit *master;
+  struct dns_transmit *slaves[MAXUDP];
+  int nslaves;
 } ;
 
+extern void dns_enable_merge(void (*logger)(const char *, const char *,
+      const char *));
+
 extern void dns_random_init(const char *);
 extern unsigned int dns_random(unsigned int);
 
diff --git a/dns_transmit.c b/dns_transmit.c
index 4d6e39f..3984776 100644
--- a/dns_transmit.c
+++ b/dns_transmit.c
@@ -7,6 +7,61 @@
 #include "byte.h"
 #include "uint16.h"
 #include "dns.h"
+#include "strerr.h"
+
+static int merge_enable;
+static void (*merge_logger)(const char *, const char *, const char *);
+void dns_enable_merge(void (*f)(const char *, const char *, const char *))
+{
+  merge_enable = 1;
+  merge_logger = f;
+}
+
+static int merge_equal(struct dns_transmit *a, struct dns_transmit *b)
+{
+  const char *ip1 = a->servers + 4 * a->curserver;
+  const char *ip2 = b->servers + 4 * b->curserver;
+  return
+    byte_equal(ip1, 4, ip2) &&
+    byte_equal(a->qtype, 2, b->qtype) &&
+    dns_domain_equal(a->query + 14, b->query + 14);
+}
+
+struct dns_transmit *inprogress[MAXUDP];
+
+static int try_merge(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (!inprogress[i]) continue;
+    if (!merge_equal(d, inprogress[i])) continue;
+    d->master = inprogress[i];
+    inprogress[i]->slaves[inprogress[i]->nslaves++] = d;
+    return 1;
+  }
+  return 0;
+}
+
+static void register_inprogress(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (!inprogress[i]) {
+      inprogress[i] = d;
+      return;
+    }
+  }
+  strerr_die1x(100, "BUG: out of inprogress slots");
+}
+
+static void unregister_inprogress(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (inprogress[i] == d)
+      inprogress[i] = 0;
+  }
+}
 
 static int serverwantstcp(const char *buf,unsigned int len)
 {
@@ -59,8 +114,28 @@ static void packetfree(struct dns_transmit *d)
   d->packet = 0;
 }
 
+static void mergefree(struct dns_transmit *d)
+{
+  int i;
+  if (merge_enable)
+    unregister_inprogress(d);
+  /* unregister us from our master */
+  if (d->master) {
+    for (i = 0; i < d->master->nslaves; i++)
+      if (d->master->slaves[i] == d)
+        d->master->slaves[i] = 0;
+  }
+  /* and unregister all of our slaves from us */
+  for (i = 0; i < d->nslaves; i++) {
+    if (d->slaves[i])
+      d->slaves[i]->master = NULL;
+  }
+  d->nslaves = 0;
+}
+
 static void queryfree(struct dns_transmit *d)
 {
+  mergefree(d);
   if (!d->query) return;
   alloc_free(d->query);
   d->query = 0;
@@ -99,11 +174,18 @@ static int thisudp(struct dns_transmit *d)
   const char *ip;
 
   socketfree(d);
+  mergefree(d);
 
   while (d->udploop < 4) {
     for (;d->curserver < 16;++d->curserver) {
       ip = d->servers + 4 * d->curserver;
       if (byte_diff(ip,4,"\0\0\0\0")) {
+        if (merge_enable && try_merge(d)) {
+          if (merge_logger)
+            merge_logger(ip, d->qtype, d->query + 14);
+          return 0;
+        }
+
 	d->query[2] = dns_random(256);
 	d->query[3] = dns_random(256);
   
@@ -118,6 +200,8 @@ static int thisudp(struct dns_transmit *d)
             taia_uint(&d->deadline,timeouts[d->udploop]);
             taia_add(&d->deadline,&d->deadline,&now);
             d->tcpstate = 0;
+            if (merge_enable)
+              register_inprogress(d);
             return 0;
           }
   
@@ -226,8 +310,12 @@ void dns_transmit_io(struct dns_transmit *d,iopause_fd *x,struct taia *deadline)
   x->fd = d->s1 - 1;
 
   switch(d->tcpstate) {
-    case 0: case 3: case 4: case 5:
-      x->events = IOPAUSE_READ;
+    case 0:
+      if (d->master) return;
+      if (d->packet) { taia_now(deadline); return; }
+      /* otherwise, fall through */
+    case 3: case 4: case 5:
+        x->events = IOPAUSE_READ;
       break;
     case 1: case 2:
       x->events = IOPAUSE_WRITE;
@@ -244,10 +332,14 @@ int dns_transmit_get(struct dns_transmit *d,const iopause_fd *x,const struct tai
   unsigned char ch;
   int r;
   int fd;
+  int i;
 
   errno = error_io;
   fd = d->s1 - 1;
 
+  if (d->tcpstate == 0 && d->master) return 0;
+  if (d->tcpstate == 0 && d->packet) return 1;
+
   if (!x->revents) {
     if (taia_less(when,&d->deadline)) return 0;
     errno = error_timeout;
@@ -279,6 +371,15 @@ have sent query to curserver on UDP socket s
     d->packet = alloc(d->packetlen);
     if (!d->packet) { dns_transmit_free(d); return -1; }
     byte_copy(d->packet,d->packetlen,udpbuf);
+
+    for (i = 0; i < d->nslaves; i++) {
+      if (!d->slaves[i]) continue;
+      d->slaves[i]->packetlen = d->packetlen;
+      d->slaves[i]->packet = alloc(d->packetlen);
+      if (!d->slaves[i]->packet) { dns_transmit_free(d->slaves[i]); continue; }
+      byte_copy(d->slaves[i]->packet,d->packetlen,udpbuf);
+    }
+
     queryfree(d);
     return 1;
   }
diff --git a/dnscache.c b/dnscache.c
index 8c899a3..c8db179 100644
--- a/dnscache.c
+++ b/dnscache.c
@@ -54,7 +54,6 @@ uint64 numqueries = 0;
 
 static int udp53;
 
-#define MAXUDP 200
 static struct udpclient {
   struct query q;
   struct taia start;
@@ -131,7 +130,6 @@ void u_new(void)
 
 static int tcp53;
 
-#define MAXTCP 20
 struct tcpclient {
   struct query q;
   struct taia start;
@@ -435,6 +433,8 @@ int main()
     response_hidettl();
   if (env_get("FORWARDONLY"))
     query_forwardonly();
+  if (env_get("MERGEQUERIES"))
+    dns_enable_merge(log_merge);
 
   if (!roots_init())
     strerr_die2sys(111,FATAL,"unable to read servers: ");
diff --git a/log.c b/log.c
index c43e8b0..3e1c674 100644
--- a/log.c
+++ b/log.c
@@ -150,6 +150,12 @@ void log_tx(const char *q,const char qtype[2],const char *control,const char ser
   line();
 }
 
+void log_merge(const char *addr, const char qtype[2], const char *q)
+{
+  string("merge "); ip(addr); space(); logtype(qtype); space(); name(q);
+  line();
+}
+
 void log_cachedanswer(const char *q,const char type[2])
 {
   string("cached "); logtype(type); space();
diff --git a/log.h b/log.h
index fe62fa3..dd6408a 100644
--- a/log.h
+++ b/log.h
@@ -18,6 +18,7 @@ extern void log_cachednxdomain(const char *);
 extern void log_cachedns(const char *,const char *);
 
 extern void log_tx(const char *,const char *,const char *,const char *,unsigned int);
+extern void log_merge(const char *, const char *, const char *);
 
 extern void log_nxdomain(const char *,const char *,unsigned int);
 extern void log_nodata(const char *,const char *,const char *,unsigned int);
-- 
1.6.2.1.591.geb450