feat(java): warm up java engine

wren-main/src/main/java/io/wren/main/PreviewService.java

@@ -28,6 +28,7 @@
 
 package io.wren.main;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import com.google.inject.Inject;
 import io.airlift.log.Logger;
@@ -39,6 +40,7 @@
 import io.wren.base.client.duckdb.DuckDBConfig;
 import io.wren.base.config.ConfigManager;
 import io.wren.base.config.WrenConfig;
+import io.wren.base.dto.Manifest;
 import io.wren.base.sql.SqlConverter;
 import io.wren.base.sqlrewrite.WrenPlanner;
 import io.wren.main.metadata.Metadata;
@@ -49,6 +51,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
+import static io.wren.base.dto.Model.model;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
 
@@ -61,6 +64,8 @@ public class PreviewService
     private final ConfigManager configManager;
     private final ExecutorService connectionPool;
 
+    private boolean isWarmed;
+
     @Inject
     public PreviewService(
             Metadata metadata,
@@ -131,4 +136,26 @@ public CompletableFuture<List<Column>> dryRun(WrenMDL mdl, String sql)
             return metadata.describeQuery(converted, List.of());
         }, connectionPool);
     }
+
+    public boolean isWarmed()
+    {
+        return isWarmed;
+    }
+
+    public void warmUp()
+    {
+        WrenMDL mdl = WrenMDL.fromManifest(
+                Manifest.builder()
+                        .setCatalog("default")
+                        .setSchema("default")
+                        .setModels(ImmutableList.of(
+                                model("orders", "SELECT * FROM tpch.tiny.orders", ImmutableList.of())))
+                        .build());
+        dryPlan(mdl, "SELECT * FROM orders", true)
+                .thenRun(() -> {
+                    isWarmed = true;
+                    LOG.info("Warm up done");
+                })
+                .join();
+    }
 }

wren-main/src/main/java/io/wren/main/web/SystemResource.java

@@ -0,0 +1,33 @@
+package io.wren.main.web;
+
+import com.google.inject.Inject;
+import io.wren.main.PreviewService;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.container.AsyncResponse;
+import jakarta.ws.rs.container.Suspended;
+import jakarta.ws.rs.core.Response;
+
+@Path("/v2")
+public class SystemResource
+{
+    private final PreviewService previewService;
+
+    @Inject
+    public SystemResource(PreviewService previewService)
+    {
+        this.previewService = previewService;
+    }
+
+    @GET
+    @Path("/health")
+    public void health(@Suspended AsyncResponse asyncResponse)
+    {
+        if (previewService.isWarmed()) {
+            asyncResponse.resume(Response.ok().build());
+        }
+        else {
+            asyncResponse.resume(Response.status(Response.Status.SERVICE_UNAVAILABLE).build());
+        }
+    }
+}

wren-server/src/main/java/io/wren/server/WrenServer.java

@@ -15,12 +15,14 @@
 package io.wren.server;
 
 import com.google.common.collect.ImmutableList;
+import com.google.inject.Injector;
 import com.google.inject.Module;
 import io.airlift.event.client.EventModule;
 import io.airlift.http.server.HttpServerModule;
 import io.airlift.jaxrs.JaxrsModule;
 import io.airlift.json.JsonModule;
 import io.airlift.node.NodeModule;
+import io.wren.main.PreviewService;
 import io.wren.main.WrenModule;
 import io.wren.main.server.Server;
 import io.wren.server.module.DuckDBConnectorModule;
@@ -49,4 +51,15 @@ protected Iterable<? extends Module> getAdditionalModules()
                 new WrenModule(),
                 new WebModule());
     }
+
+    @Override
+    protected void configure(Injector injector)
+    {
+        warmUp(injector);
+    }
+
+    private void warmUp(Injector injector)
+    {
+        injector.getInstance(PreviewService.class).warmUp();
+    }
 }

wren-server/src/main/java/io/wren/server/module/WebModule.java

@@ -25,6 +25,7 @@
 import io.wren.main.web.DuckDBResource;
 import io.wren.main.web.MDLResource;
 import io.wren.main.web.MDLResourceV2;
+import io.wren.main.web.SystemResource;
 import io.wren.main.web.WrenExceptionMapper;
 
 import static io.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
@@ -41,6 +42,7 @@ protected void setup(Binder binder)
         jaxrsBinder(binder).bind(AnalysisResourceV2.class);
         jaxrsBinder(binder).bind(ConfigResource.class);
         jaxrsBinder(binder).bind(DuckDBResource.class);
+        jaxrsBinder(binder).bind(SystemResource.class);
         jaxrsBinder(binder).bindInstance(new WrenExceptionMapper());
         binder.bind(PreviewService.class).in(Scopes.SINGLETON);
         binder.bind(ValidationService.class).in(Scopes.SINGLETON);