Fix race condition on non-blocking queries

http://www.prjhub.com/#/issues/6217
orientechnologies · Feb 16, 2016 · a4aeaf8 · a4aeaf8
1 parent 6776f28
commit a4aeaf8
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 22 deletions.
diff --git a/core/src/main/java/com/orientechnologies/orient/core/sql/query/OSQLNonBlockingQuery.java b/core/src/main/java/com/orientechnologies/orient/core/sql/query/OSQLNonBlockingQuery.java
@@ -39,10 +39,9 @@
  * SQL asynchronous query. When executed the caller does not wait for the execution, rather the listener will be called for each
  * item found in the query. OSQLAsynchQuery has been built on top of this. NOTE: if you're working with remote databases don't
  * execute any remote call inside the callback function because the network channel is locked until the query command has finished.
- * 
- * @author Luca Garulli
- * 
+ *
  * @param <T>
+ * @author Luca Garulli
  * @see com.orientechnologies.orient.core.sql.query.OSQLSynchQuery
  */
 public class OSQLNonBlockingQuery<T extends Object> extends OSQLQuery<T> implements OCommandRequestAsynch {
@@ -243,18 +242,26 @@ public <RET> RET execute(final Object... iArgs) {
     final ONonBlockingQueryFuture future = new ONonBlockingQueryFuture();
 
     if (database instanceof ODatabaseDocumentTx) {
+      ODatabaseDocumentInternal currentThreadLocal = ODatabaseRecordThreadLocal.INSTANCE.getIfDefined();
+      final ODatabaseDocumentTx db = ((ODatabaseDocumentTx) database).copy();
+      if (currentThreadLocal != null) {
+        currentThreadLocal.activateOnCurrentThread();
+      } else {
+        ODatabaseRecordThreadLocal.INSTANCE.set(null);
+      }
+
       Thread t = new Thread(new Runnable() {
         @Override
         public void run() {
-          final ODatabaseDocumentTx db = ((ODatabaseDocumentTx) database).copy();
+          db.activateOnCurrentThread();
           try {
             OSQLNonBlockingQuery.super.execute(iArgs);
-          } catch(RuntimeException e){
-            if(getResultListener()!=null){
+          } catch (RuntimeException e) {
+            if (getResultListener() != null) {
               getResultListener().end();
             }
             throw e;
-          }finally {
+          } finally {
             if (db != null) {
               try {
                 db.close();

diff --git a/core/src/test/java/com/orientechnologies/orient/core/sql/ONonBlockingQueryTest.java b/core/src/test/java/com/orientechnologies/orient/core/sql/ONonBlockingQueryTest.java
@@ -12,10 +12,10 @@
 @Test
 public class ONonBlockingQueryTest {
 
-  static class MyResultListener implements OCommandResultListener{
+  static class MyResultListener implements OCommandResultListener {
 
-    public int numResults = 0;
-    public boolean finished = false;
+    public int     numResults = 0;
+    public boolean finished   = false;
 
     @Override public boolean result(Object iRecord) {
       numResults++;
@@ -36,11 +36,9 @@ public void testExceptionManagement() {
     db.activateOnCurrentThread();
     db.create();
 
+    db.getMetadata().getSchema().createClass("test");
+    MyResultListener listener = new MyResultListener();
     try {
-      db.getMetadata().getSchema().createClass("test");
-      MyResultListener listener = new MyResultListener();
-
-
       db.command(new OCommandSQL("insert into test set name = 'foo', surname = 'bar'")).execute();
 
       db.query(new OSQLNonBlockingQuery<Object>("select from test bla blu", listener));
@@ -54,16 +52,16 @@ public void testExceptionManagement() {
       listener = new MyResultListener();
       db.query(new OSQLNonBlockingQuery<Object>("select from test", listener));
 
-      try {
-        Thread.sleep(3000);
-      } catch (InterruptedException e) {
-        e.printStackTrace();
-      }
-      Assert.assertEquals(listener.numResults, 1);
-
     } finally {
-      db.drop();
+      db.close();
     }
+    try {
+      Thread.sleep(3000);
+    } catch (InterruptedException e) {
+      e.printStackTrace();
+    }
+    Assert.assertEquals(listener.numResults, 1);
+
   }
 
 }