Cleanup async mozStorage callback management.

Also avoid raw pointers in mozStorageAsyncStatementExecution.cpp.

2 files changed, 88 insertions, 140 deletions

diff --git a/storage/mozStorageAsyncStatementExecution.cpp b/storage/mozStorageAsyncStatementExecution.cpp
index add32131a3..e1d344aca6 100644
--- a/storage/mozStorageAsyncStatementExecution.cpp
+++ b/storage/mozStorageAsyncStatementExecution.cpp
@@ -40,122 +40,6 @@ namespace storage {
 #define MAX_ROWS_PER_RESULT 15
 
 ////////////////////////////////////////////////////////////////////////////////
-//// Local Classes
-
-namespace {
-
-typedef AsyncExecuteStatements::ExecutionState ExecutionState;
-typedef AsyncExecuteStatements::StatementDataArray StatementDataArray;
-
-/**
- * Notifies a callback with a result set.
- */
-class CallbackResultNotifier : public Runnable
-{
-public:
-  CallbackResultNotifier(mozIStorageStatementCallback *aCallback,
-                         mozIStorageResultSet *aResults,
-                         AsyncExecuteStatements *aEventStatus) :
-      mCallback(aCallback)
-    , mResults(aResults)
-    , mEventStatus(aEventStatus)
-  {
-  }
-
-  NS_IMETHOD Run() override
-  {
-    NS_ASSERTION(mCallback, "Trying to notify about results without a callback!");
-
-    if (mEventStatus->shouldNotify()) {
-      // Hold a strong reference to the callback while notifying it, so that if
-      // it spins the event loop, the callback won't be released and freed out
-      // from under us.
-      nsCOMPtr<mozIStorageStatementCallback> callback = mCallback;
-
-      (void)callback->HandleResult(mResults);
-    }
-
-    return NS_OK;
-  }
-
-private:
-  mozIStorageStatementCallback *mCallback;
-  nsCOMPtr<mozIStorageResultSet> mResults;
-  RefPtr<AsyncExecuteStatements> mEventStatus;
-};
-
-/**
- * Notifies the calling thread that an error has occurred.
- */
-class ErrorNotifier : public Runnable
-{
-public:
-  ErrorNotifier(mozIStorageStatementCallback *aCallback,
-                mozIStorageError *aErrorObj,
-                AsyncExecuteStatements *aEventStatus) :
-      mCallback(aCallback)
-    , mErrorObj(aErrorObj)
-    , mEventStatus(aEventStatus)
-  {
-  }
-
-  NS_IMETHOD Run() override
-  {
-    if (mEventStatus->shouldNotify() && mCallback) {
-      // Hold a strong reference to the callback while notifying it, so that if
-      // it spins the event loop, the callback won't be released and freed out
-      // from under us.
-      nsCOMPtr<mozIStorageStatementCallback> callback = mCallback;
-
-      (void)callback->HandleError(mErrorObj);
-    }
-
-    return NS_OK;
-  }
-
-private:
-  mozIStorageStatementCallback *mCallback;
-  nsCOMPtr<mozIStorageError> mErrorObj;
-  RefPtr<AsyncExecuteStatements> mEventStatus;
-};
-
-/**
- * Notifies the calling thread that the statement has finished executing.  Takes
- * ownership of the StatementData so it is released on the proper thread.
- */
-class CompletionNotifier : public Runnable
-{
-public:
-  /**
-   * This takes ownership of the callback and the StatementData.  They are
-   * released on the thread this is dispatched to (which should always be the
-   * calling thread).
-   */
-  CompletionNotifier(mozIStorageStatementCallback *aCallback,
-                     ExecutionState aReason)
-    : mCallback(aCallback)
-    , mReason(aReason)
-  {
-  }
-
-  NS_IMETHOD Run() override
-  {
-    if (mCallback) {
-      (void)mCallback->HandleCompletion(mReason);
-      NS_RELEASE(mCallback);
-    }
-
-    return NS_OK;
-  }
-
-private:
-  mozIStorageStatementCallback *mCallback;
-  ExecutionState mReason;
-};
-
-} // namespace
-
-////////////////////////////////////////////////////////////////////////////////
 //// AsyncExecuteStatements
 
 /* static */
@@ -208,16 +92,19 @@ AsyncExecuteStatements::AsyncExecuteStatements(StatementDataArray &aStatements,
 , mCancelRequested(false)
 , mMutex(aConnection->sharedAsyncExecutionMutex)
 , mDBMutex(aConnection->sharedDBMutex)
-  , mRequestStartDate(TimeStamp::Now())
+, mRequestStartDate(TimeStamp::Now())
 {
   (void)mStatements.SwapElements(aStatements);
   NS_ASSERTION(mStatements.Length(), "We weren't given any statements!");
-  NS_IF_ADDREF(mCallback);
 }
 
 AsyncExecuteStatements::~AsyncExecuteStatements()
 {
+  MOZ_ASSERT(!mCallback, "Never called the Completion callback!");
   MOZ_ASSERT(!mHasTransaction, "There should be no transaction at this point");
+  if (mCallback) {
+    NS_ProxyRelease(mCallingThread, mCallback.forget());
+  }
 }
 
 bool
@@ -255,7 +142,7 @@ AsyncExecuteStatements::bindExecuteAndProcessStatement(StatementData &aData,
   BindingParamsArray::iterator end = paramsArray->end();
   while (itr != end && continueProcessing) {
     // Bind the data to our statement.
-    nsCOMPtr<IStorageBindingParamsInternal> bindingInternal = 
+    nsCOMPtr<IStorageBindingParamsInternal> bindingInternal =
       do_QueryInterface(*itr);
     nsCOMPtr<mozIStorageError> error = bindingInternal->bind(aStatement);
     if (error) {
@@ -462,16 +349,30 @@ AsyncExecuteStatements::notifyComplete()
     mHasTransaction = false;
   }
 
-  // Always generate a completion notification; it is what guarantees that our
-  // destruction does not happen here on the async thread.
-  RefPtr<CompletionNotifier> completionEvent =
-    new CompletionNotifier(mCallback, mState);
-
-  // We no longer own mCallback (the CompletionNotifier takes ownership).
-  mCallback = nullptr;
+  // This will take ownership of mCallback and make sure its destruction will
+  // happen on the owner thread.
+  Unused << mCallingThread->Dispatch(
+    NewRunnableMethod(this, &AsyncExecuteStatements::notifyCompleteOnCallingThread),
+    NS_DISPATCH_NORMAL);
 
-  (void)mCallingThread->Dispatch(completionEvent, NS_DISPATCH_NORMAL);
+  return NS_OK;
+}
 
+nsresult
+AsyncExecuteStatements::notifyCompleteOnCallingThread() {
+#ifdef DEBUG
+  bool onCallingThread = false;
+  (void)mCallingThread->IsOnCurrentThread(&onCallingThread);
+  MOZ_ASSERT(onCallingThread);
+#endif
+  // Take ownership of mCallback and responsibility for freeing it when we
+  // release it.  Any notifyResultsOnCallingThread and notifyErrorOnCallingThread
+  // calls on the stack spinning the event loop have guaranteed their safety by
+  // creating their own strong reference before invoking the callback.
+  nsCOMPtr<mozIStorageStatementCallback> callback = mCallback.forget();
+  if (callback) {
+    Unused << callback->HandleCompletion(mState);
+  }
   return NS_OK;
 }
 
@@ -500,27 +401,63 @@ AsyncExecuteStatements::notifyError(mozIStorageError *aError)
   if (!mCallback)
     return NS_OK;
 
-  RefPtr<ErrorNotifier> notifier =
-    new ErrorNotifier(mCallback, aError, this);
-  NS_ENSURE_TRUE(notifier, NS_ERROR_OUT_OF_MEMORY);
+  Unused << mCallingThread->Dispatch(
+    NewRunnableMethod<nsCOMPtr<mozIStorageError>>(this, &AsyncExecuteStatements::notifyErrorOnCallingThread, aError),
+    NS_DISPATCH_NORMAL);
 
-  return mCallingThread->Dispatch(notifier, NS_DISPATCH_NORMAL);
+  return NS_OK;
+}
+
+nsresult
+AsyncExecuteStatements::notifyErrorOnCallingThread(mozIStorageError *aError) {
+#ifdef DEBUG
+  bool onCallingThread = false;
+  (void)mCallingThread->IsOnCurrentThread(&onCallingThread);
+  MOZ_ASSERT(onCallingThread);
+#endif
+  // Acquire our own strong reference so that if the callback spins a nested
+  // event loop and notifyCompleteOnCallingThread is executed, forgetting
+  // mCallback, we still have a valid/strong reference that won't be freed until
+  // we exit.
+  nsCOMPtr<mozIStorageStatementCallback> callback = mCallback;
+  if (shouldNotify() && callback) {
+    Unused << callback->HandleError(aError);
+  }
+  return NS_OK;
 }
 
 nsresult
 AsyncExecuteStatements::notifyResults()
 {
   mMutex.AssertNotCurrentThreadOwns();
-  NS_ASSERTION(mCallback, "notifyResults called without a callback!");
+  MOZ_ASSERT(mCallback, "notifyResults called without a callback!");
 
-  RefPtr<CallbackResultNotifier> notifier =
-    new CallbackResultNotifier(mCallback, mResultSet, this);
-  NS_ENSURE_TRUE(notifier, NS_ERROR_OUT_OF_MEMORY);
+  // This takes ownership of mResultSet, a new one will be generated in
+  // buildAndNotifyResults() when further results will arrive.
+  Unused << mCallingThread->Dispatch(
+    NewRunnableMethod<RefPtr<ResultSet>>(this, &AsyncExecuteStatements::notifyResultsOnCallingThread, mResultSet.forget()),
+    NS_DISPATCH_NORMAL);
 
-  nsresult rv = mCallingThread->Dispatch(notifier, NS_DISPATCH_NORMAL);
-  if (NS_SUCCEEDED(rv))
-    mResultSet = nullptr; // we no longer own it on success
-  return rv;
+  return NS_OK;
+}
+
+nsresult
+AsyncExecuteStatements::notifyResultsOnCallingThread(ResultSet *aResultSet)
+{
+#ifdef DEBUG
+  bool onCallingThread = false;
+  (void)mCallingThread->IsOnCurrentThread(&onCallingThread);
+  MOZ_ASSERT(onCallingThread);
+#endif
+  // Acquire our own strong reference so that if the callback spins a nested
+  // event loop and notifyCompleteOnCallingThread is executed, forgetting
+  // mCallback, we still have a valid/strong reference that won't be freed until
+  // we exit.
+  nsCOMPtr<mozIStorageStatementCallback> callback = mCallback;
+  if (shouldNotify() && callback) {
+    Unused << callback->HandleResult(aResultSet);
+  }
+  return NS_OK;
 }
 
 NS_IMPL_ISUPPORTS(
diff --git a/storage/mozStorageAsyncStatementExecution.h b/storage/mozStorageAsyncStatementExecution.h
index c8493fd771..14ea49c2d9 100644
--- a/storage/mozStorageAsyncStatementExecution.h
+++ b/storage/mozStorageAsyncStatementExecution.h
@@ -82,6 +82,14 @@ public:
    */
   bool shouldNotify();
 
+  /**
+   * Used by notifyComplete(), notifyError() and notifyResults() to notify on
+   * the calling thread.
+   */
+  nsresult notifyCompleteOnCallingThread();
+  nsresult notifyErrorOnCallingThread(mozIStorageError *aError);
+  nsresult notifyResultsOnCallingThread(ResultSet *aResultSet);
+
 private:
   AsyncExecuteStatements(StatementDataArray &aStatements,
                          Connection *aConnection,
@@ -186,7 +194,10 @@ private:
   RefPtr<Connection> mConnection;
   sqlite3 *mNativeConnection;
   bool mHasTransaction;
-  mozIStorageStatementCallback *mCallback;
+  // Note, this may not be a threadsafe object - never addref/release off
+  // the calling thread.  We take a reference when this is created, and
+  // release it in the CompletionNotifier::Run() call back to this thread.
+  nsCOMPtr<mozIStorageStatementCallback> mCallback;
   nsCOMPtr<nsIThread> mCallingThread;
   RefPtr<ResultSet> mResultSet;