feat: add ability to wait for parse cancellation in TSParser

Author: itsaky

android-tree-sitter/src/main/java/com/itsaky/androidide/treesitter/TSParser.java

@@ -21,6 +21,7 @@
 import com.itsaky.androidide.treesitter.string.UTF16StringFactory;
 import com.itsaky.androidide.treesitter.util.TSObjectFactoryProvider;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 
 /**
@@ -30,6 +31,7 @@
 public class TSParser extends TSNativeObject {
 
   protected final ReentrantLock parseLock = new ReentrantLock();
+  protected final Condition parseCondition = parseLock.newCondition();
   protected final AtomicBoolean isParsing = new AtomicBoolean(false);
   protected final AtomicBoolean isCancellationRequested = new AtomicBoolean(false);
 
@@ -157,9 +159,9 @@ public TSTree parseString(TSTree oldTree, String source) {
    * method.
    * <p>
    * Throws {@link ParseInProgressException} if the parser is currently parsing a syntax tree and
-   * the cancellation was NOT requested using {@link #requestCancellation()}. This method blocks the
-   * current thread if the previous parse was requested to be cancelled but the parse operation has
-   * not been cancelled yet.
+   * the cancellation was NOT requested using {@link #requestCancellationAsync()}. This method
+   * blocks the current thread if the previous parse was requested to be cancelled but the parse
+   * operation has not been cancelled yet.
    *
    * @param oldTree The previously parsed syntax tree.
    * @param source  The source code to parse.
@@ -192,6 +194,7 @@ public TSTree parseString(TSTree oldTree, UTF16String source) {
       return createTree(tree);
     } finally {
       unsetParsingFlag();
+      parseCondition.signalAll();
       parseLock.unlock();
     }
   }
@@ -245,17 +248,38 @@ protected boolean unsetParsingFlag() {
    * Request the parsing operation to be cancelled if the parser is in the process of parsing a
    * syntax tree.
    * <p>
-   * The parse operation is NOT immediately cancelled.
+   * This is an asynchronous operation and the previous parse call may NOT be cancelled immediately.
+   * Use {@link #requestCancellationAndWait()} for a blocking cancellation request.
    *
    * @return <code>true</code> if the cancellation was requested successfully, <code>false</code>
    * otherwise.
    */
-  public boolean requestCancellation() {
+  public boolean requestCancellationAsync() {
     final var requested = Native.requestCancellation();
     setCancellationRequested(requested);
     return requested;
   }
 
+  /**
+   * If the parser is parsing syntax tree, sets the cancellation flag and blocks the current thread
+   * until the parse operation returns. Does nothing if {@link #requestCancellationAsync()} returns
+   * false.
+   */
+  public void requestCancellationAndWait() {
+    if (requestCancellationAsync()) {
+      parseLock.lock();
+      try {
+        while (isParsing()) {
+          parseCondition.await();
+        }
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      } finally {
+        parseLock.unlock();
+      }
+    }
+  }
+
   protected synchronized void setCancellationRequested(boolean isRequested) {
     this.isCancellationRequested.set(isRequested);
   }

android-tree-sitter/src/test/java/com/itsaky/androidide/treesitter/ParserTest.java

@@ -322,7 +322,7 @@ public void testParserCancellation() {
       final var parseFuture = executor.schedule(() -> {
 
         // cancel the parsing after 200ms
-        executor.schedule(() -> assertThat(parser.requestCancellation()).isTrue(), 200,
+        executor.schedule(() -> assertThat(parser.requestCancellationAsync()).isTrue(), 200,
           TimeUnit.MILLISECONDS);
 
         // parsing the View.java.txt file takes 300-600ms
@@ -336,6 +336,8 @@ public void testParserCancellation() {
         parseFuture.get();
       } catch (InterruptedException | ExecutionException e) {
         throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
       }
     }
   }
@@ -385,6 +387,8 @@ public void testParserParseCallShouldFailIfAnotherParseIsInProgress() {
         parseFuture2.get();
       } catch (Throwable e) {
         throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
       }
     }
   }
@@ -419,7 +423,7 @@ public void testParserParseCallShouldSucceedIfAnotherParseIsInProgressAndCancell
         assertThat(parser.isParsing()).isTrue();
 
         // request the cancellation
-        assertThat(parser.requestCancellation()).isTrue();
+        assertThat(parser.requestCancellationAsync()).isTrue();
 
         // the next parse call should wait for the previous parse call to return
         try (var tree = parser.parseString(fileContent)) {
@@ -436,6 +440,8 @@ public void testParserParseCallShouldSucceedIfAnotherParseIsInProgressAndCancell
         parseFuture2.get();
       } catch (Throwable e) {
         throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
       }
     }
   }
@@ -488,6 +494,63 @@ public void testParserParseCallShouldFailIfAnotherParseIsInProgressAndCancellati
         parseFuture2.get();
       } catch (Throwable e) {
         throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
+      }
+    }
+  }
+
+  @Test
+  public void testAwaitedCancellation() {
+    try (final var parser = TSParser.create(); final var mainParseContent = UTF16StringFactory.newString()) {
+      parser.setLanguage(TSLanguageJava.getInstance());
+
+      // Read the content before starting the threads
+      final var fileContent = readResource("View.java.txt");
+      mainParseContent.append(fileContent);
+      mainParseContent.append(fileContent);
+      mainParseContent.append(fileContent);
+
+      final var executor = Executors.newScheduledThreadPool(20);
+
+      // start the main parse operation immediately
+      final var parseFuture1 = executor.schedule(() -> {
+        try (final var tree = parser.parseString(mainParseContent)) {
+          // This parse was cancelled and another parse was requested
+          // so this should fail
+          assertThat(tree).isNull();
+        }
+      }, 0, TimeUnit.MICROSECONDS);
+
+      // delay the second parse so that the parser is in the 'parsing' state when this is executed
+      final var secondParseDelayMs = 100;
+      final var parseFuture2 = executor.schedule(() -> {
+
+        // the parser should be in the 'parsing' state by now
+        assertThat(parser.isParsing()).isTrue();
+
+        // request parse cancellation and wait till the parse returns
+        final var start = System.currentTimeMillis();
+        parser.requestCancellationAndWait();
+        System.err.println("cancelAndWait() waited for " + (System.currentTimeMillis() - start) + "ms");
+
+        // request another parse
+        try (var tree = parser.parseString(fileContent)) {
+          // A parse was already in progress when this parse was requested
+          // however, we cancelled that parse and requested this one
+          // so this parseString call should succeed
+          assertThat(tree).isNotNull();
+          assertThat(tree.canAccess()).isTrue();
+        }
+      }, secondParseDelayMs, TimeUnit.MILLISECONDS);
+
+      try {
+        parseFuture1.get();
+        parseFuture2.get();
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
       }
     }
   }