Improve performance of file revision data streaming

Author: malte-hansen

src/main/java/net/explorviz/code/analysis/export/DataExporter.java

@@ -1,5 +1,7 @@
 package net.explorviz.code.analysis.export;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CountDownLatch;
@@ -17,14 +19,61 @@
  */
 public interface DataExporter {
 
+  /**
+   * Number of files collected before a single {@code PersistFiles} streaming call
+   * is made.
+   */
+  int FILE_PERSIST_BATCH_SIZE = 250;
+
   StateData getStateData(final String repositoryName, final String branchName, final String token,
       final Map<String, String> applicationPaths);
 
   void persistFile(final FileData fileData);
 
+  /**
+   * Persists a batch of files in a single operation. The default implementation
+   * calls {@link #persistFile} for each file sequentially; exporters that support
+   * client-streaming gRPC should override this to use {@code PersistFiles}.
+   */
+  default void persistFilesBatch(final List<FileData> files) {
+    files.forEach(this::persistFile);
+  }
+
+  /**
+   * Drains {@code completedFiles} as files arrive and dispatches them to
+   * {@link #persistFilesBatch} in chunks of {@link #FILE_PERSIST_BATCH_SIZE}.
+   * Runs
+   * concurrently with the analysis pipeline; call
+   * {@code analysisFinished.countDown()}
+   * once all producers have finished to signal the final flush.
+   */
+  default void persistFilesFromQueueInBatches(final BlockingQueue<FileData> completedFiles,
+      final CountDownLatch analysisFinished) {
+    final List<FileData> batch = new ArrayList<>(FILE_PERSIST_BATCH_SIZE);
+    try {
+      while (analysisFinished.getCount() > 0 || !completedFiles.isEmpty()) {
+        final FileData fileData = completedFiles.poll(100, TimeUnit.MILLISECONDS);
+        if (fileData == null) {
+          continue;
+        }
+        batch.add(fileData);
+        if (batch.size() >= FILE_PERSIST_BATCH_SIZE) {
+          persistFilesBatch(List.copyOf(batch));
+          batch.clear();
+        }
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    if (!batch.isEmpty()) {
+      persistFilesBatch(List.copyOf(batch));
+    }
+  }
+
   /**
    * Persists files as they become available on the queue, while analysis is still
-   * running. Uses {@link Runtime#availableProcessors()} minus one when parallelism
+   * running. Uses {@link Runtime#availableProcessors()} minus one when
+   * parallelism
    * is not positive.
    */
   default void persistFilesFromQueue(final BlockingQueue<FileData> completedFiles,

src/main/java/net/explorviz/code/analysis/export/GrpcExporter.java

@@ -1,13 +1,16 @@
 package net.explorviz.code.analysis.export;
 
 import io.quarkus.grpc.GrpcClient;
+import io.smallrye.mutiny.Multi;
 import jakarta.enterprise.context.ApplicationScoped;
+import java.util.List;
 import java.util.Map;
 import net.explorviz.code.proto.CommitData;
 import net.explorviz.code.proto.CommitServiceGrpc;
 import net.explorviz.code.proto.ContributorServiceGrpc;
 import net.explorviz.code.proto.FileData;
 import net.explorviz.code.proto.FileDataServiceGrpc;
+import net.explorviz.code.proto.MutinyFileDataServiceGrpc;
 import net.explorviz.code.proto.StateData;
 import net.explorviz.code.proto.StateDataRequest;
 import net.explorviz.code.proto.StateDataServiceGrpc;
@@ -30,6 +33,9 @@ public final class GrpcExporter implements DataExporter {
   @GrpcClient(GRPC_CLIENT_NAME)
   /* package */ FileDataServiceGrpc.FileDataServiceBlockingStub fileDataGrpcClient;
 
+  @GrpcClient(GRPC_CLIENT_NAME)
+  /* package */ MutinyFileDataServiceGrpc.MutinyFileDataServiceStub fileDataMutinyGrpcClient;
+
   @GrpcClient(GRPC_CLIENT_NAME)
   /* package */ CommitServiceGrpc.CommitServiceBlockingStub commitDataGrpcClient;
 
@@ -90,6 +96,23 @@ public void persistFile(final FileData fileData) {
     }
   }
 
+  @Override
+  public void persistFilesBatch(final List<FileData> files) {
+    if (files.isEmpty()) {
+      return;
+    }
+    LOGGER.info("Sending batch of {} files via streaming RPC", files.size());
+    try {
+      fileDataMutinyGrpcClient
+          .persistFiles(Multi.createFrom().items(files.stream()))
+          .await()
+          .indefinitely();
+    } catch (final Exception e) {
+      LOGGER.error("Failed to send batch of {} files: {}", files.size(), e.getMessage());
+      throw new RuntimeException("Failed to send file batch to landscape-service", e);
+    }
+  }
+
   @Override
   public void persistCommit(final CommitData commitData) {
     LOGGER.info("Sending commit data on {}", commitData.getCommitId());

src/main/java/net/explorviz/code/analysis/parser/AntlrCParserService.java

@@ -10,6 +10,7 @@
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.ParseTreeWalker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,10 +54,12 @@ public CFileDataHandler parseFile(final String pathToFile, final String fileHash
   private CFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
     final CLexer lexer = new CLexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final CParser parser = new CParser(tokens);
 
-    final CParser.CompilationUnitContext compilationUnit = parser.compilationUnit();
+    final ParseTree compilationUnit =
+        AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName, parser::compilationUnit);
 
     final CFileDataHandler fileDataHandler = new CFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);

src/main/java/net/explorviz/code/analysis/parser/AntlrCSharpParserService.java

@@ -10,6 +10,7 @@
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.ParseTreeWalker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,10 +54,12 @@ public CSharpFileDataHandler parseFile(final String pathToFile, final String fil
   private CSharpFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
     final CSharpLexer lexer = new CSharpLexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final CSharpParser parser = new CSharpParser(tokens);
 
-    final CSharpParser.ProgContext prog = parser.prog();
+    final ParseTree prog =
+        AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName, parser::prog);
 
     final CSharpFileDataHandler fileDataHandler = new CSharpFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);

src/main/java/net/explorviz/code/analysis/parser/AntlrCppParserService.java

@@ -10,6 +10,7 @@
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.ParseTreeWalker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -52,19 +53,17 @@ public CppFileDataHandler parseFile(final String pathToFile, final String fileHa
 
   private CppFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
-    // Create lexer and parser
     final CPP14Lexer lexer = new CPP14Lexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final CPP14Parser parser = new CPP14Parser(tokens);
 
-    // Parse the translation unit (entry point for C/C++)
-    final CPP14Parser.TranslationUnitContext translationUnit = parser.translationUnit();
+    final ParseTree translationUnit =
+        AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName, parser::translationUnit);
 
-    // Create C/C++ file data handler
     final CppFileDataHandler fileDataHandler = new CppFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);
 
-    // Create and execute the listener
     final CppFileDataListener listener = new CppFileDataListener(fileDataHandler, tokens);
     final ParseTreeWalker walker = new ParseTreeWalker();
     walker.walk(listener, translationUnit);

src/main/java/net/explorviz/code/analysis/parser/AntlrGoParserService.java

@@ -10,6 +10,7 @@
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.ParseTreeWalker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,10 +54,12 @@ public GoFileDataHandler parseFile(final String pathToFile, final String fileHas
   private GoFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
     final GoLexer lexer = new GoLexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final GoParser parser = new GoParser(tokens);
 
-    final GoParser.SourceFileContext sourceFile = parser.sourceFile();
+    final ParseTree sourceFile =
+        AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName, parser::sourceFile);
 
     final GoFileDataHandler fileDataHandler = new GoFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);

src/main/java/net/explorviz/code/analysis/parser/AntlrKotlinParserService.java

@@ -54,11 +54,13 @@ public KotlinFileDataHandler parseFile(final String pathToFile, final String fil
   private KotlinFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
     final KotlinLexer lexer = new KotlinLexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final KotlinParser parser = new KotlinParser(tokens);
 
     final boolean isScript = fileName.toLowerCase().endsWith(".kts");
-    final ParseTree tree = isScript ? parser.script() : parser.kotlinFile();
+    final ParseTree tree = AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName,
+        () -> isScript ? parser.script() : parser.kotlinFile());
 
     final KotlinFileDataHandler fileDataHandler = new KotlinFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);

src/main/java/net/explorviz/code/analysis/parser/AntlrParserService.java

@@ -10,6 +10,7 @@
 import org.antlr.v4.runtime.CharStream;
 import org.antlr.v4.runtime.CharStreams;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.ParseTreeWalker;
 import org.eclipse.microprofile.config.inject.ConfigProperty;
 import org.slf4j.Logger;
@@ -56,19 +57,17 @@ public JavaFileDataHandler parseFile(final String pathToFile, final String fileH
 
   private JavaFileDataHandler parse(final CharStream charStream, final String fileName,
       final String fileHash) {
-    // Create lexer and parser
     final JavaLexer lexer = new JavaLexer(charStream);
+    AntlrParserUtils.configureLexer(lexer);
     final CommonTokenStream tokens = new CommonTokenStream(lexer);
     final JavaParser parser = new JavaParser(tokens);
 
-    // Parse the compilation unit
-    final JavaParser.CompilationUnitContext compilationUnit = parser.compilationUnit();
+    final ParseTree compilationUnit =
+        AntlrParserUtils.parseTwoStage(parser, tokens, LOGGER, fileName, parser::compilationUnit);
 
-    // Create Java file data handler
     final JavaFileDataHandler fileDataHandler = new JavaFileDataHandler(fileName);
     fileDataHandler.setFileHash(fileHash);
 
-    // Create and execute the listener
     final JavaFileDataListener listener = new JavaFileDataListener(fileDataHandler,
         wildcardImportProperty, tokens);
     final ParseTreeWalker walker = new ParseTreeWalker();