fix race condition when baton was deleted to early

kkaefer · kkaefer · commit 25388bcdafd4 · 2011-11-07T22:30:42.000+01:00
diff --git a/src/statement.cc b/src/statement.cc
@@ -593,13 +593,19 @@ Handle<Value> Statement::Each(const Arguments& args) {
     }
     else {
         baton->completed = Persistent<Function>::New(completed);
-        baton->async = new Async(stmt, baton, AsyncEach);
         stmt->Schedule(EIO_BeginEach, baton);
         return args.This();
     }
 }
 
 void Statement::EIO_BeginEach(Baton* baton) {
+    // Only create the Async object when we're actually going into
+    // the event loop. This prevents dangling events.
+    EachBaton* each_baton = static_cast<EachBaton*>(baton);
+    each_baton->async = new Async(each_baton->stmt, AsyncEach);
+    each_baton->async->item_cb = Persistent<Function>::New(each_baton->callback);
+    each_baton->async->completed_cb = Persistent<Function>::New(each_baton->completed);
+
     STATEMENT_BEGIN(Each);
 }
 
@@ -649,6 +655,7 @@ void Statement::EIO_Each(eio_req *req) {
 
 void Statement::CloseCallback(uv_handle_t* handle) {
     assert(handle != NULL);
+    assert(handle->data != NULL);
     Async* async = static_cast<Async*>(handle->data);
     delete async;
     handle->data = NULL;
@@ -657,7 +664,6 @@ void Statement::CloseCallback(uv_handle_t* handle) {
 void Statement::AsyncEach(uv_async_t* handle, int status) {
     HandleScope scope;
     Async* async = static_cast<Async*>(handle->data);
-    EachBaton* baton = async->baton;
 
     while (true) {
         // Get the contents out of the data cache for us to process in the JS callback.
@@ -670,7 +676,7 @@ void Statement::AsyncEach(uv_async_t* handle, int status) {
             break;
         }
 
-        if (!baton->callback.IsEmpty() && baton->callback->IsFunction()) {
+        if (!async->item_cb.IsEmpty() && async->item_cb->IsFunction()) {
             Local<Value> argv[2];
             argv[0] = Local<Value>::New(Null());
 
@@ -679,20 +685,20 @@ void Statement::AsyncEach(uv_async_t* handle, int status) {
             for (int i = 0; it < end; it++, i++) {
                 argv[1] = RowToJS(*it);
                 async->retrieved++;
-                TRY_CATCH_CALL(async->stmt->handle_, baton->callback, 2, argv);
+                TRY_CATCH_CALL(async->stmt->handle_, async->item_cb, 2, argv);
                 delete *it;
             }
         }
     }
 
     if (async->completed) {
-        if (!baton->completed.IsEmpty() &&
-                baton->completed->IsFunction()) {
+        if (!async->completed_cb.IsEmpty() &&
+                async->completed_cb->IsFunction()) {
             Local<Value> argv[] = {
                 Local<Value>::New(Null()),
                 Integer::New(async->retrieved)
             };
-            TRY_CATCH_CALL(async->stmt->handle_, baton->completed, 2, argv);
+            TRY_CATCH_CALL(async->stmt->handle_, async->completed_cb, 2, argv);
         }
         uv_close((uv_handle_t*)handle, CloseCallback);
     }
diff --git a/src/statement.h b/src/statement.h
@@ -88,7 +88,7 @@ class Statement : public ObjectWrap {
             ev_ref(EV_DEFAULT_UC);
             callback = Persistent<Function>::New(cb_);
         }
-        ~Baton() {
+        virtual ~Baton() {
             for (unsigned int i = 0; i < parameters.size(); i++) {
                 Values::Field* field = parameters[i];
                 DELETE_FIELD(field);
@@ -124,7 +124,7 @@ class Statement : public ObjectWrap {
         EachBaton(Statement* stmt_, Handle<Function> cb_) :
             Baton(stmt_, cb_) {}
         Persistent<Function> completed;
-        Async* async;
+        Async* async; // Isn't deleted when the baton is deleted.
     };
 
     struct PrepareBaton : Database::Baton {
@@ -155,14 +155,18 @@ class Statement : public ObjectWrap {
     struct Async {
         uv_async_t watcher;
         Statement* stmt;
-        EachBaton* baton;
         Rows data;
         pthread_mutex_t mutex;
         bool completed;
         int retrieved;
 
-        Async(Statement* st, EachBaton* eb, uv_async_cb async_cb) :
-                stmt(st), baton(eb), completed(false), retrieved(0) {
+        // Store the callbacks here because we don't have
+        // access to the baton in the async callback.
+        Persistent<Function> item_cb;
+        Persistent<Function> completed_cb;
+
+        Async(Statement* st, uv_async_cb async_cb) :
+                stmt(st), completed(false), retrieved(0) {
             watcher.data = this;
             pthread_mutex_init(&mutex, NULL);
             stmt->Ref();
@@ -171,6 +175,8 @@ class Statement : public ObjectWrap {
 
         ~Async() {
             stmt->Unref();
+            item_cb.Dispose();
+            completed_cb.Dispose();
             pthread_mutex_destroy(&mutex);
         }
     };
