zookeeper 1.0.6-java → 1.1.0-java

data/ext/{zookeeper_lib.c → event_lib.c}

@@ -7,20 +7,21 @@ This file contains three sets of helpers:
 
 wickman@twitter.com
 
-*********************************************************************************
-*
-* NOTE: be *very careful* in these functions, calling *ANY* ruby interpreter
-* function when you're not in an interpreter thread can hork ruby, trigger a
-* [BUG], corrupt the stack, kill your dog, knock up your daughter, etc. etc.
-*
-*********************************************************************************
+
+NOTE: be *very careful* in these functions, calling *ANY* ruby interpreter
+function when you're not in an interpreter thread can hork ruby, trigger a
+[BUG], corrupt the stack, kill your dog, knock up your daughter, etc. etc.
+
+NOTE: the above is only true when you're running in THREADED mode, in 
+single-threaded, everything is called on an interpreter thread.
+
 
 slyphon@gmail.com
 
 */
 
 #include "ruby.h"
-#include "zookeeper_lib.h"
+#include "event_lib.h"
 #include "c-client-src/zookeeper.h"
 #include <errno.h>
 #include <stdio.h>
@@ -28,13 +29,15 @@ slyphon@gmail.com
 #include <pthread.h>
 #include <unistd.h>
 #include <inttypes.h>
+#include "common.h"
 #include "dbg.h"
 
 #define GET_SYM(str) ID2SYM(rb_intern(str))
 
 int ZKRBDebugging;
 
-// XXX(slyphon): need to check these for error, but what to do if they fail?
+#ifdef THREADED
+
 pthread_mutex_t zkrb_q_mutex = PTHREAD_MUTEX_INITIALIZER;
 
 inline static int global_mutex_lock() {
@@ -66,6 +69,42 @@ void atfork_child() {
 // and children
 /*pthread_atfork(atfork_prepare, atfork_parent, atfork_child);*/
 
+// delegates to the system malloc (we can't use xmalloc in the threaded case,
+// as we can't touch the interpreter)
+
+inline static void* zk_malloc(size_t size) {
+  return malloc(size);
+}
+
+inline static void zk_free(void *ptr) {
+  free(ptr);
+}
+
+#else
+
+inline static int global_mutex_lock() {
+  return 0;
+}
+
+inline static int global_mutex_unlock() {
+  return 0;
+}
+
+// we can use the ruby xmalloc/xfree that will raise errors
+// in the case of a failure to allocate memory, and can cycle
+// the garbage collector in some cases.
+
+inline static void* zk_malloc(size_t size) {
+  return xmalloc(size);
+}
+
+inline static void zk_free(void *ptr) {
+  xfree(ptr);
+}
+
+
+#endif /* THREADED */
+
 
 void zkrb_enqueue(zkrb_queue_t *q, zkrb_event_t *elt) {
   if (q == NULL) {
@@ -81,7 +120,7 @@ void zkrb_enqueue(zkrb_queue_t *q, zkrb_event_t *elt) {
   global_mutex_lock();
 
   q->tail->event = elt;
-  q->tail->next = (zkrb_event_ll_t *) malloc(sizeof(zkrb_event_ll_t));
+  q->tail->next = (zkrb_event_ll_t *)zk_malloc(sizeof(zkrb_event_ll_t));
   q->tail = q->tail->next;
   q->tail->event = NULL;
   q->tail->next = NULL;
@@ -134,7 +173,7 @@ zkrb_event_t* zkrb_dequeue(zkrb_queue_t *q, int need_lock) {
   if (need_lock)
     global_mutex_unlock();
 
-  free(old_root);
+  zk_free(old_root);
   return rv;
 }
 
@@ -150,7 +189,7 @@ void zkrb_signal(zkrb_queue_t *q) {
 }
 
 zkrb_event_ll_t *zkrb_event_ll_t_alloc(void) {
-  zkrb_event_ll_t *rv = malloc(sizeof(zkrb_event_ll_t));
+  zkrb_event_ll_t *rv = zk_malloc(sizeof(zkrb_event_ll_t));
 
   if (!rv) return NULL;
 
@@ -166,7 +205,7 @@ zkrb_queue_t *zkrb_queue_alloc(void) {
  
   check(pipe(pfd) == 0, "creating the signal pipe failed");
 
-  rq = malloc(sizeof(zkrb_queue_t));
+  rq = zk_malloc(sizeof(zkrb_queue_t));
   check_mem(rq);
 
   rq->orig_pid = getpid();
@@ -181,7 +220,7 @@ zkrb_queue_t *zkrb_queue_alloc(void) {
   return rq;
 
 error:
-  free(rq);
+  zk_free(rq);
   return NULL;
 }
 
@@ -193,15 +232,15 @@ void zkrb_queue_free(zkrb_queue_t *queue) {
     zkrb_event_free(elt);
   }
 
-  free(queue->head);
+  zk_free(queue->head);
   close(queue->pipe_read);
   close(queue->pipe_write);
 
-  free(queue);
+  zk_free(queue);
 }
 
 zkrb_event_t *zkrb_event_alloc(void) {
-  zkrb_event_t *rv = malloc(sizeof(zkrb_event_t));
+  zkrb_event_t *rv = zk_malloc(sizeof(zkrb_event_t));
   return rv;
 }
 
@@ -209,58 +248,58 @@ void zkrb_event_free(zkrb_event_t *event) {
   switch (event->type) {
     case ZKRB_DATA: {
       struct zkrb_data_completion *data_ctx = event->completion.data_completion;
-      free(data_ctx->data);
-      free(data_ctx->stat);
-      free(data_ctx);
+      zk_free(data_ctx->data);
+      zk_free(data_ctx->stat);
+      zk_free(data_ctx);
       break;
     }
     case ZKRB_STAT: {
       struct zkrb_stat_completion *stat_ctx = event->completion.stat_completion;
-      free(stat_ctx->stat);
-      free(stat_ctx);
+      zk_free(stat_ctx->stat);
+      zk_free(stat_ctx);
       break;
     }
     case ZKRB_STRING: {
       struct zkrb_string_completion *string_ctx = event->completion.string_completion;
-      free(string_ctx->value);
-      free(string_ctx);
+      zk_free(string_ctx->value);
+      zk_free(string_ctx);
       break;
     }
     case ZKRB_STRINGS: {
       struct zkrb_strings_completion *strings_ctx = event->completion.strings_completion;
       int k;
       if (strings_ctx->values)
-        for (k = 0; k < strings_ctx->values->count; ++k) free(strings_ctx->values->data[k]);
-      free(strings_ctx->values);
-      free(strings_ctx);
+        for (k = 0; k < strings_ctx->values->count; ++k) zk_free(strings_ctx->values->data[k]);
+      zk_free(strings_ctx->values);
+      zk_free(strings_ctx);
       break;
     }
     case ZKRB_STRINGS_STAT: {
       struct zkrb_strings_stat_completion *strings_stat_ctx = event->completion.strings_stat_completion;
       int k;
       if (strings_stat_ctx->values)
-        for (k = 0; k < strings_stat_ctx->values->count; ++k) free(strings_stat_ctx->values->data[k]);
-      free(strings_stat_ctx->values);
+        for (k = 0; k < strings_stat_ctx->values->count; ++k) zk_free(strings_stat_ctx->values->data[k]);
+      zk_free(strings_stat_ctx->values);
 
       if (strings_stat_ctx->stat)
-	free(strings_stat_ctx->stat);
-      free(strings_stat_ctx);
+	zk_free(strings_stat_ctx->stat);
+      zk_free(strings_stat_ctx);
       break;
     }
     case ZKRB_ACL: {
       struct zkrb_acl_completion *acl_ctx = event->completion.acl_completion;
       if (acl_ctx->acl) {
         deallocate_ACL_vector(acl_ctx->acl);
-        free(acl_ctx->acl);
+        zk_free(acl_ctx->acl);
       }
-      free(acl_ctx->stat);
-      free(acl_ctx);
+      zk_free(acl_ctx->stat);
+      zk_free(acl_ctx);
       break;
     }
     case ZKRB_WATCHER: {
       struct zkrb_watcher_completion *watcher_ctx = event->completion.watcher_completion;
-      free(watcher_ctx->path);
-      free(watcher_ctx);
+      zk_free(watcher_ctx->path);
+      zk_free(watcher_ctx);
       break;
     }
     case ZKRB_VOID: {
@@ -270,7 +309,7 @@ void zkrb_event_free(zkrb_event_t *event) {
       log_err("unrecognized event in event_free!");
   }
 
-  free(event);
+  zk_free(event);
 }
 
 /* this is called only from a method_get_latest_event, so the hash is
@@ -289,6 +328,7 @@ VALUE zkrb_event_to_ruby(zkrb_event_t *event) {
 
   switch (event->type) {
     case ZKRB_DATA: {
+      zkrb_debug("zkrb_event_to_ruby ZKRB_DATA\n");
       struct zkrb_data_completion *data_ctx = event->completion.data_completion;
       if (ZKRBDebugging) zkrb_print_stat(data_ctx->stat);
       rb_hash_aset(hash, GET_SYM("data"), data_ctx->data ? rb_str_new(data_ctx->data, data_ctx->data_len) : Qnil);
@@ -329,7 +369,6 @@ VALUE zkrb_event_to_ruby(zkrb_event_t *event) {
     }
     case ZKRB_WATCHER: {
       zkrb_debug("zkrb_event_to_ruby ZKRB_WATCHER\n");
-      struct zkrb_acl_completion *acl_ctx = event->completion.acl_completion;
       struct zkrb_watcher_completion *watcher_ctx = event->completion.watcher_completion;
       rb_hash_aset(hash, GET_SYM("type"), INT2FIX(watcher_ctx->type));
       rb_hash_aset(hash, GET_SYM("state"), INT2FIX(watcher_ctx->state));
@@ -365,7 +404,7 @@ void zkrb_print_stat(const struct Stat *s) {
 }
 
 zkrb_calling_context *zkrb_calling_context_alloc(int64_t req_id, zkrb_queue_t *queue) {
-  zkrb_calling_context *ctx = malloc(sizeof(zkrb_calling_context));
+  zkrb_calling_context *ctx = zk_malloc(sizeof(zkrb_calling_context));
   if (!ctx) return NULL;
 
   ctx->req_id = req_id;
@@ -374,6 +413,10 @@ zkrb_calling_context *zkrb_calling_context_alloc(int64_t req_id, zkrb_queue_t *q
   return ctx;
 }
 
+void zkrb_calling_context_free(zkrb_calling_context *ctx) {
+  zk_free(ctx);
+}
+
 void zkrb_print_calling_context(zkrb_calling_context *ctx) {
   fprintf(stderr, "calling context (%p){\n", ctx);
   fprintf(stderr, "\treq_id = %"PRId64"\n", ctx->req_id);
@@ -394,7 +437,7 @@ void zkrb_print_calling_context(zkrb_calling_context *ctx) {
   zkrb_event_t *eptr = zkrb_event_alloc();                          \
   eptr->req_id = ctx->req_id;                                       \
   zkrb_queue_t *qptr = ctx->queue;                                  \
-  if (eptr->req_id != ZKRB_GLOBAL_REQ) free(ctx)
+  if (eptr->req_id != ZKRB_GLOBAL_REQ) zk_free(ctx)
 
 void zkrb_state_callback(
     zhandle_t *zh, int type, int state, const char *path, void *calling_ctx) {
@@ -404,7 +447,7 @@ void zkrb_state_callback(
       type, state, (void *) path, path ? path : "NULL");
 
   /* save callback context */
-  struct zkrb_watcher_completion *wc = malloc(sizeof(struct zkrb_watcher_completion));
+  struct zkrb_watcher_completion *wc = zk_malloc(sizeof(struct zkrb_watcher_completion));
   wc->type  = type;
   wc->state = state;
   wc->path  = strdup(path);
@@ -416,7 +459,7 @@ void zkrb_state_callback(
   event->req_id = ctx->req_id;
   zkrb_queue_t *queue = ctx->queue;
   if (type != ZOO_SESSION_EVENT) {
-    free(ctx);
+    zk_free(ctx);
     ctx = NULL;
   }
 
@@ -426,8 +469,6 @@ void zkrb_state_callback(
   zkrb_enqueue(queue, event);
 }
 
-
-
 void zkrb_data_callback(
     int rc, const char *value, int value_len, const struct Stat *stat, const void *calling_ctx) {
 
@@ -436,18 +477,18 @@ void zkrb_data_callback(
                 rc, zerror(rc), value ? value : "NULL", value_len);
 
   /* copy data completion */
-  struct zkrb_data_completion *dc = malloc(sizeof(struct zkrb_data_completion));
+  struct zkrb_data_completion *dc = zk_malloc(sizeof(struct zkrb_data_completion));
   dc->data = NULL;
   dc->stat = NULL;
   dc->data_len = 0;
 
   if (value != NULL) {
-    dc->data = malloc(value_len);
+    dc->data = zk_malloc(value_len);  // xmalloc may raise an exception, which means the above completion will leak
     dc->data_len = value_len;
     memcpy(dc->data, value, value_len);
   }
 
-  if (stat != NULL) { dc->stat  = malloc(sizeof(struct Stat)); memcpy(dc->stat, stat, sizeof(struct Stat)); }
+  if (stat != NULL) { dc->stat = zk_malloc(sizeof(struct Stat)); memcpy(dc->stat, stat, sizeof(struct Stat)); }
 
   ZKH_SETUP_EVENT(queue, event);
   event->rc = rc;
@@ -462,9 +503,9 @@ void zkrb_stat_callback(
   zkrb_debug("ZOOKEEPER_C_STAT WATCHER "
                     "rc = %d (%s)\n", rc, zerror(rc));
 
-  struct zkrb_stat_completion *sc = malloc(sizeof(struct zkrb_stat_completion));
+  struct zkrb_stat_completion *sc = zk_malloc(sizeof(struct zkrb_stat_completion));
   sc->stat = NULL;
-  if (stat != NULL) { sc->stat = malloc(sizeof(struct Stat)); memcpy(sc->stat, stat, sizeof(struct Stat)); }
+  if (stat != NULL) { sc->stat = zk_malloc(sizeof(struct Stat)); memcpy(sc->stat, stat, sizeof(struct Stat)); }
 
   ZKH_SETUP_EVENT(queue, event);
   event->rc = rc;
@@ -480,7 +521,7 @@ void zkrb_string_callback(
   zkrb_debug("ZOOKEEPER_C_STRING WATCHER "
                     "rc = %d (%s)\n", rc, zerror(rc));
 
-  struct zkrb_string_completion *sc = malloc(sizeof(struct zkrb_string_completion));
+  struct zkrb_string_completion *sc = zk_malloc(sizeof(struct zkrb_string_completion));
   sc->value = NULL;
   if (string)
     sc->value = strdup(string);
@@ -499,7 +540,7 @@ void zkrb_strings_callback(
                     "rc = %d (%s), calling_ctx = %p\n", rc, zerror(rc), calling_ctx);
 
   /* copy string vector */
-  struct zkrb_strings_completion *sc = malloc(sizeof(struct zkrb_strings_completion));
+  struct zkrb_strings_completion *sc = zk_malloc(sizeof(struct zkrb_strings_completion));
   sc->values = (strings != NULL) ? zkrb_clone_string_vector(strings) : NULL;
 
   ZKH_SETUP_EVENT(queue, event);
@@ -515,9 +556,9 @@ void zkrb_strings_stat_callback(
   zkrb_debug("ZOOKEEPER_C_STRINGS_STAT WATCHER "
                     "rc = %d (%s), calling_ctx = %p\n", rc, zerror(rc), calling_ctx);
 
-  struct zkrb_strings_stat_completion *sc = malloc(sizeof(struct zkrb_strings_stat_completion));
+  struct zkrb_strings_stat_completion *sc = zk_malloc(sizeof(struct zkrb_strings_stat_completion));
   sc->stat = NULL;
-  if (stat != NULL) { sc->stat = malloc(sizeof(struct Stat)); memcpy(sc->stat, stat, sizeof(struct Stat)); }
+  if (stat != NULL) { sc->stat = zk_malloc(sizeof(struct Stat)); memcpy(sc->stat, stat, sizeof(struct Stat)); }
 
   sc->values = (strings != NULL) ? zkrb_clone_string_vector(strings) : NULL;
 
@@ -545,11 +586,11 @@ void zkrb_acl_callback(
     int rc, struct ACL_vector *acls, struct Stat *stat, const void *calling_ctx) {
   zkrb_debug("ZOOKEEPER_C_ACL WATCHER rc = %d (%s)\n", rc, zerror(rc));
 
-  struct zkrb_acl_completion *ac = malloc(sizeof(struct zkrb_acl_completion));
+  struct zkrb_acl_completion *ac = zk_malloc(sizeof(struct zkrb_acl_completion));
   ac->acl = NULL;
   ac->stat = NULL;
   if (acls != NULL) { ac->acl  = zkrb_clone_acl_vector(acls); }
-  if (stat != NULL) { ac->stat = malloc(sizeof(struct Stat)); memcpy(ac->stat, stat, sizeof(struct Stat)); }
+  if (stat != NULL) { ac->stat = zk_malloc(sizeof(struct Stat)); memcpy(ac->stat, stat, sizeof(struct Stat)); }
 
   ZKH_SETUP_EVENT(queue, event);
   event->rc = rc;
@@ -574,12 +615,12 @@ VALUE zkrb_acl_to_ruby(struct ACL *acl) {
   return hash;
 }
 
-#warning [wickman] TODO test zkrb_ruby_to_aclvector
-#warning [slyphon] TODO size checking on acl_ary (cast to int)
+// [wickman] TODO test zkrb_ruby_to_aclvector
+// [slyphon] TODO size checking on acl_ary (cast to int)
 struct ACL_vector * zkrb_ruby_to_aclvector(VALUE acl_ary) {
   Check_Type(acl_ary, T_ARRAY);
 
-  struct ACL_vector *v = malloc(sizeof(struct ACL_vector));
+  struct ACL_vector *v = zk_malloc(sizeof(struct ACL_vector));
   allocate_ACL_vector(v, (int)RARRAY_LEN(acl_ary));
 
   int k;
@@ -591,7 +632,7 @@ struct ACL_vector * zkrb_ruby_to_aclvector(VALUE acl_ary) {
   return v;
 }
 
-#warning [wickman] TODO test zkrb_ruby_to_aclvector
+// [wickman] TODO test zkrb_ruby_to_aclvector
 struct ACL zkrb_ruby_to_acl(VALUE rubyacl) {
   struct ACL acl;
 
@@ -603,7 +644,7 @@ struct ACL zkrb_ruby_to_acl(VALUE rubyacl) {
   return acl;
 }
 
-#warning [wickman] TODO zkrb_ruby_to_id error checking? test
+// [wickman] TODO zkrb_ruby_to_id error checking? test
 struct Id zkrb_ruby_to_id(VALUE rubyid) {
   struct Id id;
 
@@ -611,7 +652,7 @@ struct Id zkrb_ruby_to_id(VALUE rubyid) {
   VALUE ident  = rb_iv_get(rubyid, "@id");
 
   if (scheme != Qnil) {
-    id.scheme = malloc(RSTRING_LEN(scheme) + 1);
+    id.scheme = zk_malloc(RSTRING_LEN(scheme) + 1);
     strncpy(id.scheme, RSTRING_PTR(scheme), RSTRING_LEN(scheme));
     id.scheme[RSTRING_LEN(scheme)] = '\0';
   } else {
@@ -619,7 +660,7 @@ struct Id zkrb_ruby_to_id(VALUE rubyid) {
   }
 
   if (ident != Qnil) {
-    id.id = malloc(RSTRING_LEN(ident) + 1);
+    id.id = zk_malloc(RSTRING_LEN(ident) + 1);
     strncpy(id.id, RSTRING_PTR(ident), RSTRING_LEN(ident));
     id.id[RSTRING_LEN(ident)] = '\0';
   } else {
@@ -678,9 +719,9 @@ VALUE zkrb_stat_to_rhash(const struct Stat *stat) {
   return ary;
 }
 
-#warning [wickman] TODO test zkrb_clone_acl_vector
+// [wickman] TODO test zkrb_clone_acl_vector
 struct ACL_vector * zkrb_clone_acl_vector(struct ACL_vector * src) {
-  struct ACL_vector * dst = malloc(sizeof(struct ACL_vector));
+  struct ACL_vector * dst = zk_malloc(sizeof(struct ACL_vector));
   allocate_ACL_vector(dst, src->count);
   int k;
   for (k = 0; k < src->count; ++k) {
@@ -692,9 +733,9 @@ struct ACL_vector * zkrb_clone_acl_vector(struct ACL_vector * src) {
   return dst;
 }
 
-#warning [wickman] TODO test zkrb_clone_string_vector
+// [wickman] TODO test zkrb_clone_string_vector
 struct String_vector * zkrb_clone_string_vector(const struct String_vector * src) {
-  struct String_vector * dst = malloc(sizeof(struct String_vector));
+  struct String_vector * dst = zk_malloc(sizeof(struct String_vector));
   allocate_String_vector(dst, src->count);
   int k;
   for (k = 0; k < src->count; ++k) {

data/ext/{zookeeper_lib.h → event_lib.h}

@@ -1,5 +1,5 @@
-#ifndef ZOOKEEPER_LIB_H
-#define ZOOKEEPER_LIB_H
+#ifndef ZKRB_EVENT_LIB_H
+#define ZKRB_EVENT_LIB_H
 
 #include "ruby.h"
 #include "c-client-src/zookeeper.h"
@@ -128,6 +128,7 @@ typedef struct {
 
 void zkrb_print_calling_context(zkrb_calling_context *ctx);
 zkrb_calling_context *zkrb_calling_context_alloc(int64_t req_id, zkrb_queue_t *queue);
+void zkrb_calling_context_free(zkrb_calling_context *ctx);
 
 /*
   default process completions that get queued into the ruby client event queue
@@ -171,4 +172,4 @@ struct String_vector * zkrb_clone_string_vector(const struct String_vector * src
 struct ACL             zkrb_ruby_to_acl(VALUE rubyacl);
 struct Id              zkrb_ruby_to_id(VALUE rubyid);
 
-#endif  /* ZOOKEEPER_LIB_H */
+#endif  /* ZKRB_EVENT_LIB_H */

data/ext/extconf.rb

@@ -1,10 +1,11 @@
-
 require 'mkmf'
 require 'rbconfig'
+require 'fileutils'
 
 HERE = File.expand_path(File.dirname(__FILE__))
 BUNDLE = Dir.glob("zkc-*.tar.gz").first
-BUNDLE_PATH = "c"
+
+BUNDLE_PATH = File.join(HERE, 'c')
 
 $EXTRA_CONF = ''
 
@@ -24,6 +25,7 @@ if RUBY_VERSION == '1.8.7'
 end
 
 ZK_DEBUG = (ENV['DEBUG'] or ARGV.any? { |arg| arg == '--debug' })
+ZK_DEV = ENV['ZK_DEV']
 DEBUG_CFLAGS = " -O0 -ggdb3 -DHAVE_DEBUG"
 
 if ZK_DEBUG
@@ -61,7 +63,6 @@ Dir.chdir(HERE) do
 
     Dir.chdir(BUNDLE_PATH) do        
       configure = "./configure --prefix=#{HERE} --with-pic --without-cppunit --disable-dependency-tracking #{$EXTRA_CONF} 2>&1"
-      
       configure = "env CFLAGS='#{DEBUG_CFLAGS}' #{configure}" if ZK_DEBUG
 
       safe_sh(configure)
@@ -69,17 +70,20 @@ Dir.chdir(HERE) do
       safe_sh("make install 2>&1")
     end
 
-    system("rm -rf #{BUNDLE_PATH}") unless ENV['DEBUG'] or ENV['DEV']
+    system("rm -rf #{BUNDLE_PATH}") unless ZK_DEBUG or ZK_DEV
   end
 end
 
 # Absolutely prevent the linker from picking up any other zookeeper_mt
 Dir.chdir("#{HERE}/lib") do
-  system("cp -f libzookeeper_mt.a libzookeeper_mt_gem.a") 
-  system("cp -f libzookeeper_mt.la libzookeeper_mt_gem.la") 
+  %w[st mt].each do |stmt|
+    %w[a la].each do |ext|
+      system("cp -f libzookeeper_#{stmt}.#{ext} libzookeeper_#{stmt}_gem.#{ext}") 
+    end
+  end
 end
-$LIBS << " -lzookeeper_mt_gem"
-
+$LIBS << " -lzookeeper_st_gem"
 
+$CFLAGS << ' -Wall' if ZK_DEV
 create_makefile 'zookeeper_c'
 

data/ext/generate_gvl_code.rb

@@ -59,6 +59,8 @@ REGEXP = /^ZOOAPI int (zoo_[^(]+)\(([^)]+)\);$/m
 require 'forwardable'
 require 'stringio'
 
+THIS_DIR = File.expand_path('..', __FILE__)
+
 ZKRB_WRAPPER_H_PATH = File.expand_path('../zkrb_wrapper.h', __FILE__)
 ZKRB_WRAPPER_C_PATH = File.expand_path('../zkrb_wrapper.c', __FILE__)
 
@@ -292,8 +294,13 @@ end
 
 def main
   help! if ARGV.empty?
-  
-  text = File.read('c/include/zookeeper.h')
+  opts = []
+
+  zookeeper_h_path = Dir[File.join(THIS_DIR, "**/zookeeper.h")].first
+
+  raise "Could not locate zookeeper.h!" unless zookeeper_h_path
+
+  text = File.read(zookeeper_h_path)
   code = GeneratedCode.from_zookeeper_h(text)
 
   cmd = ARGV.first