#include <sys/poll.h>
#include <sys/epoll.h>

/*
  Drop in poll->epoll upgrader
  By: David J. Schwartz, Copyright 2008, <davids@webmaster.com>
  May or may not be suitable for your application! CAUTION!

  Redistribution and use in source and binary forms, with or without
  modification, are permitted. This software is provided 'as is' and
  no representation is made that it is suitable for any purpose. It has
  not been debugged or tested in any way.

  CAUTION: This drop-in replacement for 'poll' using 'epoll' may or
  may not improve performance and may or may not break your code!
  It is definitely inferior to modifying the code itself to use 'epoll'.
  It will fail horribly if the program using it generates calls to
  'poll_epoll' from more than one place with fundamentally different
  'ufds' parameters!

  CAUTION: Known issue: If we return zero from our poll function,
  we may return too early. In some cases, we will return immediately
  when we shouldn't. This occurs when the return value from epoll_wait
  is non-zero but the return value from poll_epoll is zero.

  Version: 1.4
*/

/* Absolute maximum number of file descriptors */
#ifndef EPOLL_MAX_FD
#define EPOLL_MAX_FD 65536
#endif

/* Hint maximum concurrent connections */
#ifndef EPOLL_CONC_EV
#define EPOLL_CONC_EV 32768
#endif

/* Threshold for cutover to epoll - set to 1 for testing */
#ifndef EPOLL_THRESH
#define EPOLL_THRESH 16
#endif

static int epoll_fd=-1;
static struct epoll_event epoll_ev[EPOLL_CONC_EV];
static int reg_ev[EPOLL_MAX_FD]; /* registered events, by fd */

int close_epoll(int fd)
{ /* CAUTION: To workaround a potential lost event problem,
     you should replace all 'close' calls for descriptors
     that may have been 'poll'ed with calls to close_epoll!
  */
 if(fd>=0) reg_ev[fd]=0;
 return close(fd);
}

int poll_epoll(struct pollfd *ufds, nfds_t nfds, int timeout)
{ /* CAUTION: Not re-entrant. Cannot be called from multiple threads */
 int i, j, rc, ec;
 
 /*
    First step: Make sure we have an epoll fd. Sanity checks.
    Initialization if needed.
 */

 if( (epoll_fd==-2) || (nfds<EPOLL_THRESH) ) return poll(ufds, nfds, timeout);
 if(epoll_fd<0)
 {
  epoll_fd=epoll_create(EPOLL_CONC_EV);
  if(epoll_fd<0)
  {
   epoll_fd=-2;
   return poll(ufds, nfds, timeout);
  }
  for(i=0; i<EPOLL_MAX_FD; i++) reg_ev[i]=0;
 }


 /*
   Second step: Add any missing events.
   Extra events are okay, we'll catch them in step four.
 */
 for(i=0; i<nfds; i++)
 {
  ufds[i].revents=0;
  int fd=ufds[i].fd;
  if(fd>=0)
  { /* This is a valid test for events */
   ufds[i].revents=0;
   if(reg_ev[fd]!=ufds[i].events)
   {
    int missing_events=ufds[i].events&~reg_ev[fd];
    if(missing_events!=0)
    { /* we need to add at least one event */
     struct epoll_event ev;
     ev.data.fd=fd;
     ev.events=reg_ev[fd]|missing_events;
     epoll_ctl(epoll_fd, (reg_ev[fd]==0) ? EPOLL_CTL_ADD : EPOLL_CTL_MOD,
      fd, &ev);
     reg_ev[fd]=ev.events;
    }
   }
  }
 }
 
 /*
  Third step: Call epoll. Return error, if any
 */
 tol:
 rc=epoll_wait(epoll_fd, epoll_ev, EPOLL_CONC_EV, timeout);
 if(rc<=0) return rc;
 
 /*
  Fourth step: Report events, unregister unwanted ones
 */
 ec=0;
 for(i=0; i<rc; i++)
 {
  int fd=epoll_ev[i].data.fd;
  int needed_events=0;
  int remaining_events=epoll_ev[i].events;
  for(j=0; j<nfds; j++)
   if(fd==ufds[j].fd)
   { /* we found a hit on an FD that we are interested in */
    int report=ufds[j].events|POLLHUP|POLLNVAL|POLLERR;
    needed_events|=report; // still interested in these
    remaining_events&=~report; // still not interested in these
    ufds[j].revents=epoll_ev[i].events&report;
    if(ufds[j].revents!=0) ec++;
   }
  if(remaining_events)
  { /* we found some events we could not report */
   struct epoll_event ev;
   ev.data.fd=fd;
   ev.events=reg_ev[fd]=needed_events;
   epoll_ctl(epoll_fd, (needed_events==0) ? EPOLL_CTL_DEL : EPOLL_CTL_MOD,
    fd, &ev);
  }
 }

 if ((rc==0) && (timeout==-1)) goto tol;

 return ec;
}