IGNITE-28568 SQL Calcite: Fix sync tx key waiting for DML operations

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/ModifyNode.java

@@ -29,6 +29,8 @@
 import org.apache.calcite.rel.core.TableModify;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.internal.IgniteFutureTimeoutCheckedException;
+import org.apache.ignite.internal.IgniteInternalFuture;
 import org.apache.ignite.internal.cache.context.SessionContextImpl;
 import org.apache.ignite.internal.processors.cache.GridCacheContext;
 import org.apache.ignite.internal.processors.cache.IgniteInternalCache;
@@ -50,6 +52,12 @@
  *
  */
 public class ModifyNode<Row> extends AbstractNode<Row> implements SingleNode<Row>, Downstream<Row> {
+    /**
+     * Timeout to wait for async invoke operation to complete. In case this timeout exceeded, release thread and
+     * continue execution after invoke operation finished.
+     */
+    private static final long INVOKE_TIMEOUT = 100;
+
     /** */
     protected final CacheTableDescriptor desc;
 
@@ -77,6 +85,9 @@
     /** */
     private State state = State.UPDATING;
 
+    /** */
+    private IgniteInternalFuture<Map<Object, EntryProcessorResult<Long>>> invokeFut;
+
     /**
      * @param ctx Execution context.
      * @param desc Table descriptor.
@@ -119,22 +130,23 @@
 
         waiting--;
 
-        switch (op) {
-            case DELETE:
-            case UPDATE:
-            case INSERT:
-            case MERGE:
-                tuples.add(desc.toTuple(context(), row, op, cols));
+        tuples.add(desc.toTuple(context(), row, op, cols));
 
-                flushTuples(false);
+        if (invokeFut != null) // Still waiting for previous invocation result.
+            return;
 
-                break;
-            default:
-                throw new UnsupportedOperationException(op.name());
-        }
+        flushTuples(false);
 
-        if (waiting == 0)
-            source().request(waiting = MODIFY_BATCH_SIZE);
+        if (invokeFut != null) {
+            invokeFut.listen(f -> {
+                // Push new task to execute in correct thread.
+                context().execute(() -> {
+                    processInvokeResult(f.get());
+
+                    flushTuples(false);
+                }, this::onError);
+            });
+        }
     }
 
     /** {@inheritDoc} */
@@ -171,8 +183,27 @@
             source().request(waiting = MODIFY_BATCH_SIZE);
 
         if (state == State.UPDATED && requested > 0) {
+            if (invokeFut != null) { // Still waiting for previous invocation result.
+                invokeFut.listen(f -> context().execute(this::tryEnd, this::onError));
+
+                return;
+            }
+
             flushTuples(true);
 
+            if (invokeFut != null) {
+                invokeFut.listen(f -> {
+                    // Push new task to execute in correct thread.
+                    context().execute(() -> {
+                        processInvokeResult(f.get());
+
+                        tryEnd();
+                    }, this::onError);
+                });
+
+                return;
+            }
+
             state = State.END;
 
             inLoop = true;
@@ -193,7 +224,7 @@
 
     /** */
     @SuppressWarnings("unchecked")
-    private void flushTuples(boolean force) throws IgniteCheckedException {
+    private void flushTuples(boolean force) throws Exception {
         if (F.isEmpty(tuples) || !force && tuples.size() < MODIFY_BATCH_SIZE)
             return;
 
@@ -225,15 +256,30 @@
     private void invokeOutsideTransaction(
         List<ModifyTuple> tuples,
         IgniteInternalCache<Object, Object> cache
-    ) throws IgniteCheckedException {
+    ) throws Exception {
         SessionContextImpl sesCtx = context().unwrap(SessionContextImpl.class);
         Map<String, String> sesAttrs = sesCtx == null ? null : sesCtx.attributes();
 
         if (sesAttrs != null)
             cache = cache.withApplicationAttributes(sesAttrs);
 
         Map<Object, EntryProcessor<Object, Object, Long>> map = invokeMap(tuples);
-        Map<Object, EntryProcessorResult<Long>> res = cacheForDML(cache).invokeAll(map);
+        invokeFut = cacheForDML(cache).invokeAllAsync(map);
+
+        try {
+            // Shortcut - give a chance for operation to be executed in sync mode (it will simplify workflow).
+            Map<Object, EntryProcessorResult<Long>> res = invokeFut.get(INVOKE_TIMEOUT);
+
+            processInvokeResult(res);
+        }
+        catch (IgniteFutureTimeoutCheckedException ignore) {
+            // No-op. Result processing task will be scheduled by caller if invokeFut != null.
+        }
+    }
+
+    /** */
+    private void processInvokeResult(Map<Object, EntryProcessorResult<Long>> res) throws Exception {
+        invokeFut = null;
 
         long updated = res.values().stream().mapToLong(EntryProcessorResult::get).sum();
 
@@ -247,6 +293,9 @@
         }
 
         updatedRows += updated;
+
+        if (waiting == 0)
+            source().request(waiting = MODIFY_BATCH_SIZE);
     }
 
     /**
@@ -260,7 +309,7 @@
         List<ModifyTuple> tuples,
         IgniteInternalCache<Object, Object> cache,
         GridNearTxLocal userTx
-    ) throws IgniteCheckedException {
+    ) throws Exception {
         userTx.resume();
 
         try {

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/AbstractExecutionTest.java

@@ -88,7 +88,10 @@ public class AbstractExecutionTest extends GridCommonAbstractTest {
     protected static final String PARAMS_STRING = "Task executor = {0}, Execution strategy = {1}";
 
     /** */
-    protected static final int IN_BUFFER_SIZE = AbstractNode.IN_BUFFER_SIZE;
+    public static final int IN_BUFFER_SIZE = AbstractNode.IN_BUFFER_SIZE;
+
+    /** */
+    public static final int MODIFY_BATCH_SIZE = AbstractNode.MODIFY_BATCH_SIZE;
 
     /** */
     private Throwable lastE;

modules/calcite/src/test/java/org/apache/ignite/internal/processors/tx/TxThreadLockingTest.java

@@ -0,0 +1,219 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.processors.tx;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.apache.ignite.IgniteCache;
+import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.cache.CacheAtomicityMode;
+import org.apache.ignite.configuration.CacheConfiguration;
+import org.apache.ignite.configuration.IgniteConfiguration;
+import org.apache.ignite.internal.IgniteInternalFuture;
+import org.apache.ignite.internal.processors.query.IgniteSQLException;
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.AbstractExecutionTest;
+import org.apache.ignite.internal.processors.query.calcite.integration.AbstractBasicIntegrationTest;
+import org.apache.ignite.internal.util.typedef.X;
+import org.apache.ignite.testframework.GridTestUtils;
+import org.apache.ignite.transactions.Transaction;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import static org.apache.ignite.internal.processors.query.calcite.CalciteQueryProcessor.IGNITE_CALCITE_USE_QUERY_BLOCKING_TASK_EXECUTOR;
+import static org.apache.ignite.transactions.TransactionConcurrency.PESSIMISTIC;
+import static org.apache.ignite.transactions.TransactionIsolation.READ_COMMITTED;
+
+/** */
+@RunWith(Parameterized.class)
+public class TxThreadLockingTest extends AbstractBasicIntegrationTest {
+    /** */
+    private static final long TIMEOUT = 10_000L;
+
+    /** */
+    private static final int POOL_SIZE = 10;
+
+    /** */
+    private static final int IN_BUFFER_SIZE = AbstractExecutionTest.IN_BUFFER_SIZE;
+
+    /** */
+    private static final int MODIFY_BATCH_SIZE = AbstractExecutionTest.MODIFY_BATCH_SIZE;
+
+    /** Use query blocking executor. */
+    @Parameterized.Parameter(0)
+    public boolean qryBlockingExecutor;
+
+    /** */
+    @Parameterized.Parameters(name = "qryBlockingExecutor={0}")
+    public static Collection<?> parameters() {
+        return List.of(false, true);
+    }
+
+    /** {@inheritDoc} */
+    @Override protected void beforeTestsStarted() throws Exception {
+        // No-op.
+    }
+
+    /** {@inheritDoc} */
+    @Override protected void beforeTest() throws Exception {
+        System.setProperty(IGNITE_CALCITE_USE_QUERY_BLOCKING_TASK_EXECUTOR, String.valueOf(qryBlockingExecutor));
+
+        startGrid(0);
+
+        client = startClientGrid();
+    }
+
+    /** {@inheritDoc} */
+    @Override protected void afterTest() throws Exception {
+        stopAllGrids();
+
+        System.clearProperty(IGNITE_CALCITE_USE_QUERY_BLOCKING_TASK_EXECUTOR);
+    }
+
+    /** {@inheritDoc} */
+    @Override protected IgniteConfiguration getConfiguration(String igniteInstanceName) throws Exception {
+        return super.getConfiguration(igniteInstanceName).setQueryThreadPoolSize(POOL_SIZE);
+    }
+
+    /** {@inheritDoc} */
+    @Override protected <K, V> CacheConfiguration<K, V> cacheConfiguration() {
+        return super.<K, V>cacheConfiguration().setAtomicityMode(CacheAtomicityMode.TRANSACTIONAL);
+    }
+
+    /** */
+    @Test
+    public void testThreadsLockingByDml() throws Exception {
+        createAndPopulateTable();
+
+        CountDownLatch unlockLatch = new CountDownLatch(1);
+
+        IgniteInternalFuture<?> lockKeyFut = lockKeyAndWaitForLatch(0, unlockLatch);
+
+        IgniteInternalFuture<?> updFut = GridTestUtils.runMultiThreadedAsync(() -> {
+            sql("UPDATE person SET salary = 1 WHERE id = 0");
+        }, POOL_SIZE + 1, "update-thread");
+
+        // Ensure update queries are blocked.
+        assertFalse(GridTestUtils.waitForCondition(updFut::isDone, 200));
+
+        IgniteInternalFuture<?> selectFut = GridTestUtils.runMultiThreadedAsync(() -> {
+            for (int i = 0; i < 100; i++)
+                sql("SELECT * FROM person WHERE id = 1");
+        }, POOL_SIZE + 1, "select-thread");
+
+        // Ensure other queries are not blocked by update queries.
+        selectFut.get(TIMEOUT, TimeUnit.MILLISECONDS);
+
+        unlockLatch.countDown();
+
+        lockKeyFut.get();
+
+        // Ensure update queries are released.
+        updFut.get(TIMEOUT, TimeUnit.MILLISECONDS);
+    }
+
+    /** */
+    @Test
+    public void testDifferentBlockedBatchSize() throws Exception {
+        assertTrue("Unexpected constants [MODIFY_BATCH_SIZE=" + MODIFY_BATCH_SIZE +
+            ", IN_BUFFER_SIZE=" + IN_BUFFER_SIZE + ']', MODIFY_BATCH_SIZE < IN_BUFFER_SIZE);
+
+        createAndPopulateTable();
+
+        IgniteCache<Integer, Employer> cache = client.cache(TABLE_NAME);
+
+        for (int i = 0; i < IN_BUFFER_SIZE + 1; i++)
+            cache.put(i, new Employer("Test" + i, (double)i));
+
+        for (int size : new int[] {MODIFY_BATCH_SIZE, MODIFY_BATCH_SIZE + 1, IN_BUFFER_SIZE, IN_BUFFER_SIZE + 1}) {
+            log.info("Blocked batch size: " + size);
+
+            CountDownLatch unlockLatch = new CountDownLatch(1);
+
+            IgniteInternalFuture<?> lockKeyFut = lockKeyAndWaitForLatch(0, unlockLatch);
+
+            IgniteInternalFuture<List<List<?>>> updFut = GridTestUtils.runAsync(() ->
+                sql("UPDATE person SET salary = 1 WHERE id < ?", size));
+
+            // Ensure update query is blocked.
+            assertFalse(GridTestUtils.waitForCondition(updFut::isDone, 200));
+
+            unlockLatch.countDown();
+
+            lockKeyFut.get();
+
+            List<List<?>> res = updFut.get(TIMEOUT, TimeUnit.MILLISECONDS);
+
+            assertEquals((long)size, res.get(0).get(0));
+        }
+    }
+
+    /** */
+    @Test
+    public void testErrorAfterAsyncWait() throws Exception {
+        createAndPopulateTable();
+
+        CountDownLatch unlockLatch = new CountDownLatch(1);
+
+        IgniteInternalFuture<?> lockKeyFut = lockKeyAndWaitForLatch(0, unlockLatch);
+
+        IgniteInternalFuture<List<List<?>>> insFut = GridTestUtils.runAsync(() ->
+            sql("INSERT INTO person (id, name, salary) VALUES (0, 'Test', 0)"));
+
+        // Ensure insert query is blocked.
+        assertFalse(GridTestUtils.waitForCondition(insFut::isDone, 200));
+
+        unlockLatch.countDown();
+
+        lockKeyFut.get();
+
+        // Ensure exception is propagated to query initiator after async batch execution finished.
+        try {
+            insFut.get(TIMEOUT, TimeUnit.MILLISECONDS);
+
+            fail("Exception wasn't thrown");
+        }
+        catch (IgniteCheckedException expected) {
+            assertTrue(X.hasCause(expected, "Failed to INSERT some keys because they are already in cache",
+                IgniteSQLException.class));
+        }
+    }
+
+    /** */
+    private IgniteInternalFuture<?> lockKeyAndWaitForLatch(int key, CountDownLatch unlockLatch) throws Exception {
+        CountDownLatch lockLatch = new CountDownLatch(1);
+
+        IgniteInternalFuture<?> lockKeyFut = GridTestUtils.runAsync(() -> {
+            try (Transaction tx = client.transactions().txStart(PESSIMISTIC, READ_COMMITTED)) {
+                client.cache(TABLE_NAME).put(key, new Employer("Test", 0d));
+
+                lockLatch.countDown();
+
+                assertTrue(unlockLatch.await(TIMEOUT, TimeUnit.MILLISECONDS));
+
+                tx.commit();
+            }
+        });
+
+        assertTrue(lockLatch.await(TIMEOUT, TimeUnit.MILLISECONDS));
+
+        return lockKeyFut;
+    }
+}