用timerfd实现非阻塞的sleep()/usleep()/nanosleep()

README.rst

@@ -35,6 +35,17 @@ Usage
 2. Make sure the dynamic module(e.g. libmemcached) is patched before using::
 
     assert greenify.patch_lib('/usr/lib/libmemcached.so')
+
+   By default greenify will patch these blocking operations:
+- socket `connect()`
+- socket io, etc. `read()`/`write()`/`recvfrom()`/`sendto()` and their families
+- io multiplexing, `select()` and `poll()` if available
+
+  If you also need `sleep()`/`usleep()`/`nanosleep()` calls to be patched, pass `patch_sleep=True` to `greenify.patch_lib()`::
+
+    assert greenify.patch_lib('/usr/lib/libmemcached.so', patch_sleep=True)
+
+  Note: currently patched `sleep()` calls cannot be interrupted by signals.
 
 3. Import and use the corresponding module, which is now gevent_ compatible.
 

greenify.pyx

@@ -26,6 +26,9 @@ cdef extern from "hook_greenify.h":
         FN_SENDTO
         FN_SELECT
         FN_POLL
+        FN_SLEEP
+        FN_USLEEP
+        FN_NANOSLEEP
 
     void* greenify_patch_lib(const char* library_filename, greenified_function_t fn)
 
@@ -75,16 +78,19 @@ def wait(watchers):
         for watcher in watchers:
             watcher.stop()
 
-cpdef patch_lib(library_path):
+cpdef patch_lib(library_path, patch_sleep=False):
     cdef char* path
     if isinstance(library_path, unicode):
         library_path = (<unicode>library_path).encode('utf8')
 
     path = library_path
     cdef bint result = False
-    for fn in (FN_CONNECT, FN_READ, FN_WRITE, FN_PREAD, FN_PWRITE, FN_READV,
+    fn_set = [FN_CONNECT, FN_READ, FN_WRITE, FN_PREAD, FN_PWRITE, FN_READV,
                FN_WRITEV, FN_RECV, FN_SEND, FN_RECVMSG, FN_SENDMSG,
-               FN_RECVFROM, FN_SENDTO, FN_SELECT, FN_POLL):
+               FN_RECVFROM, FN_SENDTO, FN_SELECT, FN_POLL]
+    if patch_sleep:
+        fn_set.extend((FN_SLEEP, FN_USLEEP, FN_NANOSLEEP))
+    for fn in fn_set:
         if NULL != greenify_patch_lib(path, fn):
             result = True
 

include/hook_greenify.h

@@ -25,6 +25,9 @@ typedef enum
     FN_SENDTO,
     FN_SELECT,
     FN_POLL,
+    FN_SLEEP,
+    FN_USLEEP,
+    FN_NANOSLEEP,
 } greenified_function_t;
 
 void* greenify_patch_lib(const char* library_filename, greenified_function_t fn);

include/libgreenify.h

@@ -36,6 +36,9 @@ int green_select(int nfds, fd_set *readfds, fd_set *writefds, fd_set *exceptfds,
 #ifndef NO_POLL
 int green_poll(struct pollfd *fds, nfds_t nfds, int timeout);
 #endif
+unsigned int green_sleep(unsigned int seconds);
+int green_usleep(useconds_t usec);
+int green_nanosleep(const struct timespec *duration, struct timespec *rem);
 
 struct greenify_watcher {
 	int fd;

src/hook_greenify.c

@@ -38,6 +38,12 @@ void* greenify_patch_lib(const char* library_filename, greenified_function_t fn)
         case FN_POLL:
             return _GREENIFY_PATCH_EXPAND(library_filename, poll);
 #endif
+        case FN_SLEEP:
+            return _GREENIFY_PATCH_EXPAND(library_filename, sleep);
+        case FN_USLEEP:
+            return _GREENIFY_PATCH_EXPAND(library_filename, usleep);
+        case FN_NANOSLEEP:
+            return _GREENIFY_PATCH_EXPAND(library_filename, nanosleep);
         default:
             return NULL;
     }

src/libgreenify.c

@@ -7,6 +7,7 @@
 #include <inttypes.h>
 #include <stdint.h>
 #include <time.h>
+#include <sys/timerfd.h>
 
 #ifdef DEBUG
 #include <unistd.h>
@@ -285,3 +286,63 @@ green_poll(struct pollfd *fds, nfds_t nfds, int timeout)
     return poll(fds, nfds, 0);
 }
 #endif
+
+
+#define handle_error(msg) \
+        do { perror(msg); exit(EXIT_FAILURE); } while (0)
+
+
+unsigned int green_sleep_impl(unsigned int seconds, unsigned int nanoseconds){
+    debug("Enter green_sleep_impl\n");
+    int fd = timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC);
+    if (fd == -1)
+        handle_error("timerfd create error");
+
+    struct itimerspec new_value;
+    new_value.it_interval.tv_sec = 0;
+    new_value.it_interval.tv_nsec = 0;
+
+    new_value.it_value.tv_sec = seconds;
+    new_value.it_value.tv_nsec = nanoseconds;
+
+    if (timerfd_settime(fd, 0, &new_value, NULL) == -1)
+        handle_error("timerfd settime error");
+
+    callback_single_watcher(fd, EVENT_READ, 0);
+
+    close(fd);
+
+    return 0;
+}
+
+unsigned int green_sleep(unsigned int seconds){
+    return green_sleep_impl(seconds, 0);
+}
+
+int green_usleep(useconds_t usec) {
+    unsigned int seconds = usec / 1000 / 1000;
+    unsigned int nanoseconds = usec % (1000 * 1000) * 1000;
+    return green_sleep_impl(seconds, nanoseconds);
+}
+
+int green_nanosleep(const struct timespec *duration, struct timespec *rem) {
+    debug("Enter green_nanosleep\n"); \
+    int fd = timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC);
+    if (fd == -1)
+        handle_error("timerfd create error");
+
+    struct itimerspec new_value;
+    new_value.it_interval.tv_sec = 0;
+    new_value.it_interval.tv_nsec = 0;
+
+    new_value.it_value = *duration;
+
+    if (timerfd_settime(fd, 0, &new_value, NULL) == -1)
+        handle_error("timerfd settime error");
+
+    callback_single_watcher(fd, EVENT_READ, 0);
+
+    close(fd);
+
+    return 0;
+}